chore: import upstream snapshot with attribution

python/paddle/quantization/imperative/__init__.py

@@ -0,0 +1,34 @@
+#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from . import (
+    ptq,  # noqa: F401
+    ptq_config,  # noqa: F401
+    ptq_quantizer,  # noqa: F401
+    ptq_registry,  # noqa: F401
+    qat,  # noqa: F401
+)
+from .ptq import ImperativePTQ  # noqa: F401
+from .ptq_config import PTQConfig, default_ptq_config  # noqa: F401
+from .ptq_quantizer import (  # noqa: F401
+    SUPPORT_ACT_QUANTIZERS,
+    SUPPORT_WT_QUANTIZERS,
+    AbsmaxQuantizer,
+    BaseQuantizer,
+    HistQuantizer,
+    KLQuantizer,
+    PerChannelAbsmaxQuantizer,
+)
+from .ptq_registry import PTQRegistry  # noqa: F401
+from .qat import ImperativeQuantAware  # noqa: F401

python/paddle/quantization/imperative/fuse_utils.py

@@ -0,0 +1,221 @@
+#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+
+import paddle
+from paddle import nn
+
+from . import utils
+
+
+class Identity(nn.Layer):
+    '''a layer to replace bn or relu layers'''
+
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+
+    def forward(self, input):
+        return input
+
+
+def fuse_conv_bn(model):
+    is_train = False
+    if model.training:
+        model.eval()
+        is_train = True
+    fuse_list = []
+    tmp_pair = [None, None]
+    for name, layer in model.named_sublayers():
+        if isinstance(layer, nn.Conv2D):
+            tmp_pair[0] = name
+        if isinstance(layer, nn.BatchNorm2D):
+            tmp_pair[1] = name
+
+        if tmp_pair[0] and tmp_pair[1] and len(tmp_pair) == 2:
+            fuse_list.append(tmp_pair)
+            tmp_pair = [None, None]
+    model = fuse_layers(model, fuse_list)
+    if is_train:
+        model.train()
+
+
+def fuse_layers(model, layers_to_fuse, inplace=False):
+    '''
+    fuse layers in layers_to_fuse
+
+    Args:
+        model(paddle.nn.Layer): The model to be fused.
+        layers_to_fuse(list): The layers' names to be fused. For
+            example,"fuse_list = [["conv1", "bn1"], ["conv2", "bn2"]]".
+            A TypeError would be raised if "fuse" was set as
+            True but "fuse_list" was None.
+                              Default: None.
+        inplace(bool): Whether apply fusing to the input model.
+                       Default: False.
+
+    Return
+        fused_model(paddle.nn.Layer): The fused model.
+    '''
+    if inplace is False:
+        model = copy.deepcopy(model)
+    for layers in layers_to_fuse:
+        _fuse_layers(model, layers)
+    return model
+
+
+def _fuse_layers(model, layers_list):
+    '''fuse all the layers in layers_list'''
+    layer_list = []
+    for layer_name in layers_list:
+        parent_layer, sub_name = utils.find_parent_layer_and_sub_name(
+            model, layer_name
+        )
+        layer_list.append(getattr(parent_layer, sub_name))
+    new_layers = _fuse_func(layer_list)
+    for i, item in enumerate(layers_list):
+        parent_layer, sub_name = utils.find_parent_layer_and_sub_name(
+            model, item
+        )
+        setattr(parent_layer, sub_name, new_layers[i])
+
+
+def _fuse_func(layer_list):
+    '''choose the fuse method and fuse layers'''
+    types = tuple(type(m) for m in layer_list)
+    fusion_method = types_to_fusion_method.get(types, None)
+    new_layers = [None] * len(layer_list)
+    fused_layer = fusion_method(*layer_list)
+    for handle_id, pre_hook_fn in layer_list[0]._forward_pre_hooks.items():
+        fused_layer.register_forward_pre_hook(pre_hook_fn)
+        del layer_list[0]._forward_pre_hooks[handle_id]
+    for handle_id, hook_fn in layer_list[-1]._forward_post_hooks.items():
+        fused_layer.register_forward_post_hook(hook_fn)
+        del layer_list[-1]._forward_post_hooks[handle_id]
+    new_layers[0] = fused_layer
+    for i in range(1, len(layer_list)):
+        identity = Identity()
+        identity.training = layer_list[0].training
+        new_layers[i] = identity
+    return new_layers
+
+
+def _fuse_conv_bn(conv, bn):
+    '''fuse conv and bn for train or eval'''
+    assert conv.training == bn.training, (
+        "Conv and BN both must be in the same mode (train or eval)."
+    )
+    if conv.training:
+        assert bn._num_features == conv._out_channels, (
+            'Output channel of Conv2d must match num_features of BatchNorm2d'
+        )
+        raise NotImplementedError
+    else:
+        return _fuse_conv_bn_eval(conv, bn)
+
+
+def _fuse_conv_bn_eval(conv, bn):
+    '''fuse conv and bn for eval'''
+    assert not (conv.training or bn.training), "Fusion only for eval!"
+    fused_conv = copy.deepcopy(conv)
+
+    fused_weight, fused_bias = _fuse_conv_bn_weights(
+        fused_conv.weight,
+        fused_conv.bias,
+        bn._mean,
+        bn._variance,
+        bn._epsilon,
+        bn.weight,
+        bn.bias,
+    )
+    fused_conv.weight.set_value(fused_weight)
+    if fused_conv.bias is None:
+        fused_conv.bias = paddle.create_parameter(
+            shape=[fused_conv._out_channels], is_bias=True, dtype=bn.bias.dtype
+        )
+    fused_conv.bias.set_value(fused_bias)
+    return fused_conv
+
+
+def _fuse_conv_bn_weights(conv_w, conv_b, bn_rm, bn_rv, bn_eps, bn_w, bn_b):
+    '''fuse weights and bias of conv and bn'''
+    if conv_b is None:
+        conv_b = paddle.zeros_like(bn_rm)
+    if bn_w is None:
+        bn_w = paddle.ones_like(bn_rm)
+    if bn_b is None:
+        bn_b = paddle.zeros_like(bn_rm)
+    bn_var_rsqrt = paddle.rsqrt(bn_rv + bn_eps)
+    conv_w = conv_w * (bn_w * bn_var_rsqrt).reshape(
+        [-1] + [1] * (len(conv_w.shape) - 1)
+    )
+    conv_b = (conv_b - bn_rm) * bn_var_rsqrt * bn_w + bn_b
+    return conv_w, conv_b
+
+
+def _fuse_linear_bn(linear, bn):
+    '''fuse linear and bn'''
+    assert linear.training == bn.training, (
+        "Linear and BN both must be in the same mode (train or eval)."
+    )
+    if linear.training:
+        assert bn._num_features == linear.weight.shape[1], (
+            'Output channel of Linear must match num_features of BatchNorm'
+        )
+        raise NotImplementedError
+    else:
+        return _fuse_linear_bn_eval(linear, bn)
+
+
+def _fuse_linear_bn_eval(linear, bn):
+    '''fuse linear and bn for eval'''
+    assert not (linear.training or bn.training), "Fusion only for eval!"
+    fused_linear = copy.deepcopy(linear)
+
+    fused_weight, fused_bias = _fuse_linear_bn_weights(
+        fused_linear.weight,
+        fused_linear.bias,
+        bn._mean,
+        bn._variance,
+        bn._epsilon,
+        bn.weight,
+        bn.bias,
+    )
+    fused_linear.weight.set_value(fused_weight)
+    if fused_linear.bias is None:
+        fused_linear.bias = paddle.create_parameter(
+            shape=[fused_linear.weight.shape[1]],
+            is_bias=True,
+            dtype=bn.bias.dtype,
+        )
+    fused_linear.bias.set_value(fused_bias)
+    return fused_linear
+
+
+def _fuse_linear_bn_weights(
+    linear_w, linear_b, bn_rm, bn_rv, bn_eps, bn_w, bn_b
+):
+    '''fuse weights and bias of linear and bn'''
+    if linear_b is None:
+        linear_b = paddle.zeros_like(bn_rm)
+    bn_scale = bn_w * paddle.rsqrt(bn_rv + bn_eps)
+    fused_w = linear_w * bn_scale.unsqueeze(-1)
+    fused_b = (linear_b - bn_rm) * bn_scale + bn_b
+    return fused_w, fused_b
+
+
+types_to_fusion_method = {
+    (nn.Conv2D, nn.BatchNorm2D): _fuse_conv_bn,
+    (nn.Linear, nn.BatchNorm1D): _fuse_linear_bn,
+}

python/paddle/quantization/imperative/ptq.py

