Finer-grained quantization disabling

Author: pablomlago

src/brevitas/graph/calibrate.py

@@ -3,7 +3,7 @@
 
 from functools import partial
 import sys
-from typing import List, Optional
+from typing import List, Optional, Tuple, Type
 
 import torch
 from torch import nn
@@ -13,6 +13,7 @@
 from brevitas.nn import QuantLinear
 from brevitas.nn.quant_layer import QuantWeightBiasInputOutputLayer as QuantWBIOL
 from brevitas.proxy.parameter_quant import BiasQuantProxyFromInjectorBase
+from brevitas.proxy.parameter_quant import ParameterQuantProxyFromInjector
 from brevitas.proxy.parameter_quant import WeightQuantProxyFromInjectorBase
 from brevitas.proxy.runtime_quant import ActQuantProxyFromInjectorBase
 from brevitas.proxy.runtime_quant import ClampQuantProxyFromInjector
@@ -30,6 +31,7 @@
 
 _PARAM_PROXIES = (WeightQuantProxyFromInjectorBase, BiasQuantProxyFromInjectorBase)
 
+_WEIGHT_PROXIES = (WeightQuantProxyFromInjectorBase)
 _BIAS_PROXIES = (BiasQuantProxyFromInjectorBase)
 
 _ACC_PROXIES = (TruncQuantProxyFromInjector, ClampQuantProxyFromInjector)
@@ -205,17 +207,40 @@ def disable_act_quantization(self, model, is_training):
                 module.train(is_training)
                 module.disable_quant = True
 
-    def disable_param_quantization(self, model, is_training):
-        for module in model.modules():
-            if isinstance(module, _PARAM_PROXIES):
-                module.train(is_training)
-                module.disable_quant = True
-
-    def disable_bias_quantization(self, model, is_training):
+    def _set_param_quantization(
+            self,
+            model: nn.Module,
+            is_training: bool,
+            disable_quant: bool,
+            quant_proxies: Tuple[Type[ParameterQuantProxyFromInjector]] = _PARAM_PROXIES) -> None:
         for module in model.modules():
-            if isinstance(module, _BIAS_PROXIES):
+            if isinstance(module, quant_proxies):
                 module.train(is_training)
-                module.disable_quant = True
+                module.disable_quant = disable_quant
+
+    def disable_param_quantization(self, model: nn.Module, is_training: bool) -> None:
+        self._set_param_quantization(
+            model=model,
+            is_training=is_training,
+            disable_quant=True,
+            quant_proxies=_PARAM_PROXIES,
+        )
+
+    def disable_bias_quantization(self, model: nn.Module, is_training: bool):
+        self._set_param_quantization(
+            model=model,
+            is_training=is_training,
+            disable_quant=True,
+            quant_proxies=_BIAS_PROXIES,
+        )
+
+    def disable_weight_quantization(self, model: nn.Module, is_training: bool):
+        self._set_param_quantization(
+            model=model,
+            is_training=is_training,
+            disable_quant=True,
+            quant_proxies=_WEIGHT_PROXIES,
+        )
 
     def enable_act_quantization(self, model, is_training):
         for module in model.modules():
@@ -229,17 +254,29 @@ def enable_act_quantization(self, model, is_training):
                     if hasattr(m, 'observer_only'):
                         m.observer_only = False
 
-    def enable_param_quantization(self, model, is_training):
-        for module in model.modules():
-            if isinstance(module, _PARAM_PROXIES):
-                module.disable_quant = False
-                module.train(is_training)
-
-    def enable_bias_quantization(self, model, is_training):
-        for module in model.modules():
-            if isinstance(module, _BIAS_PROXIES):
-                module.disable_quant = False
-                module.train(is_training)
+    def enable_param_quantization(self, model: nn.Module, is_training: bool):
+        self._set_param_quantization(
+            model=model,
+            is_training=is_training,
+            disable_quant=False,
+            quant_proxies=_PARAM_PROXIES,
+        )
+
+    def enable_bias_quantization(self, model: nn.Module, is_training: bool):
+        self._set_param_quantization(
+            model=model,
+            is_training=is_training,
+            disable_quant=False,
+            quant_proxies=_BIAS_PROXIES,
+        )
+
+    def enable_weight_quantization(self, model: nn.Module, is_training: bool):
+        self._set_param_quantization(
+            model=model,
+            is_training=is_training,
+            disable_quant=False,
+            quant_proxies=_WEIGHT_PROXIES,
+        )
 
     def apply(self, model, is_training, quantization_enabled):
         if not quantization_enabled:
@@ -258,7 +295,10 @@ class disable_enable_quantization:
     Args:
         model (nn.Module): module for which quantization will be enabled/
             disabled
-        disable_quant (bool): whether to disable quantization
+        disable_quant_act (bool): whether to disable activation quantization
+        disable_weight_quant (bool): whether to disable weight quantization
+        disable_bias_quant (bool): whether to disable bias quantization
+        disable_out_quant (bool): whether to disable output quantization
         excluded_modules (list): list of submodules of modules to be excluded
             from quantization disabling
     """
@@ -267,12 +307,14 @@ def __init__(
             self,
             model: nn.Module,
             disable_act_quant: bool = True,
-            disable_param_quant: bool = True,
+            disable_weight_quant: bool = True,
+            disable_bias_quant: bool = True,
             disable_out_quant: bool = True,
             excluded_modules: Optional[List[nn.Module]] = None):
         self.model = model
         self.disable_act_quant = disable_act_quant
-        self.disable_param_quant = disable_param_quant
+        self.disable_weight_quant = disable_weight_quant
+        self.disable_bias_quant = disable_bias_quant
         self.disable_out_quant = disable_out_quant
         self.excluded_modules = excluded_modules if excluded_modules is not None else []
         self.disable_quant_class = DisableEnableQuantization()
@@ -281,8 +323,10 @@ def __init__(
     def __enter__(self):
         if self.disable_act_quant:
             self.disable_quant_class.disable_act_quantization(self.model, False)
-        if self.disable_param_quant:
-            self.disable_quant_class.disable_param_quantization(self.model, False)
+        if self.disable_weight_quant:
+            self.disable_quant_class.disable_weight_quantization(self.model, False)
+        if self.disable_bias_quant:
+            self.disable_quant_class.disable_bias_quantization(self.model, False)
         if self.disable_out_quant:
             self.return_quant_tensor_state = disable_return_quant_tensor(self.model)
         # Re-enable quantization for excluded modules
@@ -294,8 +338,10 @@ def __enter__(self):
     def __exit__(self, type, value, traceback):
         if self.disable_act_quant:
             self.disable_quant_class.enable_act_quantization(self.model, False)
-        if self.disable_param_quant:
-            self.disable_quant_class.enable_param_quantization(self.model, False)
+        if self.disable_weight_quant:
+            self.disable_quant_class.enable_weight_quantization(self.model, False)
+        if self.disable_bias_quant:
+            self.disable_quant_class.enable_bias_quantization(self.model, False)
         if self.disable_out_quant:
             restore_return_quant_tensor(self.model, self.return_quant_tensor_state)