[Refactor] Re-write the (EMA)MSEObserver to accelerate the calibration

phase.

Author: fanyunqian

mqbench/observer.py

@@ -528,11 +528,12 @@ def __init__(self, dtype=torch.quint8, qscheme=torch.per_tensor_affine,
                                           ch_axis, pot_scale, factory_kwargs)
         self.p = p
 
-    def lp_loss(self, pred, tgt):
+    def lp_loss(self, pred, tgt, dim=None):
         """
         loss function measured in L_p Norm
         """
-        return (pred - tgt).abs().pow(self.p).mean()
+        return (pred - tgt).abs().pow(self.p).mean(dim) if dim else (pred - tgt).abs().pow(self.p).mean()
+
 
     def mse(self, x: torch.Tensor, x_min: torch.Tensor, x_max: torch.Tensor, iter=80):
         best_score = 1e+10
@@ -552,6 +553,26 @@ def mse(self, x: torch.Tensor, x_min: torch.Tensor, x_max: torch.Tensor, iter=80
                 best_min, best_max = new_min, new_max
         return best_min, best_max
 
+    def mse_perchannel(self, x: torch.Tensor, x_min: torch.Tensor, x_max: torch.Tensor, iter=80, ch_axis=0):
+        assert x_min.shape == x_max.shape
+        assert ch_axis >= 0, f'{ch_axis}'
+        best_score = 1e+10 * torch.ones_like(x_min)
+        best_min, best_max = x_min.clone(), x_max.clone()
+        reduce_dim = tuple([i for i in range(len(x.shape)) if i != ch_axis])
+        for i in range(iter):
+            new_min = x_min * (1.0 - (i * 0.01))
+            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, 
+                self.quant_min, self.quant_max)
+            score = self.lp_loss(x_q, x, reduce_dim)
+            update_idx = (score < best_score)
+            best_score[update_idx] = score[update_idx]
+            best_min[update_idx] = new_min[update_idx]
+            best_max[update_idx] = new_max[update_idx]
+        return best_min, best_max
+
     def forward(self, x_orig):
         r"""Records the running minimum and maximum of ``x``."""
         if x_orig.numel() == 0:
@@ -568,8 +589,7 @@ def forward(self, x_orig):
             x_channel = x.permute(new_axis_list)
             y = torch.flatten(x_channel, start_dim=1)
             min_val_cur, max_val_cur = torch._aminmax(y, 1)
-            for ch, val in enumerate(min_val_cur):
-                min_val_cur[ch], max_val_cur[ch] = self.mse(x_channel[ch], min_val_cur[ch], max_val_cur[ch], iter=80)
+            min_val_cur, max_val_cur = self.mse_perchannel(x, min_val_cur, max_val_cur, iter=80, ch_axis=self.ch_axis)
 
         self.min_val = torch.min(self.min_val, min_val_cur)
         self.max_val = torch.max(self.max_val, max_val_cur)
@@ -588,11 +608,11 @@ def __init__(self, dtype=torch.quint8, qscheme=torch.per_tensor_affine,
         self.ema_ratio = ema_ratio
         self.p = p
 
-    def lp_loss(self, pred, tgt):
+    def lp_loss(self, pred, tgt, dim=None):
         """
         loss function measured in L_p Norm
         """
-        return (pred - tgt).abs().pow(self.p).mean()
+        return (pred - tgt).abs().pow(self.p).mean(dim) if dim else (pred - tgt).abs().pow(self.p).mean()
 
     def mse(self, x: torch.Tensor, x_min: torch.Tensor, x_max: torch.Tensor, iter=80):
         best_score = 1e+10
@@ -612,6 +632,26 @@ def mse(self, x: torch.Tensor, x_min: torch.Tensor, x_max: torch.Tensor, iter=80
                 best_min, best_max = new_min, new_max
         return best_min, best_max
 
+    def mse_perchannel(self, x: torch.Tensor, x_min: torch.Tensor, x_max: torch.Tensor, iter=80, ch_axis=0):
+        assert x_min.shape == x_max.shape
+        assert ch_axis >= 0, f'{ch_axis}'
+        best_score = 1e+10 * torch.ones_like(x_min)
+        best_min, best_max = x_min.clone(), x_max.clone()
+        reduce_dim = tuple([i for i in range(len(x.shape)) if i != ch_axis])
+        for i in range(iter):
+            new_min = x_min * (1.0 - (i * 0.01))
+            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, 
+                self.quant_min, self.quant_max)
+            score = self.lp_loss(x_q, x, reduce_dim)
+            update_idx = (score < best_score)
+            best_score[update_idx] = score[update_idx]
+            best_min[update_idx] = new_min[update_idx]
+            best_max[update_idx] = new_max[update_idx]
+        return best_min, best_max
+
     def forward(self, x_orig):
         r"""Records the running minimum and maximum of ``x``."""
         if x_orig.numel() == 0:
@@ -628,9 +668,7 @@ def forward(self, x_orig):
             x_channel = x.permute(new_axis_list)
             y = torch.flatten(x_channel, start_dim=1)
             min_val_cur, max_val_cur = torch._aminmax(y, 1)
-            for ch, val in enumerate(min_val_cur):
-                min_val_cur[ch], max_val_cur[ch] = self.mse(x_channel[ch], min_val_cur[ch],
-                                                            max_val_cur[ch], iter=80)
+            min_val_cur, max_val_cur = self.mse_perchannel(x, min_val_cur, max_val_cur, iter=80, ch_axis=self.ch_axis)
 
         if self.max_val.numel() <= 1 and self.max_val.isinf():
             self.min_val = min_val_cur