@@ -0,0 +1,485 @@
+#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+import logging
+import os
+
+import numpy as np
+
+import paddle
+from paddle.nn.quant import quant_layers
+
+from ...static.log_helper import get_logger
+from ...static.quantization.utils import (
+    _get_input_name_index,
+    _get_op_input_var_names,
+    _get_op_output_var_names,
+    _get_output_name_index,
+)
+from . import fuse_utils, ptq_config, ptq_hooks, ptq_quantizer, utils
+from .ptq_registry import PTQRegistry
+
+INFER_MODEL_SUFFIX = ".pdmodel"
+INFER_PARAMS_SUFFIX = ".pdiparams"
+
+_logger = get_logger(
+    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
+)
+
+
+class ImperativePTQ:
+    """
+    Static post training quantization.
+    """
+
+    def __init__(self, quant_config=ptq_config.default_ptq_config):
+        """
+        Constructor.
+
+        Args:
+            quant_config(PTQConfig): the config of post training quantization.
+                The config has weight_quantizer and activation_quantizer.
+                In default, the weight_quantizer is PerChannelAbsmaxQuantizer
+                and the activation_quantizer is KLQuantizer.
+        """
+        super().__init__()
+
+        assert isinstance(quant_config, ptq_config.PTQConfig)
+
+        self._quant_config = quant_config
+
+    def quantize(self, model, inplace=False, fuse=False, fuse_list=None):
+        """
+        Add quant config and hook to the target layer.
+
+        Args:
+            model(paddle.nn.Layer): The model to be quantized.
+            inplace(bool): Whether apply quantization to the input model.
+                           Default: False.
+            fuse(bool): Whether to fuse layers.
+                        Default: False.
+            fuse_list(list): The layers' names to be fused. For example,
+                "fuse_list = [["conv1", "bn1"], ["conv2", "bn2"]]".
+                A TypeError would be raised if "fuse" was set as
+                True but "fuse_list" was None.
+                Default: None.
+        Return
+            quantized_model(paddle.nn.Layer): The quantized model.
+        """
+        assert isinstance(model, paddle.nn.Layer), (
+            "The model must be the instance of paddle.nn.Layer."
+        )
+        if not inplace:
+            model = copy.deepcopy(model)
+        if fuse:
+            model.eval()
+            model = fuse_utils.fuse_layers(model, fuse_list)
+        for name, layer in model.named_sublayers():
+            if (
+                PTQRegistry.is_supported_layer(layer)
+                and utils.is_leaf_layer(layer)
+                and not self._is_skip_layer(layer)
+            ):
+                # Add quant config
+                quant_config = copy.deepcopy(self._quant_config)
+                if PTQRegistry.is_simulated_quant_layer(layer):
+                    quant_config.enable_in_act_quantizer = True
+                layer._quant_config = quant_config
+
+                # register hook
+                hook = ptq_hooks.quant_forward_post_hook
+                quant_hook_handle = layer.register_forward_post_hook(hook)
+                quant_config.quant_hook_handle = quant_hook_handle
+                layer._forward_post_hooks.move_to_end(
+                    quant_hook_handle._hook_id, last=False
+                )
+
+        return model
+
+    def save_quantized_model(self, model, path, input_spec=None, **config):
+        """
+        1. Convert the quantized model
+        2. Call jit.save to save the inference model
+        3. Post process the inference model.
+
+        Args:
+            model (Layer): The model to be saved.
+            path (str): The path prefix to save model. The format is
+                ``dirname/file_prefix`` or ``file_prefix``.
+            input_spec (list[InputSpec|Tensor], optional): Describes the input
+                of the saved model's forward method, which can be described by
+                InputSpec or example Tensor. If None, all input variables of
+                the original Layer's forward method would be the inputs of
+                the saved model. Default None.
+            **config (dict, optional): Other save configuration options for
+                compatibility. We do not recommend using these configurations,
+                they may be removed in the future. If not necessary, DO NOT use
+                them. Default None.
+                The following options are currently supported:
+                (1) output_spec (list[Tensor]): Selects the output targets of
+                the saved model. By default, all return variables of original
+                Layer's forward method are kept as the output of the saved model.
+                If the provided ``output_spec`` list is not all output variables,
+                the saved model will be pruned according to the given
+                ``output_spec`` list.
+
+        Returns:
+            None
+        """
+
+        assert isinstance(model, paddle.nn.Layer), (
+            "The model must be the instance of paddle.nn.Layer."
+        )
+
+        # Convert and save dygraph quantized model
+        self._convert(model)
+
+        paddle.jit.save(layer=model, path=path, input_spec=input_spec, **config)
+
+        # Load inference program
+        is_dynamic_mode = False
+        if paddle.in_dynamic_mode():
+            is_dynamic_mode = True
+            paddle.enable_static()
+
+        place = paddle.CPUPlace()
+        scope = paddle.static.global_scope()
+        exe = paddle.static.Executor(place)
+
+        dirname = os.path.dirname(path)
+        basename = os.path.basename(path)
+        model_filename = basename + INFER_MODEL_SUFFIX
+        params_filename = basename + INFER_PARAMS_SUFFIX
+
+        [
+            infer_program,
+            feed_target_names,
+            fetch_targets,
+        ] = paddle.static.load_inference_model(
+            path_prefix=dirname,
+            executor=exe,
+            model_filename=model_filename,
+            params_filename=params_filename,
+        )
+
+        # Process inference program
+        self._clean_up(infer_program)
+        self._gather_input_thresholds(infer_program, scope)
+        self._remove_scale_op(infer_program)
+
+        # Save final program
+        model_name = None
+        if model_filename is None:
+            model_name = "model"
+        elif model_filename.endswith(".pdmodel"):
+            model_name = model_filename.rsplit(".", 1)[0]
+        else:
+            model_name = model_filename
+        path_prefix = os.path.join(dirname, model_name)
+        feed_vars = [
+            infer_program.global_block().var(name) for name in feed_target_names
+        ]
+        paddle.static.save_inference_model(
+            path_prefix,
+            feed_vars,
+            fetch_targets,
+            executor=exe,
+            program=infer_program.clone(),
+        )
+
+        if is_dynamic_mode:
+            paddle.disable_static()
+
+    def _convert(self, model):
+        """
+        Convert the quantized model.
+
+        Args:
+            model(paddle.nn.Layer): The quantized model.
+            inplace(bool): Whether apply conversion to the input model.
+                           Default: False.
+        Returns:
+            None
+        """
+
+        for name, sub_layer in model.named_sublayers():
+            if self._is_quant_layer(sub_layer):
+                sub_layer._quant_config.quant_hook_handle.remove()
+
+        self._cal_thresholds(model)
+
+        for name, sub_layer in model.named_sublayers():
+            if self._is_quant_layer(sub_layer):
+                self._save_output_thresholds(sub_layer, sub_layer._quant_config)
+
+        self._wrap_simulated_layers(model)
+
+    def _cal_thresholds(self, model):
+        """
+        Calculate the thresholds of inputs and outputs.
+
+        Args:
+            model(paddle.nn.Layer): The quantized model.
+        Returns:
+            None
+        """
+        assert isinstance(model, paddle.nn.Layer), (
+            "The input model must be the instance of paddle.nn.Layer."
+        )
+
+        total_num = 0
+        cur_num = 0
+        for name, sub_layer in model.named_sublayers():
+            if self._is_quant_layer(sub_layer):
+                total_num += 1
+
+        for name, sub_layer in model.named_sublayers():
+            if self._is_quant_layer(sub_layer):
+                cur_num += 1
+                if cur_num % 5 == 0:
+                    _logger.info(f"Process the {cur_num} / {total_num} layer")
+
+                quant_config = sub_layer._quant_config
+
+                if quant_config.enable_in_act_quantizer:
+                    quant_config.in_act_quantizer.cal_thresholds()
+                quant_config.out_act_quantizer.cal_thresholds()
+
+                if PTQRegistry.is_simulated_quant_layer(sub_layer):
+                    weights = (sub_layer.weight,)
+                    quant_config.wt_quantizer.sample_data(sub_layer, weights)
+                    quant_config.wt_quantizer.cal_thresholds()
+
+    def _save_output_thresholds(self, sub_layer, quant_config):
+        """
+        Save the output thresholds to the layer.
+
+        Args:
+            sub_layer(paddle.nn.Layer): The quantized layer.
+            quant_config(PTQConfig): the quant config for the layer.
+        Returns:
+            None
+        """
+        assert isinstance(sub_layer, paddle.nn.Layer), (
+            "The input model must be the instance of paddle.nn.Layer."
+        )
+
+        layer_info = PTQRegistry.layer_info(sub_layer)
+
+        output_names = layer_info.output_names
+        output_thresholds = quant_config.out_act_quantizer.thresholds
+        assert len(output_names) == 1
+        if len(output_thresholds) == 1:
+            save_name = output_names[0] + str(0) + "_threshold"
+            sub_layer._set_op_attrs({save_name: output_thresholds[0]})
+            sub_layer._set_op_attrs({"out_threshold": output_thresholds[0]})
+        else:
+            _logger.warning(
+                f"output_thresholds shape of {output_names[0]} need to be 1, but received {len(output_thresholds)}"
+            )
+
+    def _wrap_simulated_layers(self, model):
+        """
+        Replace conv2d and linear with the quantized layers, and save
+        thresholds into the fake layers.
+        Args:
+            model(paddle.nn.Layer): The model to be quantized.
+        Returns:
+            None
+        """
+        assert isinstance(model, paddle.nn.Layer), (
+            "The input model must be the instance of paddle.nn.Layer."
+        )
+
+        for name, sub_layer in model.named_sublayers():
+            if self._is_quant_layer(
+                sub_layer
+            ) and PTQRegistry.is_simulated_quant_layer(sub_layer):
+                quant_config = sub_layer._quant_config
+                assert quant_config.enable_in_act_quantizer is True
+                wt_quantizer = quant_config.wt_quantizer
+                in_act_quantizer = quant_config.in_act_quantizer
+
+                # create layer
+                quant_layer_name = None
+                for key, value in utils.layer_name_map.items():
+                    if isinstance(sub_layer, value):
+                        quant_layer_name = 'Quantized' + key
+                        break
+                assert quant_layer_name is not None
+
+                if isinstance(wt_quantizer, ptq_quantizer.AbsmaxQuantizer):
+                    weight_quantize_type = "abs_max"
+                else:
+                    weight_quantize_type = "channel_wise_abs_max"
+                kwargs = {
+                    "weight_quantize_type": weight_quantize_type,
+                    "activation_quantize_type": "moving_average_abs_max",
+                    "weight_bits": wt_quantizer.quant_bits,
+                    "activation_bits": in_act_quantizer.quant_bits,
+                }
+
+                quant_layer = quant_layers.__dict__[quant_layer_name](
+                    sub_layer, **kwargs
+                )
+
+                # save the input thresholds
+                assert hasattr(quant_layer, "_fake_quant_input")
+                assert hasattr(quant_layer._fake_quant_input, "_scale")
+                if len(in_act_quantizer.thresholds) == 1:
+                    input_threshold = np.array(
+                        [in_act_quantizer.thresholds[0]], dtype=np.float32
+                    )
+                    quant_layer._fake_quant_input._scale.set_value(
+                        input_threshold
+                    )
+
+                assert hasattr(quant_layer, "_fake_quant_weight")
+                assert hasattr(quant_layer._fake_quant_weight, "_scale")
+                assert len(wt_quantizer.thresholds) == 1
+                weight_threshold = wt_quantizer.thresholds[0]
+                if isinstance(weight_threshold, list):
+                    weight_threshold = np.array(
+                        weight_threshold, dtype=np.float32
+                    )
+                else:
+                    weight_threshold = np.array(
+                        [weight_threshold], dtype=np.float32
+                    )
+                quant_layer._fake_quant_weight._scale.set_value(
+                    weight_threshold
+                )
+
+                # save the output thresholds
+                self._save_output_thresholds(quant_layer, quant_config)
+
+                # replace the layer
+                parent_layer, sub_name = utils.find_parent_layer_and_sub_name(
+                    model, name
+                )
+                setattr(parent_layer, sub_name, quant_layer)
+
+    def _gather_input_thresholds(self, program, scope):
+        """
+        Get and save input thresholds from the front ops.
+
+        Args:
+            program(Program): the input infer program.
+            scope(Scope): the corresponding scope for the program.
+        Returns:
+            None
+        """
+        for op in utils.program_all_ops(program):
+            for in_var_name in _get_op_input_var_names(op):
+                previous_op = utils.find_previous_op(op.block, in_var_name)
+                if previous_op is None:
+                    continue
+
+                if (
+                    "quantize_dequantize" in previous_op.type
+                    or previous_op.type == "moving_average_abs_max_scale"
+                ):
+                    attr_name = previous_op.output('OutScale')[0]
+                    in_threshold = utils.load_variable_data(scope, attr_name)
+                    in_threshold = utils.fp_numpy_to_naive(in_threshold)
+                    argname, index = _get_input_name_index(op, in_var_name)
+                    op._set_attr(
+                        argname + str(index) + "_threshold", in_threshold
+                    )
+                    op._set_attr("with_quant_attr", True)
+                else:
+                    for out_var_name in _get_op_output_var_names(previous_op):
+                        if out_var_name != in_var_name:
+                            continue
+                        argname, index = _get_output_name_index(
+                            previous_op, out_var_name
+                        )
+                        attr_name = argname + str(index) + "_threshold"
+                        if not previous_op.has_attr(attr_name):
+                            continue
+                        threshold = previous_op.attr(attr_name)
+
+                        argname, index = _get_input_name_index(op, in_var_name)
+                        attr_name = argname + str(index) + "_threshold"
+                        op._set_attr(attr_name, threshold)
+                        op._set_attr("with_quant_attr", True)
+
+    def _clean_up(self, program):
+        """
+        Remove useless thresholds which are added in jit.save.
+
+        Args:
+            program(Program): the input infer program.
+        Returns:
+            None
+        """
+
+        def _helper(op, next_op, old_attr_name, new_attr_name):
+            if (
+                op.has_attr(old_attr_name)
+                and next_op.has_attr(old_attr_name)
+                and op.attr(old_attr_name) == next_op.attr(old_attr_name)
+            ):
+                threshold = op.attr(old_attr_name)
+                op._remove_attr(old_attr_name)
+                next_op._remove_attr(old_attr_name)
+                next_op._set_attr(new_attr_name, threshold)
+                next_op._set_attr("with_quant_attr", True)
+
+        for op in utils.program_all_ops(program):
+            if "quantize_dequantize" in op.type:
+                # remove the thresholds in fake ops
+                for attr_name in op.attr_names:
+                    if "_threshold" in attr_name:
+                        op._remove_attr(attr_name)
+            elif op.type in ["conv2d", "matmul"]:
+                # change the thresholds in conv2d/matmul + eleadd
+                arg_name = "Output" if op.type == "conv2d" else "Out"
+                out_var_name = op.output(arg_name)[0]
+                next_ops = utils.find_next_ops(op.block, out_var_name)
+                if len(next_ops) > 1 or next_ops[0].type != "elementwise_add":
+                    continue
+                next_op = next_ops[0]
+
+                argname, index = _get_output_name_index(op, out_var_name)
+                old_attr_name = argname + str(index) + "_threshold"
+
+                argname, index = _get_output_name_index(
+                    next_op, next_op.output("Out")[0]
+                )
+                new_attr_name = argname + str(index) + "_threshold"
+
+                _helper(op, next_op, old_attr_name, new_attr_name)
+                _helper(op, next_op, "out_threshold", "out_threshold")
+
+    def _remove_scale_op(self, program):
+        """
+        Remove the moving_average_abs_max_scale op.
+        """
+        for op in utils.program_all_ops(program):
+            if op.type == "moving_average_abs_max_scale":
+                in_var_name = op.input("X")[0]
+                out_var_name = op.output("Out")[0]
+                next_ops = utils.find_next_ops(op.block, out_var_name)
+                for next_op in next_ops:
+                    next_op._rename_input(out_var_name, in_var_name)
+
+    @staticmethod
+    def _is_skip_layer(layer):
+        return hasattr(layer, "skip_quant") and layer.skip_quant is True
+
+    @staticmethod
+    def _is_quant_layer(layer):
+        return hasattr(layer, "_quant_config")

