Fix (calibrate): fix for minifloat act calibration (#966)

Xilinx · Jun 3, 2024 · 02f5b6b · 02f5b6b
1 parent 0f60606
commit 02f5b6b
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 10 deletions.
diff --git a/src/brevitas/graph/calibrate.py b/src/brevitas/graph/calibrate.py
@@ -15,7 +15,7 @@
 from brevitas.nn.quant_layer import QuantWeightBiasInputOutputLayer as QuantWBIOL
 from brevitas.proxy.parameter_quant import BiasQuantProxyFromInjectorBase
 from brevitas.proxy.parameter_quant import WeightQuantProxyFromInjectorBase
-from brevitas.proxy.runtime_quant import ActQuantProxyFromInjector
+from brevitas.proxy.runtime_quant import ActQuantProxyFromInjectorBase
 from brevitas.proxy.runtime_quant import ClampQuantProxyFromInjector
 from brevitas.proxy.runtime_quant import TruncQuantProxyFromInjector
 from brevitas.quant_tensor import QuantTensor
@@ -188,7 +188,7 @@ def disable_act_quantization(self, model, is_training):
         # will be discarded through the hook. It is useful for collecting activation stats,
         # for example during activation calibration in PTQ
         for module in model.modules():
-            if isinstance(module, ActQuantProxyFromInjector):
+            if isinstance(module, ActQuantProxyFromInjectorBase):
                 module.train(is_training)
                 if self.call_act_quantizer_impl:
                     hook = module.register_forward_hook(self.disable_act_quant_hook)
@@ -216,7 +216,7 @@ def enable_act_quantization(self, model, is_training):
             if isinstance(module, _ACC_PROXIES):
                 module.train(is_training)
                 module.disable_quant = False
-            elif isinstance(module, ActQuantProxyFromInjector):
+            elif isinstance(module, ActQuantProxyFromInjectorBase):
                 module.disable_quant = False
                 module.train(is_training)
         for hook in self.disable_act_quant_hooks:

diff --git a/src/brevitas_examples/imagenet_classification/ptq/ptq_common.py b/src/brevitas_examples/imagenet_classification/ptq/ptq_common.py
@@ -141,13 +141,14 @@ class CNNInt8DynamicActPerTensorFloat(Int8DynamicActPerTensorFloat):
                         'sym': CNNInt8DynamicActPerTensorFloat,
                         'asym': CNNShiftedUint8DynamicActPerTensorFloat}}}}},
     'float': {
-        'float_scale': {
-            'stats': {
-                'per_tensor': {
-                    'sym': Fp8e4m3ActPerTensorFloat}},
-            'mse': {
-                'per_tensor': {
-                    'sym': Fp8e4m3ActPerTensorFloatMSE}}}}}
+        'static': {
+            'float_scale': {
+                'stats': {
+                    'per_tensor': {
+                        'sym': Fp8e4m3ActPerTensorFloat}},
+                'mse': {
+                    'per_tensor': {
+                        'sym': Fp8e4m3ActPerTensorFloatMSE}}}}}}
 
 
 def quantize_model(