Repository: microsoft/nni
File: nni/algorithms/compression/pytorch/quantization/quantizers.py (Python)

Code before:
        return out

    def load_calibration_config(self, calibration_config):
        for name, module in self.bound_model.named_modules():
            if name not in calibration_config:
                if hasattr(module, 'weight_bits') or hasattr(module, 'output_bits'):
                    logger.warning(f"Can not find module {name}'s parameter in input config.")
                continue
            if hasattr(module, 'weight_bits'):
                assert calibration_config[name]['weight_bits'] == module.weight_bits, f"weight bits of module {name} fail to match"
            if hasattr(module, 'input_bits'):
                assert calibration_config[name]['input_bits'] == module.input_bits, f"input bits of module {name} fail to match"
                module.tracked_min_input = calibration_config[name]['tracked_min_input']
                module.tracked_max_input = calibration_config[name]['tracked_max_input']
            if hasattr(module, 'output_bits'):
                assert calibration_config[name]['output_bits'] == module.output_bits, f"output bits of module {name} fail to match"
                module.tracked_min_output = calibration_config[name]['tracked_min_output']
                module.tracked_max_output = calibration_config[name]['tracked_max_output']

    def export_model(self, model_path, calibration_path=None, onnx_path=None, input_shape=None, device=None):
        """
        Export quantized model weights and calibration parameters(optional)

        Parameters
        ----------
        model_path : str
            path to save quantized model weight
        calibration_path : str
            (optional) path to save quantize parameters after calibration
        onnx_path : str
            (optional) path to save onnx model
        input_shape : list or tuple
            input shape to onnx model
        device : torch.device
            device of the model, used to place the dummy input tensor for exporting onnx file.
            the tensor is placed on cpu if ```device``` is None

        Returns
        -------
        Dict
        """
        assert model_path is not None, 'model_path must be specified'
        self._unwrap_model()
        calibration_config = {}

        for name, module in self.bound_model.named_modules():
            if hasattr(module, 'weight_bits') or hasattr(module, 'output_bits'):
                calibration_config[name] = {}
            if hasattr(module, 'weight_bits'):
                calibration_config[name]['weight_bits'] = int(module.weight_bits)


Review comment (suggestion):
I think it is better to use 
```
modules_to_compress = self.get_modules_to_compress()
for layer, config in modules_to_compress:
   load calibration config...
```

Code after:
        if quant_start_step > int(self.bound_model.steps):
            current_min, current_max = torch.min(output), torch.max(output)
            module.tracked_min_output.copy_(current_min)
            module.tracked_max_output.copy_(current_max)
            return output

        # we dont update output quantization parameters in evaluation stage
        if wrapper.training:
            current_min, current_max = torch.min(output), torch.max(output)
            tracked_min_output = update_ema(module.tracked_min_output, current_min,
                                            module.ema_decay)
            tracked_max_output = update_ema(module.tracked_max_output, current_max,
                                            module.ema_decay)
            module.tracked_min_output.copy_(tracked_min_output)
            module.tracked_max_output.copy_(tracked_max_output)

        scale, zero_point = update_quantization_param(
            output_bits, module.tracked_min_output, module.tracked_max_output)
        module.scale.copy_(scale)
        module.zero_point.copy_(zero_point)

        out = self._quantize(output_bits, module, output)
        out = self._dequantize(module, out)
        return out

    def load_calibration_config(self, calibration_config):
        modules_to_compress = self.get_modules_to_compress()
        for layer, _ in modules_to_compress:
            name, module = layer.name, layer.module
            if name not in calibration_config:
                if hasattr(module, 'weight_bits') or hasattr(module, 'output_bits') or hasattr(module, 'input_bits'):
                    logger.warning(f"Can not find module {name}'s parameter in input config.")
                continue
            if hasattr(module, 'weight_bits'):
                assert calibration_config[name]['weight_bits'] == module.weight_bits, f"weight bits of module {name} fail to match"
            if hasattr(module, 'input_bits'):
                assert calibration_config[name]['input_bits'] == module.input_bits, f"input bits of module {name} fail to match"
                module.tracked_min_input.data = torch.Tensor([calibration_config[name]['tracked_min_input']])
                module.tracked_max_input.data = torch.Tensor([calibration_config[name]['tracked_max_input']])
            if hasattr(module, 'output_bits'):
                assert calibration_config[name]['output_bits'] == module.output_bits, f"output bits of module {name} fail to match"
                module.tracked_min_output.data = torch.Tensor([calibration_config[name]['tracked_min_output']])
                module.tracked_max_output.data = torch.Tensor([calibration_config[name]['tracked_max_output']])

    def export_model(self, model_path, calibration_path=None, onnx_path=None, input_shape=None, device=None):
        """
        Export quantized model weights and calibration parameters(optional)

        Parameters
        ----------
        model_path : str