python/paddle/quantization/imperative/ptq_config.py

@@ -0,0 +1,55 @@
+#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+
+from .ptq_quantizer import (
+    SUPPORT_ACT_QUANTIZERS,
+    SUPPORT_WT_QUANTIZERS,
+    KLQuantizer,
+    PerChannelAbsmaxQuantizer,
+)
+
+
+class PTQConfig:
+    """
+    The PTQ config shows how to quantize the inputs and outputs.
+    """
+
+    def __init__(self, activation_quantizer, weight_quantizer):
+        """
+        Constructor.
+
+        Args:
+            activation_quantizer(BaseQuantizer): The activation quantizer.
+                It should be the instance of BaseQuantizer.
+            weight_quantizer(BaseQuantizer): The weight quantizer.
+                It should be the instance of BaseQuantizer.
+        """
+        super().__init__()
+        assert isinstance(activation_quantizer, tuple(SUPPORT_ACT_QUANTIZERS))
+        assert isinstance(weight_quantizer, tuple(SUPPORT_WT_QUANTIZERS))
+
+        self.in_act_quantizer = copy.deepcopy(activation_quantizer)
+        self.out_act_quantizer = copy.deepcopy(activation_quantizer)
+        self.wt_quantizer = copy.deepcopy(weight_quantizer)
+
+        self.quant_hook_handle = None
+
+        # In order to wrap simulated layers, use in_act_quantizer
+        # to calculate the input thresholds for conv2d, linear and etc.
+        self.enable_in_act_quantizer = False
+
+
+default_ptq_config = PTQConfig(KLQuantizer(), PerChannelAbsmaxQuantizer())

python/paddle/quantization/imperative/ptq_hooks.py

@@ -0,0 +1,27 @@
+#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+def quant_forward_post_hook(layer, inputs, outputs):
+    """
+    The forward_post_hook for PTQ.
+    """
+    assert hasattr(layer, '_quant_config'), (
+        "The layer should have _quant_config attr"
+    )
+
+    qc = layer._quant_config
+    if qc.enable_in_act_quantizer:
+        qc.in_act_quantizer.sample_data(layer, inputs)
+    qc.out_act_quantizer.sample_data(layer, (outputs,))

python/paddle/quantization/imperative/ptq_quantizer.py

