[QAT] Make Fused modules torchscriptable (#136285)

Ninja91 · facebook-github-bot · commit 92e4eb3dbbd1 · 2024-09-27T17:59:15.000-07:00
Summary: Pull Request resolved: #136285 Same as title. Just moved around code to avoid torchscript failing on - False branch not having `conv_output` variable - type inference issue Inspired by: https://pytorch.org/tutorials/recipes/script_optimized.html#fix-common-errors-when-using-the-script-method Test Plan: CI Reviewed By: jerryzh168 Differential Revision: D62980019
diff --git a/test/quantization/core/test_workflow_module.py b/test/quantization/core/test_workflow_module.py
@@ -1,68 +1,66 @@
 # Owner(s): ["oncall: quantization"]
 
 # Torch
+# Standard library
+import copy
+import io
+import itertools
+import math
+import unittest
+
+import numpy as np
 import torch
+
+import torch.nn as nn
+import torch.testing._internal.hypothesis_utils as hu
+
+# Testing utils
+from hypothesis import given, settings, strategies as st
 from torch.ao.quantization import (
+    convert,
+    default_debug_qconfig,
+    default_histogram_observer,
+    default_observer,
+    default_per_channel_weight_observer,
+    FakeQuantize,
+    FixedQParamsObserver,
+    FusedMovingAvgObsFakeQuantize,
+    get_embedding_qat_module_mappings,
+    get_embedding_static_quant_module_mappings,
+    HistogramObserver,
     MinMaxObserver,
-    PerChannelMinMaxObserver,
     MovingAverageMinMaxObserver,
     MovingAveragePerChannelMinMaxObserver,
-    HistogramObserver,
-    RecordingObserver,
-    PlaceholderObserver,
     NoopObserver,
-    FakeQuantize,
-    FixedQParamsObserver,
-    default_debug_qconfig,
-    default_observer,
-    default_histogram_observer,
-    default_per_channel_weight_observer,
+    PerChannelMinMaxObserver,
+    PlaceholderObserver,
     prepare,
     prepare_qat,
-    convert,
     QConfig,
-    FusedMovingAvgObsFakeQuantize,
-    get_embedding_qat_module_mappings,
-    get_embedding_static_quant_module_mappings,
+    RecordingObserver,
 )
 from torch.ao.quantization.quantize import _get_observer_dict
 
-import torch.nn as nn
-
-# Standard library
-import copy
-import io
-import itertools
-import unittest
-import math
-import numpy as np
-
-# Testing utils
-from hypothesis import given, settings
-from hypothesis import strategies as st
-import torch.testing._internal.hypothesis_utils as hu
 hu.assert_deadline_disabled()
-from torch.testing._internal.common_cuda import TEST_MULTIGPU, TEST_CUDA
-from torch.testing._internal.common_utils import TestCase, skipIfTorchDynamo
+from torch.testing._internal.common_cuda import TEST_CUDA, TEST_MULTIGPU
+
 from torch.testing._internal.common_quantization import (
-    QuantizationTestCase,
     AnnotatedSingleLayerLinearModel,
-    test_only_eval_fn,
+    DeFusedEmbeddingBagLinear,
+    QuantizationTestCase,
     SingleLayerLinearModel,
+    test_only_eval_fn,
 )
 
 from torch.testing._internal.common_quantized import (
+    _fake_quantize_per_channel_affine_grad_reference,
+    _fake_quantize_per_channel_affine_reference,
+    override_qengines,
     override_quantized_engine,
     supported_qengines,
-    override_qengines,
-    _fake_quantize_per_channel_affine_reference,
-    _fake_quantize_per_channel_affine_grad_reference,
     to_tensor,
 )
-
-from torch.testing._internal.common_quantization import (
-    DeFusedEmbeddingBagLinear,
-)
+from torch.testing._internal.common_utils import skipIfTorchDynamo, TestCase
 
 NP_RANDOM_SEED = 19
 tolerance = 1e-6
@@ -915,6 +913,83 @@ def _get_buffer_ids(module):
     """
     return [id(v) for k, v in module._buffers.items()]
 
+class TestFusedModuleScriptable(QuantizationTestCase):
+    def test_fx_qat_convbn_fused_jit_scriptable(self):
+        """
+        Tests jit scriptability works for fused ConvBN.
+        """
+        for qengine in ['fbgemm', 'qnnpack']:
+            with override_quantized_engine(qengine):
+                # create conv-bn
+                class Model(nn.Module):
+                    def __init__(self) -> None:
+                        super().__init__()
+                        self.conv = nn.Conv2d(4, 1, 3, padding=1)
+                        self.bn = nn.BatchNorm2d(1)
+
+                    def forward(self, x):
+                        x = self.conv(x)
+                        x = self.bn(x)
+                        return x
+
+                model = Model()
+                model = torch.fx.symbolic_trace(model)
+
+                # fuse it
+                fused_model = torch.ao.quantization.fuse_modules_qat(
+                    model,
+                    [['conv', 'bn']],
+                )
+                # convert to QAT
+                qconfig_mapping = torch.ao.quantization.get_default_qat_qconfig_mapping(qengine)
+
+                quantizable_model = torch.ao.quantization.quantize_fx.prepare_qat_fx(fused_model,
+                                                                                     qconfig_mapping,
+                                                                                     example_inputs=None)
+                assert isinstance(quantizable_model.conv, torch.ao.nn.intrinsic.qat.ConvBn2d)
+
+                # jit script
+                scripted_model = torch.jit.script(quantizable_model)
+
+                self.assertTrue(
+                    isinstance(scripted_model, torch.jit.ScriptModule),
+                    "Expected prepared model with to be scriptable")
+
+    def test_qat_convbn_fused_jit_scriptable(self):
+        """
+        Tests jit scriptability works for fused ConvBN.
+        """
+        for qengine in ['fbgemm', 'qnnpack']:
+            with override_quantized_engine(qengine):
+                # create conv-bn
+                class Model(nn.Module):
+                    def __init__(self) -> None:
+                        super().__init__()
+                        self.conv = nn.Conv2d(4, 1, 3, padding=1)
+                        self.bn = nn.BatchNorm2d(1)
+
+                    def forward(self, x):
+                        x = self.conv(x)
+                        x = self.bn(x)
+                        return x
+
+                model = Model()
+
+                # fuse it
+                fused_model = torch.ao.quantization.fuse_modules_qat(
+                    model,
+                    [['conv', 'bn']],
+                )
+                # convert to QAT
+                fused_model.qconfig = torch.ao.quantization.get_default_qconfig(qengine)
+                torch.ao.quantization.prepare_qat(fused_model, inplace=True)
+                assert isinstance(fused_model.conv, torch.ao.nn.intrinsic.qat.ConvBn2d)
+
+                # Test jit script fails
+                # Prepared eager module fails due to observer hooks not being scriptable
+                with self.assertRaises(RuntimeError):
+                    torch.jit.script(fused_model)
+
 class TestDistributed(QuantizationTestCase):
 
     def test_observers_preserve_buffers(self):
diff --git a/torch/ao/nn/intrinsic/qat/modules/conv_fused.py b/torch/ao/nn/intrinsic/qat/modules/conv_fused.py
@@ -215,16 +215,15 @@ def _forward_slow(self, input):
                 )
                 self.bn(conv_out_bias)
 
-        # fused conv + bn without bias using bn running statistics
-        running_std = torch.sqrt(self.bn.running_var + self.bn.eps)
-        scale_factor = self.bn.weight / running_std
-        scaled_weight = self.weight_fake_quant(
-            self.weight * scale_factor.reshape(weight_shape)
-        )
-        # fused conv without bias for inference: (r * W / running_std) * X
-        conv_bn = self._conv_forward(input, scaled_weight, zero_bias)
+            # fused conv + bn without bias using bn running statistics
+            running_std = torch.sqrt(self.bn.running_var + self.bn.eps)
+            scale_factor = self.bn.weight / running_std
+            scaled_weight = self.weight_fake_quant(
+                self.weight * scale_factor.reshape(weight_shape)
+            )
+            # fused conv without bias for inference: (r * W / running_std) * X
+            conv_bn = self._conv_forward(input, scaled_weight, zero_bias)
 
-        if self.bn.training:
             avg_dims = [0] + list(range(2, len(self.weight.shape)))
             batch_mean = conv_out.mean(avg_dims)  # type: ignore[possibly-undefined]
             batch_var = torch.square(conv_out - batch_mean.reshape(bias_shape)).mean(
@@ -240,6 +239,15 @@ def _forward_slow(self, input):
             fused_mean = batch_mean
             fused_std = batch_std
         else:
+            # fused conv + bn without bias using bn running statistics
+            running_std = torch.sqrt(self.bn.running_var + self.bn.eps)
+            scale_factor = self.bn.weight / running_std
+            scaled_weight = self.weight_fake_quant(
+                self.weight * scale_factor.reshape(weight_shape)
+            )
+            # fused conv without bias for inference: (r * W / running_std) * X
+            conv_bn = self._conv_forward(input, scaled_weight, zero_bias)
+
             fused_mean = self.bn.running_mean - (
                 self.bias if self.bias is not None else 0
             )
@@ -563,7 +571,7 @@ def __init__(
         )
 
     def forward(self, input):
-        return F.relu(ConvBn1d._forward(self, input))
+        return F.relu(self._forward(input))
 
     @classmethod
     def from_float(cls, mod, use_precomputed_fake_quant=False):
@@ -760,7 +768,7 @@ def __init__(
         )
 
     def forward(self, input):
-        return F.relu(ConvBn2d._forward(self, input))
+        return F.relu(self._forward(input))
 
     @classmethod
     def from_float(cls, mod, use_precomputed_fake_quant=False):
