add bias correction

add bias correction

rectify

inprogressing

resnet18 ok

resnet18 ok

deit ok

deit ok

Author: Jiang-Stan

examples/post_training_quantization/imagenet1k/deit/qconfig.yaml

@@ -15,4 +15,8 @@ A:
     BIT: 8
   OBSERVER:
     TYPE: MINMAX
-    LAYOUT: NCHW
+    LAYOUT: NLC
+  SPECIFIC: [{
+    "patch_embed_proj": ["OBSERVER.LAYOUT", "NCHW"],
+    "head": ["OBSERVER.LAYOUT", "NCHW"],
+  }]

sparsebit/quantization/quant_config.py

@@ -9,6 +9,7 @@
 
 _C.SCHEDULE = CN()
 _C.SCHEDULE.FUSE_BN = False  # use ``with torch.no_grad()`` if it's enabled
+_C.SCHEDULE.BIAS_CORRECTION = False
 _C.SCHEDULE.BN_TUNING = False
 _C.SCHEDULE.DISABLE_UNNECESSARY_QUANT = True
 

sparsebit/quantization/quant_model.py

@@ -185,13 +185,13 @@ def prepare_calibration(self):
         from sparsebit.quantization.tools.calibration import CalibrationRunner
 
         self.eval()
-        self.calibration_runner = CalibrationRunner(self.model)
+        self.calibration_runner = CalibrationRunner(self.model, self.cfg.SCHEDULE.BIAS_CORRECTION)
         self.calibration_runner.prepare_calibration()
 
     def calc_qparams(self):
         assert hasattr(self, "calibration_runner"), "run self.prepare_calibration first"
-        self.calibration_runner.feature_layerwise_calibration(self.device)
         self.calibration_runner.weight_calibration()
+        self.calibration_runner.feature_layerwise_calibration(self.device)
         del self.calibration_runner
 
     def init_QAT(self):

sparsebit/quantization/tools/calibration.py

@@ -1,4 +1,5 @@
 import torch
+import torch.nn as nn
 from functools import partial
 
 from sparsebit.quantization.modules import QuantOpr
@@ -7,8 +8,9 @@
 
 
 class CalibrationRunner(object):
-    def __init__(self, model):
+    def __init__(self, model, bias_correction=False):
         self.model = fx_symbolic_trace(model)
+        self.bias_correction = bias_correction
 
     def prepare_calibration(self):
         input_names_cache = set(
@@ -102,6 +104,30 @@ def feature_layerwise_calibration(self, device):
                     # more time for less cuda memory occupation
                     outputs.append(to_cpu(module(*args, **kwargs)))
             self.builder.storage.set_output(node.target, outputs)
+
+            if self.bias_correction:
+                if isinstance(module, QuantOpr) and getattr(module, "weight_quantizer", None):
+                    for inp_node in node.all_input_nodes:
+                        inp_tensors = self.builder.storage.get_output(inp_node.target)
+                        float_outputs = torch.Tensor([])
+                        quant_outputs = torch.Tensor([])
+                        float_outputs_cached = self.builder.storage.get_output(node.target)
+                        for idx in range(batch_num):
+                            inp_tensor = inp_tensors[idx].cuda()
+                            with torch.no_grad():
+                                float_output = float_outputs_cached[idx]
+                                module.set_quant(True, False)
+                                quant_output = module(inp_tensor).cpu()
+                                module.set_quant(False, False)
+                                float_outputs = torch.cat((float_outputs, float_output.detach()), 0)
+                                quant_outputs = torch.cat((quant_outputs, quant_output.detach()), 0)
+                        float_output_mean = float_outputs.transpose(module.input_quantizer.qdesc._ch_axis,0).flatten(1).mean(-1)
+                        quant_output_mean = quant_outputs.transpose(module.input_quantizer.qdesc._ch_axis,0).flatten(1).mean(-1)
+                        bias = quant_output_mean - float_output_mean
+                        if module.bias is None:
+                            module.bias = nn.Parameter(data=torch.zeros(module.weight.size(0), dtype=torch.float32, device=device), requires_grad=False)
+                        module.bias.data = module.bias.data-bias.cuda()
+
             self.builder.storage.finish_node(node.target)
 
     def weight_calibration(self):