@@ -0,0 +1,264 @@
+#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import abc
+import math
+
+import numpy as np
+
+import paddle
+
+from ...static.quantization.cal_kl_threshold import cal_kl_threshold
+from . import utils
+
+
+def abs_max_value(tensor):
+    return float(paddle.max(paddle.abs(tensor)))
+
+
+def merge_max_value(old, new):
+    """
+    Merge the max element one by one in two lists.
+    """
+    assert isinstance(old, list) and isinstance(new, list)
+    if old != []:
+        assert len(old) == len(new)
+        for i in range(len(old)):
+            assert type(old[i]) == type(new[i])
+            if isinstance(old[i], list):
+                new[i] = merge_max_value(old[i], new[i])
+            else:
+                new[i] = max(new[i], old[i])
+    return new
+
+
+def combine_abs_max_and_hist(
+    tensor, origin_max, origin_hist, bins, upsample_bins
+):
+    """ """
+
+    new_max = abs_max_value(tensor)
+
+    if new_max == 0.0:
+        return origin_max, origin_hist
+    elif origin_max == 0.0:
+        new_hist, _ = np.histogram(
+            paddle.abs(tensor).numpy(False), range=(0, new_max), bins=bins
+        )
+        new_hist = new_hist.astype(np.float32)
+        return new_max, new_hist
+    elif new_max <= origin_max:
+        new_hist, _ = np.histogram(
+            paddle.abs(tensor).numpy(False), range=(0, origin_max), bins=bins
+        )
+        new_hist = new_hist.astype(np.float32)
+        new_hist += origin_hist
+        return origin_max, new_hist
+    else:
+        # bin_width = origin_max / (bins * upsample_bins)
+        #           = new_max / (bins * downsample_bins)
+        bin_width = origin_max / (bins * upsample_bins)
+        downsample_bins = int(math.ceil(new_max / (bins * bin_width)))
+        new_max = bins * bin_width * downsample_bins
+
+        upsampled_hist = np.repeat(origin_hist, upsample_bins)
+        expanded_hist = np.zeros((bins * downsample_bins), dtype=np.float32)
+        expanded_hist[0 : bins * upsample_bins] = upsampled_hist
+        cumsumed_hist = np.cumsum(expanded_hist, dtype=np.float64)[
+            downsample_bins - 1 :: downsample_bins
+        ]
+        shift_cumsumed_hist = np.zeros((bins), dtype=np.float64)
+        shift_cumsumed_hist[1:] = cumsumed_hist[0:-1]
+        sampled_hist = (cumsumed_hist - shift_cumsumed_hist) / upsample_bins
+        sampled_hist = sampled_hist.astype(np.float32)
+
+        new_hist, _ = np.histogram(
+            paddle.abs(tensor).numpy(False), range=(0, new_max), bins=bins
+        )
+        new_hist = new_hist.astype(np.float32)
+        new_hist += sampled_hist
+
+        return new_max, new_hist
+
+
+class BaseQuantizer(metaclass=abc.ABCMeta):
+    """
+    Base quantizer for activation and weight.
+    """
+
+    def __init__(self, quant_bits=8):
+        super().__init__()
+        assert isinstance(quant_bits, int)
+        assert quant_bits > 0 and quant_bits <= 16
+
+        self.quant_bits = quant_bits
+
+        self.abs_max_vals = []
+        self.thresholds = []
+
+    @abc.abstractmethod
+    def sample_data(self, layer, tensors):
+        pass
+
+    @abc.abstractmethod
+    def cal_thresholds(self):
+        pass
+
+
+class AbsmaxQuantizer(BaseQuantizer):
+    """
+    Per-tensor abs max quantizer.
+    """
+
+    def __init__(self, quant_bits=8):
+        super().__init__(quant_bits)
+
+    def sample_data(self, layer, tensors):
+        assert isinstance(tensors, tuple)
+
+        abs_max_vals = [abs_max_value(t) for t in tensors]
+        self.abs_max_vals = merge_max_value(self.abs_max_vals, abs_max_vals)
+
+    def cal_thresholds(self):
+        self.thresholds = self.abs_max_vals
+
+
+class PerChannelAbsmaxQuantizer(BaseQuantizer):
+    """
+    Per channel abs max quantizer.
+    """
+
+    def __init__(self, quant_bits=8):
+        super().__init__(quant_bits)
+
+    def sample_data(self, layer, tensors):
+        assert isinstance(layer, paddle.nn.Layer)
+        assert isinstance(tensors, tuple)
+
+        abs_max_vals_list = []
+        for idx, tensor in enumerate(tensors):
+            if isinstance(layer, tuple(utils.spec_channel_axis_layers)):
+                abs_max_vals = [
+                    abs_max_value(tensor[:, i]) for i in range(tensor.shape[1])
+                ]
+                abs_max_vals_list.append(abs_max_vals)
+            else:
+                abs_max_vals = [
+                    abs_max_value(tensor[i]) for i in range(tensor.shape[0])
+                ]
+                abs_max_vals_list.append(abs_max_vals)
+
+        self.abs_max_vals = merge_max_value(
+            self.abs_max_vals, abs_max_vals_list
+        )
+
+    def cal_thresholds(self):
+        self.thresholds = self.abs_max_vals
+
+
+class BaseHistQuantizer(BaseQuantizer, metaclass=abc.ABCMeta):
+    """ """
+
+    def __init__(self, quant_bits=8, bins=1024, upsample_bins=64):
+        super().__init__(quant_bits)
+        self.bins = bins
+        self.upsample_bins = upsample_bins
+
+        self.hists = []
+
+    def sample_data(self, layer, tensors):
+        assert isinstance(tensors, tuple)
+
+        if self.abs_max_vals == []:
+            abs_max_vals = [abs_max_value(t) for t in tensors]
+            self.abs_max_vals = abs_max_vals
+
+            for idx, tensor in enumerate(tensors):
+                if abs_max_vals[idx] == 0.0:
+                    self.hists.append(None)
+                else:
+                    hist, _ = np.histogram(
+                        paddle.abs(tensor).numpy(False),
+                        range=(0.0, abs_max_vals[idx]),
+                        bins=self.bins,
+                    )
+                    hist = hist.astype(np.float32)
+                    self.hists.append(hist)
+        else:
+            assert len(self.abs_max_vals) == len(tensors)
+            assert len(self.hists) == len(tensors)
+
+            for idx, tensor in enumerate(tensors):
+                new_abs_max, new_hist = combine_abs_max_and_hist(
+                    tensor,
+                    self.abs_max_vals[idx],
+                    self.hists[idx],
+                    self.bins,
+                    self.upsample_bins,
+                )
+                self.abs_max_vals[idx] = new_abs_max
+                self.hists[idx] = new_hist
+
+    @abc.abstractmethod
+    def cal_thresholds(self):
+        pass
+
+
+class HistQuantizer(BaseHistQuantizer):
+    """ """
+
+    def __init__(
+        self, quant_bits=8, bins=1024, upsample_bins=64, hist_percent=0.99999
+    ):
+        super().__init__(quant_bits, bins, upsample_bins)
+        self.hist_percent = hist_percent
+
+    def cal_thresholds(self):
+        def _helper(abs_max, hist, percent):
+            assert hist.ndim == 1 and percent < 1.0
+            hist = hist / np.sum(hist, dtype=np.float64)
+            cumsumed_hist = np.cumsum(hist)
+            index = np.argwhere(cumsumed_hist >= percent)[0]
+            return float((index - 0.5) * (abs_max / hist.shape[0]))
+
+        for idx in range(len(self.hists)):
+            if self.hists[idx] is None:
+                self.thresholds.append(self.abs_max_vals[idx])
+            else:
+                threshold = _helper(
+                    self.abs_max_vals[idx], self.hists[idx], self.hist_percent
+                )
+                self.thresholds.append(threshold)
+
+
+class KLQuantizer(BaseHistQuantizer):
+    """ """
+
+    def __init__(self, quant_bits=8, bins=1024, upsample_bins=64):
+        super().__init__(quant_bits, bins, upsample_bins)
+
+    def cal_thresholds(self):
+        for idx in range(len(self.hists)):
+            if self.hists[idx] is None:
+                self.thresholds.append(self.abs_max_vals[idx])
+            else:
+                hist = self.hists[idx]
+                abs_max_val = self.abs_max_vals[idx]
+                bin_width = abs_max_val / hist.shape[0]
+                threshold = cal_kl_threshold(hist, bin_width, self.quant_bits)
+                self.thresholds.append(threshold)
+
+
+SUPPORT_ACT_QUANTIZERS = [AbsmaxQuantizer, HistQuantizer, KLQuantizer]
+SUPPORT_WT_QUANTIZERS = [AbsmaxQuantizer, PerChannelAbsmaxQuantizer]

python/paddle/quantization/imperative/ptq_registry.py

