[Fix] mse observer uses per_channel affine (ModelTC#150)

(cherry picked from commit d4abf393b017ecb9b470a9824d2a5c0397f082e8)

Co-authored-by: fanyunqian <[email protected]>

mqbench/fake_quantize/adaround_quantizer.py

@@ -1,11 +1,9 @@
 import torch
 from torch.nn.parameter import Parameter
 
-from mqbench.fake_quantize.quantize_base import QuantizeBase
+from mqbench.fake_quantize.quantize_base import QuantizeBase, _version_under_1100 
 from mqbench.utils.hook import PerChannelLoadHook
 
-_version_under_1100 = int(torch.__version__.split('.')[1]) < 10
-
 def _rectified_sigmoid(alpha, zeta, gamma):
     """Function to generate rounding mask.
 

mqbench/fake_quantize/quantize_base.py

@@ -5,6 +5,7 @@
 
 from mqbench.utils import is_symmetric_quant
 
+_version_under_1100 = int(torch.__version__.split('.')[1]) < 10
 
 class QuantizeBase(FakeQuantizeBase):
     r""" This is an extension of the FakeQuantize module in fake_quantize.py, which

mqbench/observer.py

@@ -4,6 +4,7 @@
 import torch
 from torch.quantization.observer import _ObserverBase
 
+from mqbench.fake_quantize.quantize_base import _version_under_1100 
 from mqbench.utils import sync_tensor, pot_quantization, is_symmetric_quant
 from mqbench.utils.logger import logger
 from mqbench.utils.hook import PerChannelLoadHook
@@ -523,7 +524,7 @@ def mse_perchannel(self, x: torch.Tensor, x_min: torch.Tensor, x_max: torch.Tens
             new_max = x_max * (1.0 - (i * 0.01))
             scale, zero_point = self._calculate_qparams(new_min, new_max)
             x_q = torch.fake_quantize_per_channel_affine(
-                x, scale, zero_point.long(), ch_axis, 
+                x, scale, zero_point.long() if _version_under_1100 else zero_point, ch_axis, 
                 self.quant_min, self.quant_max)
             score = self.lp_loss(x_q, x, reduce_dim)
             update_idx = (score < best_score)
@@ -602,7 +603,7 @@ def mse_perchannel(self, x: torch.Tensor, x_min: torch.Tensor, x_max: torch.Tens
             new_max = x_max * (1.0 - (i * 0.01))
             scale, zero_point = self._calculate_qparams(new_min, new_max)
             x_q = torch.fake_quantize_per_channel_affine(
-                x, scale, zero_point.long(), ch_axis, 
+                x, scale, zero_point.long() if _version_under_1100 else zero_point, ch_axis, 
                 self.quant_min, self.quant_max)
             score = self.lp_loss(x_q, x, reduce_dim)
             update_idx = (score < best_score)