optimum icon indicating copy to clipboard operation
optimum copied to clipboard

Add batchnorm folding transformations (fx)

Open fxmarty opened this issue 3 years ago • 0 comments

What does this PR do?

Add the folding of nn.BatchNorm2d into nn.Conv2d, and of nn.BatchNorm1d into nn.Linear.

The former is useful for example for resnet, beit, regnet, dpt, levit, segformer and others. The latter is useful solely for groupvit, and could be for levit; but there is a flatten inbetween so the current transformation is not directly usable: https://github.com/huggingface/transformers/blob/ab2006e3d6db88654526a4169e65d4bfc52da2e3/src/transformers/models/levit/modeling_levit.py#L153-L156

Before submitting

  • [x] Did you make sure to update the documentation with your changes?
  • [x] Did you write any new necessary tests?

fxmarty avatar Aug 10 '22 16:08 fxmarty