@@ -0,0 +1,143 @@
+#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import paddle
+
+
+class LayerInfo:
+    """
+    Store the arg names of the inputs and outputs.
+    """
+
+    def __init__(self, layer, input_names, weight_names, output_names):
+        super().__init__()
+        self.layer = layer
+        self.input_names = input_names
+        self.weight_names = weight_names
+        self.output_names = output_names
+
+
+PTQ_LAYERS_INFO = [
+    LayerInfo(paddle.nn.Conv2D, ['Input'], ['Filter'], ['Output']),
+    LayerInfo(paddle.nn.Linear, ['X'], ['Y'], ['Out']),
+    LayerInfo(paddle.nn.BatchNorm2D, ['X'], [], ['Y']),
+    LayerInfo(paddle.nn.AdaptiveMaxPool2D, ['X'], [], ['Out']),
+    LayerInfo(paddle.nn.AdaptiveAvgPool2D, ['X'], [], ['Out']),
+    LayerInfo(paddle.nn.AvgPool2D, ['X'], [], ['Out']),
+    LayerInfo(paddle.nn.MaxPool2D, ['X'], [], ['Out']),
+    LayerInfo(paddle.nn.ReLU, ['X'], [], ['Out']),
+    LayerInfo(paddle.nn.ReLU6, ['X'], [], ['Out']),
+    LayerInfo(paddle.nn.Hardswish, ['X'], [], ['Out']),
+    LayerInfo(paddle.nn.Swish, ['X'], [], ['Out']),
+    LayerInfo(paddle.nn.Sigmoid, ['X'], [], ['Out']),
+    LayerInfo(paddle.nn.Softmax, ['X'], [], ['Out']),
+    LayerInfo(paddle.nn.Tanh, ['X'], [], ['Out']),
+    LayerInfo(paddle.nn.quant.add, ['X', 'Y'], [], ['Out']),
+]
+
+QUANT_LAYERS_INFO = [
+    LayerInfo(
+        paddle.nn.quant.quant_layers.QuantizedConv2D,
+        ['Input'],
+        ['Filter'],
+        ['Output'],
+    ),
+    LayerInfo(
+        paddle.nn.quant.quant_layers.QuantizedLinear, ['X'], ['Y'], ['Out']
+    ),
+]
+
+SIMULATED_LAYERS = [paddle.nn.Conv2D, paddle.nn.Linear]
+
+
+class PTQRegistry:
+    """
+    Register the supported layers for PTQ and provide layers info.
+    """
+
+    supported_layers_map = {}
+    registered_layers_map = {}
+    is_inited = False
+
+    def __init__(self):
+        super().__init__()
+
+    @classmethod
+    def _init(cls):
+        if not cls.is_inited:
+            for layer_info in PTQ_LAYERS_INFO:
+                cls.supported_layers_map[layer_info.layer] = layer_info
+
+            all_layers_info = PTQ_LAYERS_INFO + QUANT_LAYERS_INFO
+            for layer_info in all_layers_info:
+                cls.registered_layers_map[layer_info.layer] = layer_info
+        cls.is_inited = True
+
+    @classmethod
+    def is_supported_layer(cls, layer):
+        """
+        Analyze whether the layer supports quantization.
+        Args:
+            layer(Layer): The input layer can be a python class or an instance.
+        Returns:
+            flag(bool): Whether the layer is supported.
+        """
+        cls._init()
+        return layer in cls.supported_layers_map or isinstance(
+            layer, tuple(cls.supported_layers_map.keys())
+        )
+
+    @classmethod
+    def is_registered_layer(cls, layer):
+        """
+        Analyze whether the layer is register layer_info.
+        Args:
+            layer(Layer): The input layer can be a python class or an instance.
+        Returns:
+            flag(bool): Whether the layer is register layer_info.
+        """
+        cls._init()
+        return layer in cls.registered_layers_map or isinstance(
+            layer, tuple(cls.registered_layers_map.keys())
+        )
+
+    @classmethod
+    def is_simulated_quant_layer(cls, layer):
+        """
+        Analyze whether the layer is simulated quant layer.
+        Args:
+            layer(Layer): The input layer can be a python class or an instance.
+        Returns:
+            flag(bool): Whether the layer is supported.
+        """
+        return layer in SIMULATED_LAYERS or isinstance(
+            layer, tuple(SIMULATED_LAYERS)
+        )
+
+    @classmethod
+    def layer_info(cls, layer):
+        """
+        Get the information for the layer.
+        Args:
+            layer(Layer): The input layer can be a python class or an instance.
+        Returns:
+            layer_info(LayerInfo): The layer info of the input layer.
+        """
+        assert cls.is_registered_layer(layer), (
+            "The input layer is not register."
+        )
+
+        for layer_key, layer_info in cls.registered_layers_map.items():
+            if layer == layer_key or isinstance(layer, layer_key):
+                return layer_info

python/paddle/quantization/imperative/qat.py

