[ONNX] training parameter of batch_norm is ignored

[shinh]

🚀 Feature

Provide a way to honor training passed to batch_norm

Motivation

After this PR, ONNX export starts using training parameter passed to torch.onnx.export instead of the one for batch_norm. This made it impossible to export frozen BN in training mode.

Pitch

How about honoring training of batch_norm when training=torch.onnx.TrainingMode.PRESERVE?

Alternatives

I personally like the behavior before this change, so reverting this is the best for me. However, changing this for TrainingMode.EVAL and TrainingMode.TRAINING might be confusing. This is why I am proposing changing the behavior only for TrainingMode.PRESERVE.

Additional context

I think it's common to use fixed BN even during training when you train detection/segmentation model with pre-trained classification backbone.

See frozen BN in detectron2 for example: https://github.com/facebookresearch/detectron2/blob/84a09834f6d838534951907fd9ef90fec73614d2/detectron2/layers/batch_norm.py#L14

cc @houseroad @spandantiwari @lara-hdr @BowenBao @neginraoof

[jspisak]

cc @SplitInfinity .. :)

[spandantiwari]

@shinh - There are a few different variables to consider here, such as model mode (model.eval() vs model.train()), export API training argument, and as you highlight, the training argument in the torch.nn.functional.batch_norm API. We can consider this, but I don't have an ETA on this just now.

I think it's common to use fixed BN even during training when you train detection/segmentation model with pre-trained classification backbone.

Regarding this, have you tried setting track_running_stats=False in the BatchNorm*d op? I am wondering if that solves this use case or not.

[shinh]

Unfortunately, torch.onnx outputs an invalid ONNX with track_running_stats=False :(

import torch
import onnx

class BNModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.bn = torch.nn.BatchNorm2d(1024, track_running_stats=False)

    def forward(self, x):
        return self.bn(x)

x = torch.rand(1, 1024, 4, 4)
torch.onnx.export(BNModule(), x, 'bn')

# ==> Context: Bad node spec: input: "input" input: "bn.weight" input: "bn.bias" input: "" input: "" output: "5" name: "BatchNormalization_0" op_type: "BatchNormalization" attribute { name: "epsilon" f: 1e-05 type: FLOAT } attribute { name: "momentum" f: 0.9 type: FLOAT }
torch.onnx.export(BNModule(), x, 'bn.onnx')

It outputs BN with mean/var unfilled, which current ONNX disallows. Anyway, I think track_running_stats=False is not what we want. It still uses mean/var of the input while BN in eval mode uses mean/var stored in nn.BatchNorm2d.

As for ETA, I understand this is not your priority. I'm happy to send a PR once we decide how to fix this. For example, what do you think of my proposal? I think TrainingMode.PRESERVE would imply "do what the model says" and it'd make some sense to honor training parameter in this mode. I'm guessing there are less users of TrainingMode.PRESERVE than other two options, so changing it slightly would not confuse users a lot?

[neginraoof]

@shinh
Thanks for reporting this. I'm looking into sending a fix for 'training' argument.
However, about the track_running_stats parameter, it should be supported by exporter correctly. We actually do have tests for this.
Can you specify which torch version and which onnx version you are using?

[shinh]

Now I confirmed it can be still reproduced by torch==1.7.1 with both onnx==1.8.0 and onnx==1.8.1.

When I wrote #49229 (comment), I believe I was using torch==1.6.0 and probably onnx==1.8.0.

It looks like #47903 is merged recently. I guess it is fixed but the fix is not in torch 1.7.1 yet?

[neginraoof]

@shinh Correct, this PR is not in the stable release yet. I believe 1.8 will be out soon, and it includes this fix.
Meanwhile, we work on the fix for this feature.

[neginraoof]

Here is the proposal to fix this issue: #52867

[thiagocrepaldi]

Resolved by #52758