Add QuantizationMixin.ignore API

Signed-off-by: Kyunggeun Lee <quic_kyunggeu@quicinc.com>
Co-authored-by: Kyunggeun Lee <quic_kyunggeu@quicinc.com>

Author: 2 people

Docs/apiref/torch/nn.rst

@@ -21,7 +21,7 @@ for :class:`torch.nn.Conv2d` or :class:`QuantizedSoftmax` for :class:`torch.nn.S
 
 .. autoclass:: QuantizationMixin
    :noindex:
-   :members: __quant_init__, forward, compute_encodings
+   :members: __quant_init__, forward, compute_encodings, ignore
 
 Configuration
 -------------

TrainingExtensions/torch/src/python/aimet_torch/v2/nn/base.py

@@ -107,6 +107,8 @@ def generate_err_msg(self) -> str:
             f"Please register the quantized module definition of {module_cls} "
             f"using `@{mixin_cls.__name__}.implements({module_cls.__name__})` decorator.\n\n"
             f"For example:\n\n{code_example}\n\n"
+            "If you believe this module need not be quantized, please exclude it from quantization by calling "
+            f"`QuantizationMixin.ignore({module_cls.__name__})`.\n\n"
             f"For more details, please refer to the official API reference:\n{self.api_reference_url}"
         )
 
@@ -231,6 +233,8 @@ class BaseQuantizationMixin(abc.ABC):
     cls_to_qcls: dict
     qcls_to_cls: dict
 
+    _ignored_module_types = set()
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.__quant_init__()
@@ -365,6 +369,16 @@ def wrap(cls, module_cls: Type[nn.Module]):
         Wrap a regular module class into a quantized module class
         """
 
+    @classmethod
+    def ignore(cls, module_cls):
+        """
+        Exempt given module type from quantization
+        """
+        if not issubclass(module_cls, torch.nn.Module):
+            raise RuntimeError
+
+        cls._ignored_module_types.add(module_cls)
+
     @classmethod
     def implements(cls, module_cls):
         """
@@ -414,6 +428,8 @@ def from_module(cls, module: nn.Module):
         qtzn_module_cls = cls.cls_to_qcls.get(module_cls, None)
 
         if not qtzn_module_cls:
+            if module_cls in cls._ignored_module_types:
+                return module
             raise UnknownModuleError(module_cls, cls)
 
         qtzn_module = cls.__new__(qtzn_module_cls)

TrainingExtensions/torch/src/python/aimet_torch/v2/nn/true_quant.py

@@ -373,6 +373,26 @@ def from_module(cls, module: nn.Module):
         """
         return super().from_module(module)
 
+    @classmethod
+    def ignore(cls, module_cls):
+        """
+        Exempt given module type from quantization
+
+        Example:
+
+            >>> class MyModule(torch.nn.Module):
+                ...     def forward(self, x):
+                ...         return x ** 2
+            >>> QuantizationMixin.ignore(MyModule)
+            >>> model = torch.nn.Sequential(MyModule())
+            >>> sim = aimet_torch.QuantizationSimModel(model, torch.randn(10, 10))
+            >>> print(sim.model)
+            Sequential(
+              (0): MyModule()
+            )
+        """
+        super().ignore(module_cls)
+
     @classmethod
     def implements(cls, module_cls):
         r"""

TrainingExtensions/torch/src/python/aimet_torch/v2/quantsim/quantsim.py

@@ -185,6 +185,10 @@ def _convert_to_qmodule(module: torch.nn.Module):
                 f"using `@{e.mixin_cls.__name__}.implements()` decorator.",
                 "For example:",
                 *(e.generate_code_example() for e in exceptions.values()),
+                "If you believe these modules need not be quantized, "
+                "please exclude them from quantization by calling `QuantizationMixin.ignore` "
+                f"(for example, `QuantizationMixin.ignore({e.module_cls.__name__})`) "
+                "before creating QuantizationSimModel.",
                 f"For more details, please refer to the official API reference:\n{e.api_reference_url}",
             ]
         )

TrainingExtensions/torch/test/python/v2/nn/test_true_quant.py

@@ -1802,3 +1802,17 @@ def test_fold_param_quantizers(device, requires_grad):
     assert isinstance(qlinear.bias, torch.Tensor)
     assert isinstance(qlinear.bias, torch.nn.Parameter)
     assert torch.equal(qlinear.bias, original_bias)
+
+
+def test_ignore():
+    class MyModule(torch.nn.Module):
+        def forward(self, x):
+            return x**2
+
+    QuantizationMixin.ignore(MyModule)
+
+    model = torch.nn.Sequential(MyModule())
+    sim = aimet_torch.QuantizationSimModel(
+        model, dummy_input=torch.randn(1, 3, 224, 224)
+    )
+    assert type(sim.model[0]) == MyModule