Small clean-up

Xilinx · May 27, 2024 · ffa5059 · ffa5059
1 parent ca864ca
commit ffa5059
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/src/brevitas/core/stats/stats_op.py b/src/brevitas/core/stats/stats_op.py
@@ -443,9 +443,18 @@ def _set_local_loss_mode(module, enabled):
 
 
 def _set_quantize_zero_point(module, enabled):
+    prev_state = dict()
     for m in module.modules():
         if hasattr(m, 'quantize_zero_point'):
+            prev_state[m] = m.quantize_zero_point
             m.quantize_zero_point = enabled
+    return prev_state
+
+
+def _restore_quantize_zero_point(module, prev_state):
+    for m in module.modules():
+        if hasattr(m, 'quantize_zero_point'):
+            m.quantize_zero_point = prev_state[m]
 
 
 class MSE(torch.nn.Module):
@@ -673,6 +682,7 @@ def __init__(
         super(HalfQuadraticOptimizerZeroPoint, self).__init__()
         self.hqo_init_op_zp = hqo_init_op_zp
         self.input_view_shape_impl = inner_stats_input_view_shape_impl
+        self.proxy_module = proxy_module
         self.proxy_forward = proxy_module.forward
         self.set_local_loss_mode = lambda enabled: _set_local_loss_mode(proxy_module, enabled)
         self.set_quantize_zero_point = lambda enabled: _set_quantize_zero_point(
@@ -696,10 +706,10 @@ def parameter_search(self, xl, x):
             for i in range(0, self.hqo_iters):
                 self.internal_candidate = candidate
                 self.set_local_loss_mode(True)
-                self.set_quantize_zero_point(False)
+                prev_state = _set_quantize_zero_point(self.proxy_module, False)
                 quant_tensor = self.proxy_forward(x).detach()
                 self.set_local_loss_mode(False)
-                self.set_quantize_zero_point(True)
+                _restore_quantize_zero_point(self.proxy_module, prev_state)
                 loss = torch.abs(quant_tensor.value - x).mean()
                 best_candidate = torch.where(loss < best_loss, candidate, best_candidate)
                 if loss >= best_loss:

diff --git a/src/brevitas_examples/imagenet_classification/ptq/ptq_common.py b/src/brevitas_examples/imagenet_classification/ptq/ptq_common.py
@@ -23,6 +23,7 @@
 from brevitas.inject import value
 import brevitas.nn as qnn
 from brevitas.quant.experimental.float import Fp8e4m3ActPerTensorFloat
+from brevitas.quant.experimental.float import Fp8e4m3ActPerTensorFloatMSE
 from brevitas.quant.experimental.float import Fp8e4m3WeightPerChannelFloat
 from brevitas.quant.experimental.float import Fp8e4m3WeightPerChannelFloatMSE
 from brevitas.quant.experimental.float import Fp8e4m3WeightPerTensorFloat