optimum
optimum copied to clipboard
Add batchnorm folding transformations (fx)
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?