@@ -0,0 +1,765 @@
+#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+import paddle
+from paddle.base.framework import IrGraph
+from paddle.framework import core
+from paddle.nn.quant import quant_layers
+
+from ...static.quantization.quantization_pass import (
+    QuantWeightPass,
+    ReplaceFakeQuantDequantPass,
+)
+from ...static.quantization.utils import (
+    _get_input_name_index,
+    _get_op_input_var_names,
+    _get_output_name_index,
+    move_persistable_var_to_global_block,
+)
+from . import fuse_utils, utils
+
+INFER_MODEL_SUFFIX = ".pdmodel"
+INFER_PARAMS_SUFFIX = ".pdiparams"
+
+
+def lazy_import_fleet(layer_name_map, fake_quant_input_layers):
+    from paddle.distributed import fleet
+
+    layer_name_map['ColumnParallelLinear'] = (
+        fleet.meta_parallel.parallel_layers.mp_layers.ColumnParallelLinear
+    )
+    layer_name_map['RowParallelLinear'] = (
+        fleet.meta_parallel.parallel_layers.mp_layers.RowParallelLinear
+    )
+    fake_quant_input_layers.append(fleet.meta_parallel.RowParallelLinear)
+    fake_quant_input_layers.append(fleet.meta_parallel.ColumnParallelLinear)
+    return layer_name_map, fake_quant_input_layers
+
+
+class ImperativeQuantAware:
+    """
+    Applying quantization aware training (QAT) to the dygraph model.
+    """
+
+    def __init__(
+        self,
+        quantizable_layer_type=[
+            'Conv2D',
+            'Linear',
+            'Conv2DTranspose',
+            'ColumnParallelLinear',
+            'RowParallelLinear',
+        ],
+        weight_quantize_type='abs_max',
+        activation_quantize_type='moving_average_abs_max',
+        weight_bits=8,
+        activation_bits=8,
+        moving_rate=0.9,
+        fuse_conv_bn=False,
+        weight_preprocess_layer=None,
+        act_preprocess_layer=None,
+        weight_quantize_layer=None,
+        act_quantize_layer=None,
+        onnx_format=False,
+    ):
+        """
+        The constructor for ImperativeQuantAware.
+
+        Args:
+            quantizable_layer_type(list[str | layer]): List the type of
+                layers that will be quantized. Default is ['Conv2D', 'Linear'].
+            weight_quantize_type(str): quantization type for weights,
+                which supports 'abs_max' and 'channel_wise_abs_max'.
+            activation_quantize_type(str): quantization type for activations,
+                which supports 'abs_max' and 'moving_average_abs_max' now.
+                If using 'abs_max' mode, the quantization scale will be
+                calculated dynamically each step in both training and testing
+                period. If using 'moving_average_abs_max', the static
+                quantization scale will be calculated during training and
+                used in inference.
+            weight_bits(int): quantization bit number for weights, whereas
+                the bias is not quantized.
+            activation_bits(int): quantization bit number for activations.
+            moving_rate(float): the parameter for 'moving_average_abs_max'
+                quantization.
+            fuse_conv_bn(bool): Whether to fuse conv and bn, default is False.
+            weight_preprocess_layer(paddle.nn.Layer, optional): A paddle
+                Layer that defines how to preprocess weight before quantization.
+                Using this can quickly test if user's preprocess method works
+                or not. The input is non-quantized weight and function returns
+                processed weight to be quantized.
+                If None, the weight will be quantized directly.
+                Default is None.
+            act_preprocess_layer(paddle.nn.Layer, optional): A paddle Layer
+                that defines how to preprocess activation before quantization.
+                Using this can quickly test if user's preprocess method works
+                or not. The input is non-quantized activation and function returns
+                processed activation to be quantized.
+                If None, the activation will be quantized directly.
+                Default is None.
+            weight_quantize_layer(paddle.nn.Layer, optional): A paddle Layer that
+                defines how to quantize weight.
+                Using this can quickly test if user's quantization method works or not.
+                In this layer, user should both define quantization method and
+                dequantization method, that is, the function's input is non-quantized
+                weight and returns dequantized weight.
+                If None, will use quantization op defined by 'weight_quantize_type'.
+                Default is None.
+            act_quantize_layer(paddle.nn.Layer, optional): A paddle Layer that defines
+                how to quantize activation.
+                Using this can quickly test if user's quantization method works or not.
+                In this layer, user should both define quantization method and
+                dequantization method, that is, the function's input is non-quantized
+                activation and returns dequantized activation.
+                If None, will use quantization op defined by 'activation_quantize_type'.
+                Default is None.
+            onnx_format (bool, optional): Whether to export the quantized model
+                with format of ONNX. Default is False.
+
+        Note:
+            If user sets attribute 'skip_quant' to a Layer that support dynamic
+            quantization and sets it to true, the layer would not be quantized
+            during training. If this attribute is not sets or the attribute is
+            false, the Layer would be quantized in training.
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> from paddle.static.quantization import (
+                ...     ImperativeQuantAware,
+                ... )
+                >>> from paddle.vision.models import (
+                ...     resnet,
+                ... )
+
+                >>> model = resnet.resnet50(pretrained=True)
+
+                >>> imperative_qat = ImperativeQuantAware(
+                ...     weight_quantize_type='abs_max',
+                ...     activation_quantize_type='moving_average_abs_max',
+                ... )
+
+                >>> # Add the fake quant logical.
+                >>> # The original model will be rewrite.
+                >>> # The outscale of outputs in supported layers would be calculated.
+                >>> imperative_qat.quantize(model)
+
+                >>> # Fine-tune the quantized model
+                >>> # ...
+
+                >>> # Save quant model for the inference.
+                >>> imperative_qat.save_quantized_model(
+                ...     layer=model,
+                ...     model_path="./resnet50_qat",
+                ...     input_spec=[
+                ...         paddle.static.InputSpec(shape=[None, 3, 224, 224], dtype='float32'),
+                ...     ],
+                ... )
+
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> from paddle.static.quantization import (
+                ...     ImperativeQuantAware,
+                ... )
+
+                >>> class ImperativeModel(paddle.nn.Layer):
+                ...     def __init__(self):
+                ...         super().__init__()
+                ...         # self.linear_0 would skip the quantization.
+                ...         self.linear_0 = paddle.nn.Linear(784, 400)
+                ...         self.linear_0.skip_quant = True
+
+                ...         # self.linear_1 would not skip the quantization.
+                ...         self.linear_1 = paddle.nn.Linear(400, 10)
+                ...         self.linear_1.skip_quant = False
+
+                ...     def forward(self, inputs):
+                ...         x = self.linear_0(inputs)
+                ...         x = self.linear_1(inputs)
+                ...         return x
+
+                >>> model = ImperativeModel()
+                >>> imperative_qat = ImperativeQuantAware(
+                ...     weight_quantize_type='abs_max',
+                ...     activation_quantize_type='moving_average_abs_max',
+                ... )
+
+                >>> # Add the fake quant logical.
+                >>> # The original model will be rewrite.
+                >>> #
+                >>> # There is only one Layer(self.linear1) would be added the
+                >>> # fake quant logical.
+                >>> imperative_qat.quantize(model)
+
+                >>> # Fine-tune the quantized model
+                >>> # ...
+
+                >>> # Save quant model for the inference.
+                >>> imperative_qat.save_quantized_model(
+                ...     layer=model,
+                ...     model_path="./imperative_model_qat",
+                ... )
+        """
+        super().__init__()
+        self.fuse_conv_bn = fuse_conv_bn
+
+        kwargs = {
+            "quantizable_layer_type": quantizable_layer_type,
+            "weight_quantize_type": weight_quantize_type,
+            "activation_quantize_type": activation_quantize_type,
+            "weight_bits": weight_bits,
+            "activation_bits": activation_bits,
+            "moving_rate": moving_rate,
+            "weight_preprocess_layer": weight_preprocess_layer,
+            "act_preprocess_layer": act_preprocess_layer,
+            "weight_quantize_layer": weight_quantize_layer,
+            "act_quantize_layer": act_quantize_layer,
+        }
+
+        self._quantize_inputs = ImperativeQuantizeInputs(**kwargs)
+
+        self._quantize_outputs = ImperativeQuantizeOutputs(
+            moving_rate, activation_bits, onnx_format
+        )
+
+    def quantize(self, model):
+        """
+        According to weights' and activations' quantization types,
+        the model will be added some fake quant ops, such as
+        fake_quantize_dequantize_moving_average_abs_max,
+        fake_quantize_dequantize_abs_max and so on. At the same time,
+        the out_scale value of outputs would be calculated.
+
+        Args:
+            model(paddle.nn.Layer): the model to be quantized.
+        Returns:
+            None
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> from paddle.static.quantization import (
+                ...     ImperativeQuantAware,
+                ... )
+
+                >>> class ImperativeModel(paddle.nn.Layer):
+                ...     def __init__(self):
+                ...         super().__init__()
+                ...         # self.linear_0 would skip the quantization.
+                ...         self.linear_0 = paddle.nn.Linear(784, 400)
+                ...         self.linear_0.skip_quant = True
+
+                ...         # self.linear_1 would not skip the quantization.
+                ...         self.linear_1 = paddle.nn.Linear(400, 10)
+                ...         self.linear_1.skip_quant = False
+
+                ...     def forward(self, inputs):
+                ...         x = self.linear_0(inputs)
+                ...         x = self.linear_1(inputs)
+                ...         return x
+
+                >>> model = ImperativeModel()
+                >>> imperative_qat = ImperativeQuantAware(
+                ...     weight_quantize_type='abs_max',
+                ...     activation_quantize_type='moving_average_abs_max',
+                ... )
+
+                >>> # Add the fake quant logical.
+                >>> # The original model will be rewrite.
+                >>> #
+                >>> # There is only one Layer(self.linear1) would be added the
+                >>> # fake quant logical.
+                >>> imperative_qat.quantize(model)
+        """
+        assert isinstance(model, paddle.nn.Layer), (
+            "The model must be the instance of paddle.nn.Layer."
+        )
+
+        if self.fuse_conv_bn:
+            fuse_utils.fuse_conv_bn(model)
+
+        self._quantize_inputs.apply(model)
+        self._quantize_outputs.apply(model)
+        return model
+
+    def save_quantized_model(self, layer, path, input_spec=None, **config):
+        with paddle.pir_utils.OldIrGuard():
+            self._quantize_outputs.save_quantized_model(
+                layer, path, input_spec, **config
+            )
+
+
+class ImperativeQuantizeInputs:
+    """
+    Based on the input params, add the quant_dequant computational
+    logic both for activation inputs and weight inputs.
+    """
+
+    def __init__(
+        self,
+        quantizable_layer_type=['Conv2D', 'Linear', 'Conv2DTranspose'],
+        weight_quantize_type='abs_max',
+        activation_quantize_type='moving_average_abs_max',
+        weight_bits=8,
+        activation_bits=8,
+        moving_rate=0.9,
+        weight_preprocess_layer=None,
+        act_preprocess_layer=None,
+        weight_quantize_layer=None,
+        act_quantize_layer=None,
+    ):
+        """
+        The constructor for ImperativeQuantizeInputs.
+
+        Please refer to the args of ImperativeQuantAware.
+        """
+        super().__init__()
+        self.layer_name_map, self.fake_quant_input_layers = lazy_import_fleet(
+            utils.layer_name_map, utils.fake_quant_input_layers
+        )
+
+        self._quantizable_layer_type = tuple(
+            (
+                self.layer_name_map[layer]
+                if layer in self.layer_name_map
+                else layer
+            )
+            for layer in quantizable_layer_type
+        )
+        for layer in self._quantizable_layer_type:
+            assert (
+                not isinstance(layer, str)
+                and layer in self.fake_quant_input_layers
+            ), f"{layer} is unsupported to be quantized."
+
+        quantize_type = {
+            'abs_max',
+            'moving_average_abs_max',
+            'channel_wise_abs_max',
+            'lsq_weight',
+            'channel_wise_lsq_weight',
+        }
+        act_quantize_type = {'moving_average_abs_max', 'lsq_act'}
+        assert (
+            weight_quantize_type != 'moving_average_abs_max'
+            and weight_quantize_type in quantize_type
+        ), (
+            f"Unsupported weight_quantize_type: {weight_quantize_type}. It can only "
+            "be abs_max or channel_wise_abs_max."
+        )
+        # TODO (jc): activation_quantize_type supports range_abs_max
+        assert activation_quantize_type in act_quantize_type, (
+            f"Unsupported activation_quantize_type: {activation_quantize_type}. It can "
+            "only be moving_average_abs_max or lsq_act now."
+        )
+
+        bits_check = lambda bits: (
+            isinstance(bits, int) and bits >= 0 and bits <= 16
+        )
+        assert bits_check(weight_bits), "weight_bits should be 1, 2,... or 16."
+        assert bits_check(activation_bits), (
+            "activation_bits should be 1, 2,... or 16."
+        )
+
+        layer_check = lambda method: (
+            method is None or issubclass(method, paddle.nn.Layer)
+        )
+        assert layer_check(weight_preprocess_layer), (
+            "weight_preprocess should be nn.Layer."
+        )
+        assert layer_check(act_preprocess_layer), (
+            "act_preprocess should be nn.Layer."
+        )
+        assert layer_check(weight_quantize_layer), (
+            "weight_quantize should be nn.Layer."
+        )
+        assert layer_check(act_quantize_layer), (
+            "act_quantize should be nn.Layer."
+        )
+
+        self._kwargs = {
+            "weight_quantize_type": weight_quantize_type,
+            "activation_quantize_type": activation_quantize_type,
+            "weight_bits": weight_bits,
+            "activation_bits": activation_bits,
+            "moving_rate": moving_rate,
+            "weight_pre_layer": weight_preprocess_layer,
+            "act_pre_layer": act_preprocess_layer,
+            "weight_quant_layer": weight_quantize_layer,
+            "act_quant_layer": act_quantize_layer,
+        }
+
+    def apply(self, model):
+        """
+        Quantize the weights and activations to calculate for specific
+        layers.
+
+        Args:
+            model(paddle.nn.Layer): The target model which would
+                calculate the input quantization scale.
+
+        Returns:
+            None
+        """
+
+        assert isinstance(model, paddle.nn.Layer), (
+            "The model must be the instance of paddle.nn.Layer."
+        )
+
+        for name, cur_layer in model.named_sublayers():
+            if not isinstance(cur_layer, self._quantizable_layer_type) or (
+                hasattr(cur_layer, "skip_quant")
+                and cur_layer.skip_quant is True
+            ):
+                continue
+
+            parent_layer, sub_name = utils.find_parent_layer_and_sub_name(
+                model, name
+            )
+
+            cur_quant_layer = self._get_input_quantized_layer(cur_layer)
+            setattr(parent_layer, sub_name, cur_quant_layer)
+
+    def _get_input_quantized_layer(self, layer):
+        quant_layer_name = None
+
+        for key, value in self.layer_name_map.items():
+            if isinstance(layer, value):
+                quant_layer_name = 'Quantized' + key
+                break
+        assert quant_layer_name is not None, (
+            f"The layer {layer.full_name()} is unsupported to be quantized."
+        )
+
+        return quant_layers.__dict__[quant_layer_name](layer, **self._kwargs)
+
+
+class ImperativeQuantizeOutputs:
+    """
+    Calculate the output scales for target layers.
+    """
+
+    def __init__(self, moving_rate=0.9, activation_bits=8, onnx_format=False):
+        """
+        The constructor for ImperativeQuantizeOutputs.
+
+        Args:
+            moving_rate(float): The decay coefficient of moving average.
+                                The default value is 0.9.
+            activation_bits(int, optional): quantization bit number for activation. Default is 8.
+        """
+        super().__init__()
+        self._moving_rate = moving_rate
+        self._activation_bits = activation_bits
+        self._onnx_format = onnx_format
+
+    def apply(self, model):
+        """
+        Insert the `moving_average_abs_max_scale` layers to calculate the
+        output scales for specific layers in the dygraph model.
+
+        Args:
+            model(paddle.nn.Layer): The target model which would be
+                calculate the output quantization scale.
+
+        Returns:
+            None
+        """
+        assert isinstance(model, paddle.nn.Layer), (
+            "The model must be the instance of paddle.nn.Layer."
+        )
+
+        for cur_name, cur_layer in model.named_sublayers():
+            if '_act_preprocess' in cur_name:
+                continue
+            if not self._is_target_layer(cur_layer):
+                continue
+
+            parent_layer, sub_name = utils.find_parent_layer_and_sub_name(
+                model, cur_name
+            )
+
+            reduce_type = None
+
+            if isinstance(cur_layer, tuple(utils.fake_quant_output_layers)):
+                cur_quant_layer = quant_layers.FakeQuantMAOutputScaleLayer(
+                    cur_layer, self._moving_rate, reduce_type=reduce_type
+                )
+            else:
+                cur_quant_layer = quant_layers.MAOutputScaleLayer(
+                    cur_layer, self._moving_rate, reduce_type=reduce_type
+                )
+
+            setattr(parent_layer, sub_name, cur_quant_layer)
+
+    def save_quantized_model(self, model, path, input_spec=None, **config):
+        """
+        Save the quantized model for the inference.
+
+        Args:
+            model (Layer): The model to be saved.
+            path (str): The path prefix to save model. The format is
+                ``dirname/file_prefix`` or ``file_prefix``.
+            input_spec (list[InputSpec|Tensor], optional): Describes the input
+                of the saved model's forward method, which can be described by
+                InputSpec or example Tensor. If None, all input variables of
+                the original Layer's forward method would be the inputs of
+                the saved model. Default None.
+            **config (dict, optional): Other save configuration options for
+                compatibility. We do not recommend using these configurations,
+                they may be removed in the future. If not necessary, DO NOT use
+                them. Default None.
+                The following options are currently supported:
+                (1) output_spec (list[Tensor]): Selects the output targets of
+                the saved model. By default, all return variables of original
+                Layer's forward method are kept as the output of the saved model.
+                If the provided ``output_spec`` list is not all output variables,
+                the saved model will be pruned according to the given
+                ``output_spec`` list.
+
+        Returns:
+            None
+        """
+        assert isinstance(model, paddle.nn.Layer), (
+            "The model must be the instance of paddle.nn.Layer."
+        )
+
+        if input_spec:
+            paddle.jit.to_static(model, input_spec=input_spec)
+        paddle.jit.save(layer=model, path=path, input_spec=input_spec, **config)
+
+        is_dynamic_mode = False
+        if paddle.in_dynamic_mode():
+            is_dynamic_mode = True
+            paddle.enable_static()
+
+        place = core.CPUPlace()
+        scope = paddle.static.global_scope()
+        exe = paddle.static.Executor(place)
+
+        dirname = os.path.dirname(path)
+        basename = os.path.basename(path)
+        model_filename = basename + INFER_MODEL_SUFFIX
+        params_filename = basename + INFER_PARAMS_SUFFIX
+
+        [
+            infer_program,
+            feed_target_names,
+            fetch_targets,
+        ] = paddle.static.load_inference_model(
+            dirname,
+            executor=exe,
+            model_filename=model_filename,
+            params_filename=params_filename,
+        )
+
+        if not self._onnx_format:
+            self._gather_scales(infer_program, scope, fetch_targets)
+
+            # Remove `moving_average_abs_max_scale` node in sub graphs.
+            graph = IrGraph(core.Graph(infer_program.desc), for_test=False)
+            for sub_graph in graph.all_sub_graphs():
+                for _op in sub_graph.all_op_nodes():
+                    if _op.name() == "moving_average_abs_max_scale":
+                        sub_graph.safe_remove_nodes(_op)
+                sub_graph.resolve_hazard()
+            infer_program = graph.to_program()
+
+            self._set_skip_quant_attr(infer_program)
+
+            clip_extra = False
+        else:
+            graph = IrGraph(core.Graph(infer_program.desc), for_test=False)
+            transform_pass = ReplaceFakeQuantDequantPass(
+                scope, place, quant_bits=self._activation_bits
+            )
+            for sub_graph in graph.all_sub_graphs():
+                sub_graph._for_test = True
+                transform_pass.apply(sub_graph)
+
+            quant_weight_pass = QuantWeightPass(scope, place)
+            for sub_graph in graph.all_sub_graphs():
+                sub_graph._for_test = True
+                quant_weight_pass.apply(sub_graph)
+
+            infer_program = graph.to_program()
+
+            clip_extra = True
+
+        move_persistable_var_to_global_block(infer_program)
+
+        model_name = None
+        if model_filename is None:
+            model_name = "model"
+        elif model_filename.endswith(".pdmodel"):
+            model_name = model_filename.rsplit(".", 1)[0]
+        else:
+            model_name = model_filename
+        path_prefix = os.path.join(dirname, model_name)
+        feed_vars = [
+            infer_program.global_block().var(name) for name in feed_target_names
+        ]
+        paddle.static.save_inference_model(
+            path_prefix,
+            feed_vars,
+            fetch_targets,
+            executor=exe,
+            program=infer_program.clone(),
+            clip_extra=clip_extra,
+        )
+
+        if is_dynamic_mode:
+            paddle.disable_static()
+
+    def _is_target_layer(self, layer):
+        """
+        Whether the layer needs to calculate output scales.
+        """
+        # exclude fake_quant ops in quant_layers file
+        if not isinstance(layer, paddle.nn.Layer):
+            return False
+
+        if self._onnx_format:
+            return (
+                True
+                if isinstance(layer, tuple(utils.fake_quant_wrap_layers))
+                else False
+            )
+
+        flag = False
+        if utils.is_leaf_layer(layer) and not isinstance(
+            layer, tuple(utils.fake_quant_leaf_layers)
+        ):
+            flag = True
+
+        if isinstance(layer, tuple(utils.fake_quant_wrap_layers)):
+            flag = True
+
+        if isinstance(layer, paddle.nn.quant.FloatFunctionalLayer):
+            flag = True
+
+        return flag
+
+    def _gather_scales(self, program, scope, fetch_targets):
+        """
+        Get all scales from fake ops, save them into the corresponding ops
+        and delete all moving_average_abs_max_scale ops.
+        """
+
+        def _gather_input_scale():
+            target_ops = []
+            skip_ops = [
+                *utils.fake_quantize_dequantize_op_types,
+                "moving_average_abs_max_scale",
+            ]
+            for block in program.blocks:
+                for op in block.ops:
+                    if op.type not in skip_ops:
+                        target_ops.append(op)
+
+            for op in target_ops:
+                for in_var_name in _get_op_input_var_names(op):
+                    previous_op = utils.find_previous_op(op.block, in_var_name)
+
+                    if previous_op is not None and (
+                        "quantize_dequantize" in previous_op.type
+                        or previous_op.type == "moving_average_abs_max_scale"
+                    ):
+                        scale_name = previous_op.output('OutScale')[0]
+                        in_scale = utils.load_variable_data(scope, scale_name)
+                        in_scale = utils.fp_numpy_to_naive(in_scale)
+                        argname, index = _get_input_name_index(op, in_var_name)
+                        op._set_attr(
+                            argname + str(index) + "_threshold", in_scale
+                        )
+                        op._set_attr("with_quant_attr", True)
+
+        def _gather_output_scale():
+            target_ops = []
+            for block in program.blocks:
+                for op in block.ops:
+                    if op.type == "moving_average_abs_max_scale":
+                        target_ops.append(op)
+
+            for op in target_ops:
+                in_var_name = op.input('X')[0]
+                out_var_name = op.output('Out')[0]
+                block = op.block
+                previous_op = utils.find_previous_op(block, in_var_name)
+                next_ops = utils.find_next_ops(block, out_var_name)
+
+                out_scale_name = op.output('OutScale')[0]
+                out_scale = utils.load_variable_data(scope, out_scale_name)
+                out_scale = utils.fp_numpy_to_naive(out_scale)
+
+                if previous_op.type != "feed":
+                    res = _get_output_name_index(previous_op, in_var_name)
+                    if res is not None:
+                        argname, index = res
+                        previous_op._set_attr(
+                            argname + str(index) + "_threshold", out_scale
+                        )
+                        previous_op._set_attr("out_threshold", out_scale)
+                        previous_op._set_attr("with_quant_attr", True)
+
+                for next_op in next_ops:
+                    next_op._rename_input(out_var_name, in_var_name)
+                    # If next_op is `fetch` and out_var_name in fetch_targets,
+                    # fetch_targets must update to in_var_name when rename input.
+                    for i in range(len(fetch_targets)):
+                        if fetch_targets[i].name == out_var_name:
+                            fetch_targets[i] = block.var(in_var_name)
+
+        _gather_input_scale()
+        _gather_output_scale()
+
+    def _set_skip_quant_attr(self, program):
+        """
+        Label the skip quantized ops.
+        """
+        for block in program.blocks:
+            for op in block.ops:
+                if self._is_skip_quant_op(block, op):
+                    op._set_attr("skip_quant", True)
+                    op._set_attr("with_quant_attr", True)
+
+    def _is_skip_quant_op(self, block, in_op):
+        """
+        The input op should be skipped quantization.
+        1. the type of input op should be conv2d, depthwise_conv2d or matmul
+        2. the previous ops of the input op are not fake_quantize_dequantize ops
+        """
+        target_op_types = [
+            "conv2d",
+            "depthwise_conv2d",
+            "matmul",
+            "conv2d_transpose",
+        ]
+        if in_op.type not in target_op_types:
+            return False
+
+        previous_ops = [
+            utils.find_previous_op(block, arg_name)
+            for arg_name in in_op.input_arg_names
+        ]
+        return any(
+            op is not None
+            and op.type not in utils.fake_quantize_dequantize_op_types
+            for op in previous_ops
+        )

