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

Code before:
            current_min, current_max = torch.min(inputs), torch.max(inputs)
            module.tracked_min_input = update_ema(module.tracked_min_input, current_min,
                                                                       module.ema_decay)
            module.tracked_max_input = update_ema(module.tracked_max_input, current_max,
                                                                       module.ema_decay)
        module.scale, module.zero_point = update_quantization_param(
            input_bits, module.tracked_min_input, module.tracked_max_input)
        inp = self._quantize(input_bits, module, inputs)
        inp = self._dequantize(module, inp)
        return inp

    def quantize_output(self, output, wrapper, **kwargs):
        config = wrapper.config
        module = wrapper.module
        output_bits = get_bits_length(config, 'output')
        module.output_bits = torch.Tensor([output_bits])
        quant_start_step = config.get('quant_start_step', 0)
        assert output_bits >= 1, "quant bits length should be at least 1"

        if quant_start_step > self.bound_model.steps:
            module.tracked_min_output, module.tracked_max_output = torch.min(output), torch.max(output)
            return output

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

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

    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

Review comment (suggestion):
a little confused that we use the same var `module.scale, module.zero_point` in quantizing weight, input, output

Code after:
        module.scale, module.zero_point = update_quantization_param(
            input_bits, module.tracked_min_input, module.tracked_max_input)
        inp = self._quantize(input_bits, module, inputs)
        inp = self._dequantize(module, inp)
        return inp

    def quantize_output(self, output, wrapper, **kwargs):
        config = wrapper.config
        module = wrapper.module
        output_bits = get_bits_length(config, 'output')
        module.output_bits = torch.Tensor([output_bits])
        quant_start_step = config.get('quant_start_step', 0)
        assert output_bits >= 1, "quant bits length should be at least 1"

        if quant_start_step > self.bound_model.steps:
            module.tracked_min_output, module.tracked_max_output = torch.min(output), torch.max(output)
            return output

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

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

    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