Add batchnorm folding transformations (fx)

[fxmarty]

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

Fixes microsoft/onnxruntime#12522

Before submitting

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

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[michaelbenayoun]

Left a few comments but you should be able to merge after, thanks!!!

[michaelbenayoun]

No we need to be able to also discover other files, for instance when quantization will be merged, we will need to discover those tests as well.

[fxmarty]

Yes I agree, I did this considering my comment above:

@michaelbenayoun I rolled back to having a subfolder. I use python -m pytest fx/optimization/test_transformations.py --exitfirst in the yml because test_quantization.py is currently not run on main, and the tests don't pass; independenty from this PR (python -m unittest discover -s fx -p 'test_*.py' was not capturing it, see e.g. https://github.com/huggingface/optimum/actions/runs/3068109805/jobs/4955165866 )

[jogepari]

Hey, @fxmarty! Found this PR following an old ORT issue to grab your folding implementation, thx!

Sorry for digging an old thread, just wanted to comment that there might be an error here and here:
bn_b = bn1d.bias if bn1d.bias is not None else torch.ones_like(bn1d.running_mean)

In case of affine=False, bias should probably be torch.zeros_like(bn1d.running_mean).

[fxmarty]

Hi @jogepari Thank you for notifying, will fix!