Add QuantizationMixin.ignore_unknown_modules 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, ignore
+   :members: __quant_init__, forward, compute_encodings, ignore, ignore_unknown_modules
 
 Configuration
 -------------

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

@@ -233,6 +233,7 @@ class BaseQuantizationMixin(abc.ABC):
     cls_to_qcls: dict
     qcls_to_cls: dict
 
+    _ignore_unknown_modules: bool = False
     _ignored_module_types = set()
 
     def __init__(self, *args, **kwargs):
@@ -379,6 +380,13 @@ def ignore(cls, module_cls):
 
         cls._ignored_module_types.add(module_cls)
 
+    @classmethod
+    def ignore_unknown_modules(cls, ignore: bool = True):
+        """
+        Exempt all unkown module types from quantization
+        """
+        cls._ignore_unknown_modules = ignore
+
     @classmethod
     def implements(cls, module_cls):
         """
@@ -428,7 +436,7 @@ 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:
+            if cls._ignore_unknown_modules or module_cls in cls._ignored_module_types:
                 return module
             raise UnknownModuleError(module_cls, cls)
 

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

@@ -381,8 +381,8 @@ def ignore(cls, module_cls):
         Example:
 
             >>> class MyModule(torch.nn.Module):
-                ...     def forward(self, x):
-                ...         return x ** 2
+            ...     def forward(self, x):
+            ...         return x ** 2
             >>> QuantizationMixin.ignore(MyModule)
             >>> model = torch.nn.Sequential(MyModule())
             >>> sim = aimet_torch.QuantizationSimModel(model, torch.randn(10, 10))
@@ -393,6 +393,26 @@ def ignore(cls, module_cls):
         """
         super().ignore(module_cls)
 
+    @classmethod
+    def ignore_unknown_modules(cls, ignore: bool = True):
+        """
+        Exempt all unkown module types from quantization
+
+        Example:
+
+            >>> class MyModule(torch.nn.Module):
+            ...     def forward(self, x):
+            ...         return x ** 2
+            >>> QuantizationMixin.ignore_unknown_modules(True)
+            >>> model = torch.nn.Sequential(MyModule())
+            >>> sim = aimet_torch.QuantizationSimModel(model, torch.randn(10, 10))
+            >>> print(sim.model)
+            Sequential(
+              (0): MyModule()
+            )
+        """
+        super().ignore_unknown_modules(ignore)
+
     @classmethod
     def implements(cls, module_cls):
         r"""

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

@@ -1805,6 +1805,11 @@ def test_fold_param_quantizers(device, requires_grad):
 
 
 def test_ignore():
+    """
+    When: Call QuantizationMixin.ignore
+    Then: Unknown modules should be ignored during quantization
+    """
+
     class MyModule(torch.nn.Module):
         def forward(self, x):
             return x**2
@@ -1816,3 +1821,24 @@ def forward(self, x):
         model, dummy_input=torch.randn(1, 3, 224, 224)
     )
     assert type(sim.model[0]) == MyModule
+
+    """
+    When: Call QuantizationMixin.ignore_unknown_modules
+    Then: Unknown modules should be ignored during quantization
+    """
+
+    class MyModule(torch.nn.Module):
+        def forward(self, x):
+            return x**2
+
+    orig = QuantizationMixin._ignore_unknown_modules
+    try:
+        QuantizationMixin.ignore_unknown_modules(True)
+
+        model = torch.nn.Sequential(MyModule())
+        sim = aimet_torch.QuantizationSimModel(
+            model, dummy_input=torch.randn(1, 3, 224, 224)
+        )
+        assert type(sim.model[0]) == MyModule
+    finally:
+        QuantizationMixin.ignore_unknown_modules(orig)