Fix static_quantization_tutorial error in qat_model

[bjhargrave]

The model must be in eval mode to fuse and in train mode to prepare_qat. After this fix, the tutorial runs to completion on a linux x86 system.

Fixes #1269

cc @jerryzh168 @jianyuh @sekyondaMeta @svekars @carljparker @NicolasHug @kit1980 @subramen

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/tutorials/2661

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 1726ae9 with merge base 66eaf6a ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[jerryzh168]

eval mode fuse_model might do different things comparing to train mode fuse_model, could you print the result after fuse_model before and after you add this change?

[bjhargrave]

could you print the result after fuse_model before and after you add this change

It throws an error stating fusion is only for eval:

Traceback (most recent call last):
  File "/root/tutorials/advanced_source/static_quantization_tutorial.py", line 400, in <module>
    qat_model.fuse_model()  
    ^^^^^^^^^^^^^^^^^^^^^^
  File "/root/tutorials/advanced_source/static_quantization_tutorial.py", line 175, in fuse_model
    torch.ao.quantization.fuse_modules(m, ['0', '1', '2'], inplace=True)
  File "/root/miniconda3/envs/pytorch-tutorials/lib/python3.11/site-packages/torch/ao/quantization/fuse_modules.py", line 158, in fuse_modules
    return _fuse_modules(
           ^^^^^^^^^^^^^^
  File "/root/miniconda3/envs/pytorch-tutorials/lib/python3.11/site-packages/torch/ao/quantization/fuse_modules.py", line 96, in _fuse_modules
    _fuse_modules_helper(model, modules_to_fuse, is_qat, fuser_func, fuse_custom_config_dict)
  File "/root/miniconda3/envs/pytorch-tutorials/lib/python3.11/site-packages/torch/ao/quantization/fuse_modules.py", line 84, in _fuse_modules_helper
    new_mod_list = fuser_func(mod_list, is_qat, additional_fuser_method_mapping)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/miniconda3/envs/pytorch-tutorials/lib/python3.11/site-packages/torch/ao/quantization/fuse_modules.py", line 56, in fuse_known_modules
    fused = fuser_method(is_qat, *mod_list)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/miniconda3/envs/pytorch-tutorials/lib/python3.11/site-packages/torch/ao/quantization/fuser_method_mappings.py", line 96, in fuse_conv_bn_relu
    fused_conv = nn.utils.fusion.fuse_conv_bn_eval(conv, bn)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/miniconda3/envs/pytorch-tutorials/lib/python3.11/site-packages/torch/nn/utils/fusion.py", line 27, in fuse_conv_bn_eval
    assert not (conv.training or bn.training), "Fusion only for eval!"

[jerryzh168]

I see, thanks, I think you should do:

mode.fuse_model(is_qat=True)

instead.

see: https://github.com/pytorch/vision/blob/main/torchvision/models/quantization/mobilenetv2.py#L58

[bjhargrave]

I think you should do:

mode.fuse_model(is_qat=True)

I updated the tutorial as suggested.

[jerryzh168]

thanks!