python/paddle/quantization/imperative/utils.py

@@ -0,0 +1,180 @@
+#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+
+import paddle
+from paddle.nn.quant import quant_layers
+
+layer_name_map = {
+    'Conv2DTranspose': paddle.nn.Conv2DTranspose,
+    'Conv2D': paddle.nn.Conv2D,
+    'Linear': paddle.nn.Linear,
+    'AdaptiveAvgPool2D': paddle.nn.AdaptiveAvgPool2D,
+    'AdaptiveMaxPool2D': paddle.nn.AdaptiveMaxPool2D,
+    'AvgPool2D': paddle.nn.AvgPool2D,
+    'MaxPool2D': paddle.nn.MaxPool2D,
+    'Hardswish': paddle.nn.Hardswish,
+    'LeakyReLU': paddle.nn.LeakyReLU,
+    'PReLU': paddle.nn.PReLU,
+    'ReLU': paddle.nn.ReLU,
+    'ReLU6': paddle.nn.ReLU6,
+    'Sigmoid': paddle.nn.Sigmoid,
+    'Softmax': paddle.nn.Softmax,
+    'Swish': paddle.nn.Swish,
+    'Tanh': paddle.nn.Tanh,
+    'BatchNorm': paddle.nn.BatchNorm,
+    'GroupNorm': paddle.nn.GroupNorm,
+    'LayerNorm': paddle.nn.LayerNorm,
+}
+
+# Apply fake quant for the inputs of these layers
+fake_quant_input_layers = [
+    paddle.nn.Conv2D,
+    paddle.nn.Linear,
+    paddle.nn.Conv2DTranspose,
+]
+
+# Apply fake quant for the output of these layers
+# TODO(jc): fix the problem of adding duplicate fake_quant ops
+# paddle.nn.AdaptiveAvgPool2D, paddle.nn.AvgPool2D, paddle.nn.ReLU,paddle.nn.LeakyReLU
+fake_quant_output_layers = [
+    paddle.nn.quant.add,
+    paddle.nn.quant.subtract,
+    paddle.nn.quant.multiply,
+    paddle.nn.quant.divide,
+    paddle.nn.quant.matmul,
+]
+
+fake_quant_leaf_layers = [
+    quant_layers.FakeQuantAbsMax,
+    quant_layers.FakeQuantChannelWiseAbsMax,
+    quant_layers.FakeQuantMovingAverageAbsMax,
+    quant_layers.MovingAverageAbsMaxScale,
+]
+
+fake_quant_wrap_layers = [
+    quant_layers.QuantizedConv2D,
+    quant_layers.QuantizedLinear,
+    quant_layers.QuantizedConv2DTranspose,
+    quant_layers.QuantizedColumnParallelLinear,
+    quant_layers.QuantizedRowParallelLinear,
+]
+
+# The weight format of these layers is Cin * Cout * H * W
+spec_channel_axis_layers = [paddle.nn.Conv2DTranspose, paddle.nn.Linear]
+
+weight_op_types = [
+    "conv2d",
+    "depthwise_conv2d",
+    "matmul",
+    "conv2d_transpose",
+    "depthwise_conv2d_transpose",
+]
+
+fake_quantize_dequantize_op_types = [
+    "fake_quantize_dequantize_abs_max",
+    "fake_channel_wise_quantize_dequantize_abs_max",
+    "fake_quantize_dequantize_moving_average_abs_max",
+]
+
+
+def load_variable_data(scope, var_name):
+    """
+    Load variable value from scope
+    """
+    var_node = scope.find_var(var_name)
+    assert var_node is not None, "Can not find " + var_name + " in the scope."
+    return np.array(var_node.get_tensor())
+
+
+def find_previous_op(block, var_name):
+    """
+    Find the previous op for the input variable.
+    """
+    for op in block.ops:
+        if var_name in op.output_arg_names:
+            return op
+    return None
+
+
+def find_next_ops(block, var_name):
+    """
+    Find all followed ops for the input variable.
+    """
+    res_ops = []
+    for op in block.ops:
+        if var_name in op.input_arg_names:
+            res_ops.append(op)
+    return res_ops
+
+
+def find_parent_layer_and_sub_name(model, name):
+    """
+    Given the model and the name of a layer, find the parent layer and
+    the sub_name of the layer.
+    For example, if name is 'block_1/convbn_1/conv_1', the parent layer is
+    'block_1/convbn_1' and the sub_name is `conv_1`.
+    Args:
+        model(paddle.nn.Layer): the model to be quantized.
+        name(string): the name of a layer
+
+    Returns:
+        parent_layer, subname
+    """
+    assert isinstance(model, paddle.nn.Layer), (
+        "The model must be the instance of paddle.nn.Layer."
+    )
+    assert len(name) > 0, "The input (name) should not be empty."
+
+    last_idx = 0
+    idx = 0
+    parent_layer = model
+    while idx < len(name):
+        if name[idx] == '.':
+            sub_name = name[last_idx:idx]
+            if hasattr(parent_layer, sub_name):
+                parent_layer = getattr(parent_layer, sub_name)
+                last_idx = idx + 1
+        idx += 1
+    sub_name = name[last_idx:idx]
+    return parent_layer, sub_name
+
+
+def program_all_ops(program):
+    """
+    Return all ops for the input program.
+    """
+    all_ops = []
+    for block in program.blocks:
+        for op in block.ops:
+            all_ops.append(op)
+    return all_ops
+
+
+def is_leaf_layer(layer):
+    """
+    Whether the layer is leaf layer.
+    """
+    return isinstance(layer, paddle.nn.Layer) and len(layer.sublayers()) == 0
+
+
+def fp_numpy_to_naive(x_np):
+    """
+    Convert numpy to float or list.
+    """
+    if x_np.size == 1:
+        return float(x_np)
+    else:
+        return x_np.tolist()