black

Author: Jiang-Stan

sparsebit/quantization/quant_model.py

@@ -185,7 +185,9 @@ def prepare_calibration(self):
         from sparsebit.quantization.tools.calibration import CalibrationRunner
 
         self.eval()
-        self.calibration_runner = CalibrationRunner(self.model, self.cfg.SCHEDULE.BIAS_CORRECTION)
+        self.calibration_runner = CalibrationRunner(
+            self.model, self.cfg.SCHEDULE.BIAS_CORRECTION
+        )
         self.calibration_runner.prepare_calibration()
 
     def calc_qparams(self):

sparsebit/quantization/tools/calibration.py

@@ -106,27 +106,54 @@ def feature_layerwise_calibration(self, device):
             self.builder.storage.set_output(node.target, outputs)
 
             if self.bias_correction:
-                if isinstance(module, QuantOpr) and getattr(module, "weight_quantizer", None):
+                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)
+                        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)
+                                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()
+                            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)