chore: import upstream snapshot with attribution

python/paddle/static/quantization/__init__.py

@@ -0,0 +1,45 @@
+# 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 .post_training_quantization import (  # noqa: F401
+    PostTrainingQuantization,
+    PostTrainingQuantizationProgram,
+    WeightQuantization,
+)
+from .quant2_int8_onednn_pass import (  # noqa: F401
+    Quant2Int8MkldnnPass,
+    Quant2Int8OnednnPass,
+)
+from .quant_int8_onednn_pass import (  # noqa: F401
+    QuantInt8MkldnnPass,
+    QuantInt8OnednnPass,
+)
+from .quanter import (  # noqa: F401
+    convert,
+    quant_aware,
+)
+from .quantization_pass import (  # noqa: F401
+    AddQuantDequantForInferencePass,
+    AddQuantDequantPass,
+    AddQuantDequantPassV2,
+    ConvertToInt8Pass,
+    OutScaleForInferencePass,
+    OutScaleForTrainingPass,
+    QuantizationFreezePass,
+    QuantizationTransformPass,
+    QuantizationTransformPassV2,
+    QuantWeightPass,
+    ReplaceFakeQuantDequantPass,
+    TransformForMobilePass,
+)

python/paddle/static/quantization/adaround.py

@@ -0,0 +1,377 @@
+#   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 logging
+import sys
+import time
+
+import numpy as np
+
+import paddle
+from paddle import static
+
+from ..log_helper import get_logger
+from .utils import (
+    _channelwise_quant_axis1_ops,
+    bias_correction_w,
+    calculate_quant_cos_error,
+    dequant_tensor,
+    load_variable_data,
+    quant_tensor,
+    set_variable_data,
+    stable_sigmoid,
+)
+
+_logger = get_logger(
+    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
+)
+
+GAMMA = -0.1
+ZETA = 1.1
+
+
+def compute_soft_rounding(alpha_v):
+    return paddle.clip(
+        paddle.nn.functional.sigmoid(alpha_v) * (ZETA - GAMMA) + GAMMA,
+        min=0,
+        max=1,
+    )
+
+
+def compute_soft_rounding_np(alpha_v):
+    return np.clip(
+        stable_sigmoid(alpha_v) * (ZETA - GAMMA) + GAMMA, a_min=0, a_max=1
+    )
+
+
+class AdaRoundLoss:
+    def __init__(self, reg_param=0.01, default_beta_range=(20, 2)):
+        self.default_reg_param = reg_param
+        self.default_beta_range = default_beta_range
+
+    def compute_recon_loss(self, ada_quantized_output, orig_output):
+        square_cost = paddle.nn.functional.square_error_cost(
+            ada_quantized_output, orig_output
+        )
+        recon_loss = paddle.mean(paddle.sum(square_cost, axis=-1))
+        return recon_loss
+
+    def compute_round_loss(self, alpha_v, warm_start, beta):
+        def round_loss_fn():
+            # compute rectified sigmoid of parameter 'alpha' which maps it between zero and one
+            h_v = compute_soft_rounding(alpha_v)
+
+            # calculate regularization term - which ensures parameter to converge to exactly zeros and ones
+            # at the end of optimization
+            reg_term = paddle.sum(
+                -paddle.pow(paddle.abs(2 * h_v - 1), beta) + 1
+            )
+
+            # calculate the rounding loss
+            round_loss = self.default_reg_param * reg_term
+
+            return round_loss
+
+        round_loss = static.nn.cond(
+            warm_start,
+            lambda: paddle.full(shape=[1], dtype='float32', fill_value=0.0),
+            round_loss_fn,
+        )
+
+        return round_loss
+
+    def compute_beta(self, max_iter, cur_iter, warm_start):
+        #  Start and stop beta for annealing of rounding loss (start_beta, end_beta)
+        start_beta, end_beta = self.default_beta_range
+
+        # iteration at end of warm start period, which is 20% of max iterations
+        warm_start_end_iter = warm_start * max_iter
+
+        # compute relative iteration of current iteration
+        rel_iter = (cur_iter - warm_start_end_iter) / (
+            max_iter - warm_start_end_iter
+        )
+        beta = end_beta + 0.5 * (start_beta - end_beta) * (
+            1 + np.cos(rel_iter * np.pi)
+        )
+
+        return beta
+
+
+class AdaRound:
+    def __init__(
+        self,
+        scale,
+        weight_tensor,
+        scope=None,
+        weight_var_name=None,
+        weight_op_type=None,
+        is_train=True,
+        num_iterations=1000,
+    ):
+        self.is_train = is_train
+        self.num_iterations = num_iterations
+        self.warm_start = 0.1
+        self.weight_bits = 8
+        self.offset = 0.0  # zero-point offset
+        self.adaround_loss = AdaRoundLoss()
+        self.ori_weight_tensor = weight_tensor
+        self.scale = scale
+        self.scope = scope
+        self.quant_axis = 0
+        if weight_op_type in _channelwise_quant_axis1_ops:
+            self.quant_axis = 1
+        self.weight_var_name = weight_var_name
+        self.alpha_name = weight_var_name + ".alpha"
+        self.initialize_alpha(weight_tensor.copy(), scale, weight_var_name)
+
+    def initialize_alpha(self, tensor, scale, var_name):
+        """
+        Initializes alpha parameter, same shape as the weight tensor
+        """
+        tensor_scale = quant_tensor(tensor, scale, quant_axis=self.quant_axis)
+        tensor_floor = np.floor(tensor_scale)
+        tensor = tensor_scale - tensor_floor
+        alpha = -np.log((ZETA - GAMMA) / (tensor - GAMMA) - 1)
+        self.alpha_v = paddle.create_parameter(
+            shape=alpha.shape,
+            dtype="float32",
+            name=var_name + ".alpha",
+            default_initializer=paddle.nn.initializer.Assign(alpha),
+        )
+
+    def _calculate_output_with_adarounded_weights(
+        self, program, place, exe, data, fp32_fetch_list, weight_tensor_dequant
+    ):
+        set_variable_data(
+            self.scope, place, self.weight_var_name, weight_tensor_dequant
+        )
+
+        adaround_out_tensor = exe.run(
+            program=program,
+            feed=data,
+            fetch_list=[fp32_fetch_list],
+            return_numpy=True,
+            scope=self.scope,
+        )
+        return adaround_out_tensor
+
+    def _calculate_quant_weight(self):
+        np_alpha = load_variable_data(self.scope, self.alpha_name)
+        h_alpha = compute_soft_rounding_np(np_alpha)
+
+        # Scale the tensor
+        tensor_scale = quant_tensor(
+            self.ori_weight_tensor.copy(),
+            self.scale,
+            quant_axis=self.quant_axis,
+        )
+
+        weight_tensor = np.floor(tensor_scale)
+
+        # Adaround the tensor
+        weight_tensor_quant = np.add(weight_tensor, h_alpha)
+        return weight_tensor_quant
+
+    def _calculate_adarounded_weights(self):
+        weight_tensor_quant = self._calculate_quant_weight()
+
+        # Dequantize the tensor
+        weight_tensor_dequant = dequant_tensor(
+            weight_tensor_quant + self.offset,
+            self.scale,
+            quant_axis=self.quant_axis,
+        )
+        return weight_tensor_dequant
+
+    def update_final_weights(self):
+        weight_tensor_quant = self._calculate_quant_weight()
+        return weight_tensor_quant
+
+    def get_loss(self, beta, warm_start, adaround_out_tensor, orig_out_tensor):
+        round_loss = self.adaround_loss.compute_round_loss(
+            self.alpha_v, warm_start, beta
+        )
+        recon_loss = self.adaround_loss.compute_recon_loss(
+            adaround_out_tensor, orig_out_tensor
+        )
+        loss = round_loss + recon_loss
+        losses = {
+            'loss': loss,
+            'round_loss': round_loss,
+            'recon_loss': recon_loss,
+        }
+        return losses
+
+    def update_beta_warm(self, cur_iteration):
+        warm_start = cur_iteration < self.num_iterations * self.warm_start
+        beta = self.adaround_loss.compute_beta(
+            self.num_iterations, cur_iteration, self.warm_start
+        )
+        return beta, warm_start
+
+
+def run_adaround(
+    data_loader,
+    fp32_program,
+    fetch_list,
+    exe,
+    scope,
+    place,
+    quantized_op_pairs,
+    weight_op_pairs,
+    scale_dict,
+    num_iterations=1000,
+    lr=0.001,
+    bias_correction=False,
+    fast_mode=True,
+):
+    fetch_op_name = fetch_list[0].name
+    final_weight_tensor_quant_dict = {}
+    for weight_var_name, quant_op_out_name in quantized_op_pairs.items():
+        _logger.info(f'Start adaround op: {weight_var_name}')
+        weight_op_type = weight_op_pairs[weight_var_name]
+        # get scale and weight tensor
+        weight_var_tensor = load_variable_data(scope, weight_var_name)
+        scale = scale_dict[weight_var_name]
+        fp32_fetch_list = None
+        for _op in fp32_program.global_block().ops:
+            if _op.type == "fetch":
+                _op._rename_input(fetch_op_name, quant_op_out_name)
+                fp32_fetch_list = fp32_program.global_block().var(
+                    quant_op_out_name
+                )
+                fetch_op_name = quant_op_out_name
+
+        # build adaround program
+        startup_program = static.Program()
+        train_program = static.Program()
+        with (
+            static.program_guard(train_program, startup_program),
+            paddle.utils.unique_name.guard(),
+        ):
+            # initialize adaround
+            adaround = AdaRound(
+                scale,
+                weight_var_tensor,
+                scope=scope,
+                weight_var_name=weight_var_name,
+                weight_op_type=weight_op_type,
+                num_iterations=num_iterations,
+            )
+            orig_out_tensor = static.data(
+                name='orig_out_tensor',
+                shape=(-1, *fp32_fetch_list.shape),
+                dtype='float32',
+            )
+            adaround_out_tensor = static.data(
+                name='adaround_out_tensor',
+                shape=(-1, *fp32_fetch_list.shape),
+                dtype='float32',
+            )
+            beta_tensor = static.data(
+                name='beta', shape=[-1, 1], dtype='float32'
+            )
+            warm_start_tensor = static.data(
+                name='warm_start', shape=[-1, 1], dtype='bool'
+            )
+
+            train_fetches_loss = adaround.get_loss(
+                beta_tensor,
+                warm_start_tensor,
+                adaround_out_tensor,
+                orig_out_tensor,
+            )
+            optimizer = paddle.optimizer.Adam(learning_rate=lr)
+            loss = train_fetches_loss['loss']
+            optimizer.minimize(loss)
+        exe.run(startup_program)
+
+        start_time = time.time()
+        prev_start_time = start_time
+        for i, data in enumerate(data_loader()):
+            prev_start_time = start_time
+            start_time = time.time()
+            # run fp32 model
+            np_orig_out_tensor = exe.run(
+                program=fp32_program,
+                feed=data,
+                fetch_list=[fp32_fetch_list],
+                return_numpy=True,
+                scope=scope,
+            )
+
+            adaround_weight_tensor_dequant = (
+                adaround._calculate_adarounded_weights()
+            )
+            np_adaround_out_tensor = (
+                adaround._calculate_output_with_adarounded_weights(
+                    fp32_program,
+                    place,
+                    exe,
+                    data,
+                    fp32_fetch_list,
+                    adaround_weight_tensor_dequant,
+                )
+            )
+
+            # If the cosine distance of the two tensor is small, skip training
+            cos_error = calculate_quant_cos_error(
+                np_orig_out_tensor[0], np_adaround_out_tensor[0]
+            )
+            if fast_mode and cos_error > 0.99:
+                _logger.info("The cosine error is small, skip training.")
+                break
+            beta, warm_start = adaround.update_beta_warm(i)
+            feed_dict = {
+                'orig_out_tensor': np_orig_out_tensor[0],
+                'adaround_out_tensor': np_adaround_out_tensor[0],
+                'beta': beta,
+                'warm_start': warm_start,
+            }
+            out = exe.run(
+                train_program,
+                feed=feed_dict,
+                fetch_list=[v.name for v in train_fetches_loss.values()],
+                return_numpy=True,
+            )
+            _logger.info(
+                f"Iter {i:d}, lr {lr:.5f}, loss {np.mean(out[0]):.5f}, loss_round {np.mean(out[1]):.5f}, loss_recon {np.mean(out[2]):.5f}, time {start_time - prev_start_time:.5f}s"
+            )
+            sys.stdout.flush()
+            if i == num_iterations:
+                break
+        final_weight_tensor_quant_dict[weight_var_name] = (
+            adaround.update_final_weights()
+        )
+
+        if bias_correction:
+            final_weight_tensor_quant_dict[weight_var_name] = bias_correction_w(
+                weight_var_tensor,
+                final_weight_tensor_quant_dict[weight_var_name],
+                scale,
+                adaround.quant_axis,
+                weight_bits=adaround.weight_bits,
+            )
+
+        del adaround
+
+    # update adarounded calibrated weights
+    for weight_var_name in quantized_op_pairs.keys():
+        set_variable_data(
+            scope,
+            place,
+            weight_var_name,
+            final_weight_tensor_quant_dict[weight_var_name],
+        )

python/paddle/static/quantization/cal_kl_threshold.py

@@ -0,0 +1,144 @@
+#   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 logging
+import math
+
+import numpy as np
+
+from ..log_helper import get_logger
+
+_logger = get_logger(
+    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
+)
+
+
+def expand_quantized_bins(quantized_bins, reference_bins):
+    '''
+    Expand hist bins.
+    '''
+    expanded_quantized_bins = [0] * len(reference_bins)
+    num_merged_bins = int(len(reference_bins) / len(quantized_bins))
+    j_start = 0
+    j_end = num_merged_bins
+    for idx in range(len(quantized_bins)):
+        zero_count = reference_bins[j_start:j_end].count(0)
+        num_merged_bins = j_end - j_start
+        if zero_count == num_merged_bins:
+            avg_bin_ele = 0
+        else:
+            avg_bin_ele = quantized_bins[idx] / (
+                num_merged_bins - zero_count + 0.0
+            )
+        for idx1 in range(j_start, j_end):
+            expanded_quantized_bins[idx1] = (
+                0 if reference_bins[idx1] == 0 else avg_bin_ele
+            )
+        j_start += num_merged_bins
+        j_end += num_merged_bins
+        if (idx + 1) == len(quantized_bins) - 1:
+            j_end = len(reference_bins)
+    return expanded_quantized_bins
+
+
+def safe_entropy(reference_distr_P, P_sum, candidate_distr_Q, Q_sum):
+    '''
+    Calculate the entropy.
+    '''
+    assert len(reference_distr_P) == len(candidate_distr_Q)
+    tmp_sum1 = 0
+    tmp_sum2 = 0
+    for idx in range(len(reference_distr_P)):
+        p_idx = reference_distr_P[idx]
+        q_idx = candidate_distr_Q[idx]
+        if p_idx == 0:
+            tmp_sum1 += 0
+            tmp_sum2 += 0
+        else:
+            if q_idx == 0:
+                _logger.error(
+                    "Fatal error!, idx = "
+                    + str(idx)
+                    + " qindex = 0! p_idx = "
+                    + str(p_idx)
+                )
+            tmp_sum1 += p_idx * (math.log(Q_sum * p_idx))
+            tmp_sum2 += p_idx * (math.log(P_sum * q_idx))
+    return (tmp_sum1 - tmp_sum2) / P_sum
+
+
+def cal_kl_threshold(hist, bin_width, bits):
+    '''
+    Using the KL-divergence method to get the more precise threshold.
+
+    Args:
+        hist(List): The hist of the tensor.
+        bin_width(float): The bin width for the hist.
+        bits(int): The quantization bits.
+    '''
+    assert hist.ndim == 1
+    hist_bins = hist.shape[0]
+    starting_iter = int((hist_bins - 1) * 0.5)
+    quant_range = 2 ** (bits - 1) - 1
+
+    P_sum = np.sum(np.array(hist).ravel())
+    min_kl_divergence = 0
+    min_kl_index = 0
+    kl_inited = False
+
+    for i in range(starting_iter, hist_bins):
+        reference_distr_P = hist[0:i].tolist()
+        outliers_count = sum(hist[i:])
+        if reference_distr_P[i - 1] == 0:
+            continue
+        reference_distr_P[i - 1] += outliers_count
+        reference_distr_bins = reference_distr_P[:]
+        candidate_distr_Q = hist[0:i].tolist()
+        num_merged_bins = int(i / quant_range)
+        candidate_distr_Q_quantized = [0] * quant_range
+        j_start = 0
+        j_end = num_merged_bins
+        for idx in range(quant_range):
+            candidate_distr_Q_quantized[idx] = sum(
+                candidate_distr_Q[j_start:j_end]
+            )
+            j_start += num_merged_bins
+            j_end += num_merged_bins
+            if (idx + 1) == quant_range - 1:
+                j_end = i
+        candidate_distr_Q = expand_quantized_bins(
+            candidate_distr_Q_quantized, reference_distr_bins
+        )
+        Q_sum = sum(candidate_distr_Q)
+        kl_divergence = safe_entropy(
+            reference_distr_P, P_sum, candidate_distr_Q, Q_sum
+        )
+        if not kl_inited:
+            min_kl_divergence = kl_divergence
+            min_kl_index = i
+            kl_inited = True
+        elif kl_divergence < min_kl_divergence:
+            min_kl_divergence = kl_divergence
+            min_kl_index = i
+        else:
+            pass
+    if min_kl_index == 0:
+        while starting_iter > 0:
+            if hist[starting_iter] == 0:
+                starting_iter -= 1
+                continue
+            else:
+                break
+        min_kl_index = starting_iter
+    return (min_kl_index + 0.5) * bin_width

python/paddle/static/quantization/quant2_int8_onednn_pass.py

@@ -0,0 +1,736 @@
+#   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
+
+from paddle.utils import deprecated
+
+from ...base.framework import IrGraph
+from ...framework import _get_paddle_place, core
+
+OpRole = core.op_proto_and_checker_maker.OpRole
+
+
+class Quant2Int8OnednnPass:
+    """
+    Transform a quant model IrGraph into MKL-DNN supported INT8 IrGraph.
+    The pass consists of the following transformations:
+        1. gather scale values from fake quantize/dequantize operators,
+        2. extract FP32 inference model graph from the quant graph, i.e.
+            a.  remove fake quantize/dequantize operators,
+            b.  dequantize conv2d and mul's weights,
+        3. optimize the FP32 graph using standard FP32 optimization fuses
+            (e.g. `conv2d`+`bn` -> `conv2d`),
+        4. quantize the optimized FP32 graph using standard INT8v2 quantization
+            passes (`cpu_quantize_pass`, `cpu_quantize_squash_pass`).
+    """
+
+    def __init__(
+        self,
+        _ops_to_quantize,
+        _op_ids_to_skip=None,
+        _scope=None,
+        _place=None,
+        _core=None,
+        _debug=False,
+    ):
+        self._scope = _scope
+        self._place = _get_paddle_place(_place)
+        self._core = _core
+        self._debug = _debug
+        self._fake_quantize_types = [
+            'fake_quantize_moving_average_abs_max',
+            'fake_quantize_range_abs_max',
+        ]
+        self._fake_dequantize_types = [
+            'fake_dequantize_max_abs',
+            'fake_channel_wise_dequantize_max_abs',
+        ]
+        self._fake_quantize_dequantize_types = [
+            'fake_quantize_dequantize_abs_max',
+            'fake_quantize_dequantize_moving_average_abs_max',
+            'fake_channel_wise_quantize_dequantize_abs_max',
+        ]
+        self._ops_to_quantize = _ops_to_quantize
+        self._op_ids_to_skip = (
+            _op_ids_to_skip if _op_ids_to_skip is not None else {-1}
+        )
+        self._scale_immutable_ops = [
+            'transpose2',
+            'reshape2',
+            'pool2d',
+            'slice',
+            'shape',
+            'nearest_interp',
+            'nearest_interp_v2',
+            'split',
+        ]
+        self._scale_ops = ['scale']
+        self._conv_ops = ['conv2d', 'depthwise_conv2d']
+        self._pool_ops = ['pool2d']
+        self._mul_ops = ['mul']
+        self._fc_ops = ['fc']
+        self._relu_ops = ['relu', 'relu6']
+        self._matmul_ops = ['matmul', 'matmul_v2']
+        self._gru_ops = ['fusion_gru', 'multi_gru']
+        self._lstm_ops = ['fusion_lstm']
+        self._weight_thresholds = {}
+        # Collect the Input and Output scales from Fake quant models
+        self._var_quant_scales = {}
+        self._max_range = {}
+        self._s8_max = 127
+        self._pass_idx = 0
+        self._pass_group = 'int8'
+
+    def apply(self, graph):
+        assert isinstance(graph, IrGraph), (
+            'graph must be the instance of IrGraph.'
+        )
+
+        self._reset_pass_idx_and_group('int8')
+        graph = self._label_skip_quantized_op(graph)
+        graph = self._gather_weight_thresholds_from_fake(graph)
+        graph = self._gather_input_scales_from_fake(graph)
+        graph = self._gather_output_scales_from_attr(graph)
+        graph = self._remove_fake_ops(graph)
+        graph = self._dequantize_weights(graph)
+        graph = self._optimize_fp32_graph(graph)
+        graph = self._compute_weight_scales(graph)
+        # This function causes nondeterministic quantization behavior
+        # graph = self._update_relu_output_scales(graph)
+        graph = self._propagate_scales(graph)
+        graph = self._quantize_fp32_graph(graph)
+        graph = self._cleanup(graph)
+        return graph
+
+    def prepare_and_optimize_fp32(self, graph):
+        assert isinstance(graph, IrGraph), (
+            'graph must be the instance of IrGraph.'
+        )
+
+        self._reset_pass_idx_and_group('fp32')
+        graph = self._optimize_fp32_graph(graph)
+        graph = self._cleanup(graph)
+        return graph
+
+    def _reset_pass_idx_and_group(self, group):
+        self._pass_idx = 0
+        self._pass_group = group
+
+    def _convert_scale2tensor(self, scale):
+        tensor = core.DenseTensor()
+        tensor.set(scale, core.CPUPlace())
+        return tensor
+
+    def _is_quantizing_all_ops(self):
+        return len(self._ops_to_quantize) == 0
+
+    def _is_any_of_op_types_in_graph(self, op_types, graph):
+        return any(op.name() in op_types for op in graph.all_op_nodes())
+
+    def _is_any_of_op_types_quantized(self, op_types, graph):
+        return self._is_any_of_op_types_in_graph(op_types, graph) and (
+            self._is_quantizing_all_ops()
+            or any(op_type in self._ops_to_quantize for op_type in op_types)
+        )
+
+    def _is_conv_quantized(self, graph):
+        return self._is_any_of_op_types_quantized(self._conv_ops, graph)
+
+    def _is_fc_quantized(self, graph):
+        return self._is_any_of_op_types_quantized(self._fc_ops, graph)
+
+    def _label_skip_quantized_op(self, graph):
+        """
+        For some ops(conv2d, depthwise_conv2d, mul, matmul), find and label
+        the skip quantized ops. cpu_quantize_placement_pass will use the
+        label to identify it.
+        For static models, the skip quantized ops have `skip_quant` attr.
+        Therefore, it only needs to find and label the skip quantized ops for
+        dygraph models, in which the quantized ops don't have `quantization_type`
+        attr.
+        """
+        target_ops = self._conv_ops + self._mul_ops + self._matmul_ops
+        for op_node in graph.all_op_nodes():
+            if op_node.name() in target_ops and not op_node.op().has_attr(
+                "quantization_type"
+            ):
+                is_quantized_op = True
+                for var_node in op_node.inputs:
+                    for front_op_node in var_node.inputs:
+                        if "quantize" not in front_op_node.name():
+                            is_quantized_op = False
+                if not is_quantized_op:
+                    op_node.op()._set_attr("skip_quant", True)
+        return graph
+
+    def _add_scale_for_vars(self, var_names, use_unsigned_int, lod_tensor):
+        """
+        Save quantization scales for variables. Do not overwrite.
+        """
+        scales = self._var_quant_scales
+        for var_name in var_names:
+            if var_name not in scales:
+                scales[var_name] = (use_unsigned_int, lod_tensor)
+
+    def _gather_input_scales_from_fake(self, graph):
+        # fake_quantize_dequantize_abs_max doesn't have scale value
+        fake_ops = ['fake_quantize_dequantize_moving_average_abs_max']
+        fake_ops.extend(self._fake_quantize_types)
+
+        for op in graph.all_op_nodes():
+            if op.name() in fake_ops:
+                bit_length = op.op().attr("bit_length")
+                assert bit_length == 8, (
+                    f'Unsupported number quantization bits ({bit_length}). Only 8 is supported now.'
+                )
+
+                input_name = op.input("X")[0]
+                scale_name = op.input("InScale")[0]
+                output_name = op.output("Out")[0]
+                # Gather new weight scales after folding batchnorm in convolution
+                scale = np.array(
+                    1.0 / self._load_param(self._scope, scale_name)[0]
+                ).astype(np.float64)
+                scale[scale == np.inf] = 0.0
+                lod_tensor = self._convert_scale2tensor(scale)
+                use_unsigned_int = False
+                self._add_scale_for_vars(
+                    [input_name, output_name], use_unsigned_int, lod_tensor
+                )
+
+        return graph
+
+    def _gather_weight_thresholds_from_fake(self, graph):
+        for op in graph.all_op_nodes():
+            if op.name() in self._fake_dequantize_types:
+                input_name = op.input("X")[0]
+                if op.op().has_attr("max_range"):
+                    _max_range = np.array(op.op().attr("max_range")).astype(
+                        np.float64
+                    )
+                    self._weight_thresholds[input_name] = np.array(
+                        self._s8_max * self._s8_max / _max_range
+                    ).astype(np.float64)
+                else:
+                    scale_name = op.input("Scales")[0]
+                    self._weight_thresholds[input_name] = np.array(
+                        self._load_param(self._scope, scale_name)
+                    ).astype(np.float64)
+
+        return graph
+
+    def _gather_output_scales_from_attr(self, graph):
+        for op in graph.all_op_nodes():
+            if op.op().has_attr("out_threshold"):
+                attr_scale = op.op().attr("out_threshold")
+                if attr_scale == 0.0:
+                    continue
+                scale = np.array(1.0 / attr_scale).astype(np.float64)
+                scale[scale == np.inf] = 0.0
+                scale_lod_tensor = self._convert_scale2tensor(scale)
+                use_unsigned_int = False
+                for output_name in op.op().outputs():
+                    for out_var_name in op.op().output(output_name):
+                        self._add_scale_for_vars(
+                            [out_var_name], use_unsigned_int, scale_lod_tensor
+                        )
+
+        return graph
+
+    def _propagate_scales(self, graph):
+        def _update_scale_op_in_scale(op, input, output):
+            unsigned, tensor = self._var_quant_scales[output]
+            scale = np.array(tensor) * op.op().attr("scale")
+            new_tensor = self._convert_scale2tensor(scale.astype(np.float64))
+            self._var_quant_scales[input] = (unsigned, new_tensor)
+
+        def _update_scales(graph):
+            waiting_for_scale = set()
+            for op in graph.all_op_nodes():
+                if op.name() in self._scale_immutable_ops:
+                    if op.name() == 'slice' or op.name() == 'shape':
+                        input_name = op.input("Input")[0]
+                    else:
+                        input_name = op.input("X")[0]
+                    output_name = op.output("Out")[0]
+                    tensor_names = [input_name, output_name]
+
+                    if all(
+                        name not in self._var_quant_scales
+                        for name in tensor_names
+                    ):
+                        waiting_for_scale.update(tensor_names)
+                        continue
+                    elif input_name in self._var_quant_scales:
+                        self._var_quant_scales[output_name] = (
+                            self._var_quant_scales[input_name]
+                        )
+                    elif output_name in self._var_quant_scales:
+                        self._var_quant_scales[input_name] = (
+                            self._var_quant_scales[output_name]
+                        )
+
+                elif op.name() == 'concat':
+                    output_name = op.output("Out")[0]
+                    if output_name in self._var_quant_scales:
+                        input_names = op.input("X")
+                        for input_name in input_names:
+                            self._var_quant_scales[input_name] = (
+                                self._var_quant_scales[output_name]
+                            )
+                elif op.name() in self._scale_ops:
+                    input_name = op.input("X")[0]
+                    output_name = op.output("Out")[0]
+                    if output_name in self._var_quant_scales:
+                        _update_scale_op_in_scale(op, input_name, output_name)
+            return waiting_for_scale
+
+        waiting_for_scale = _update_scales(graph)
+        waiting_for_scale_prev = set()
+
+        while (
+            len(waiting_for_scale) != 0
+            and waiting_for_scale != waiting_for_scale_prev
+        ):
+            waiting_for_scale_prev = waiting_for_scale
+            waiting_for_scale = _update_scales(graph)
+
+        return graph
+
+    def _load_param(self, scope, param_name):
+        return np.array(scope.find_var(param_name).get_tensor())
+
+    def _remove_fake_ops(self, graph):
+        for op in graph.all_op_nodes():
+            if op.name() in self._fake_quantize_types:
+                self._remove_fake_quantize(graph, op)
+            elif op.name() in self._fake_dequantize_types:
+                self._remove_fake_dequantize(graph, op)
+            elif op.name() in self._fake_quantize_dequantize_types:
+                self._remove_fake_dequantize(graph, op)
+
+        return graph
+
+    def _remove_fake_quantize(self, graph, op):
+        fake_quant_in = graph._find_node_by_name(op.inputs, op.input("X")[0])
+        fake_quant_in_scale = graph._find_node_by_name(
+            op.inputs, op.input("InScale")[0]
+        )
+        fake_quant_out = graph._find_node_by_name(
+            op.outputs, op.output("Out")[0]
+        )
+        fake_quant_out_scale = graph._find_node_by_name(
+            op.outputs, op.output("OutScale")[0]
+        )
+
+        next_ops = fake_quant_out.outputs
+        for next_op in next_ops:
+            self._swap_inputs(next_op, fake_quant_out, fake_quant_in)
+            graph.link_to(fake_quant_in, next_op)
+        graph.safe_remove_nodes(
+            {op, fake_quant_in_scale, fake_quant_out, fake_quant_out_scale}
+        )
+
+        return graph
+
+    def _remove_fake_dequantize(self, graph, op):
+        fake_dequant_in = graph._find_node_by_name(op.inputs, op.input("X")[0])
+        fake_dequant_out = graph._find_node_by_name(
+            op.outputs, op.output("Out")[0]
+        )
+
+        next_ops = fake_dequant_out.outputs
+        for next_op in next_ops:
+            self._swap_inputs(next_op, fake_dequant_out, fake_dequant_in)
+            graph.link_to(fake_dequant_in, next_op)
+        graph.safe_remove_nodes({op, fake_dequant_out})
+
+        return graph
+
+    def _swap_inputs(self, op, old_input, new_input):
+        for input_name in op.op().input_names():
+            if old_input.name() in op.input(input_name):
+                op.op().set_input(
+                    input_name,
+                    [
+                        new_input.name() if x == old_input.name() else x
+                        for x in op.input(input_name)
+                    ],
+                )
+
+    def _dequantize_weights(self, graph):
+        def _is_int8_weights(op_node, weight_name):
+            weight_var_name = op_node.input(weight_name)[0]
+            if self._scope.find_var(weight_var_name) is None:
+                return False
+            weight = self._load_param(self._scope, weight_var_name)
+            return np.all(np.mod(weight, 1) == 0)
+
+        mul_and_matmul_ops = self._mul_ops + self._matmul_ops
+        for op in graph.all_op_nodes():
+            if op.name() in self._conv_ops and _is_int8_weights(op, "Filter"):
+                self._dequantize_op_weights(graph, op, "Filter", "Output")
+            elif op.name() in mul_and_matmul_ops and _is_int8_weights(op, "Y"):
+                self._dequantize_op_weights(graph, op, "Y", "Out")
+
+        return graph
+
+    def _dequantize_op_weights(self, graph, op_node, weight_name, output_name):
+        weight_var_name = op_node.input(weight_name)[0]
+        output_var_name = op_node.output(output_name)[0]
+        # Convert int8 range weights to fp32 range weights
+        scales = self._weight_thresholds[output_var_name]
+        weight = self._load_param(self._scope, weight_var_name)
+        if scales.size == 1 or scales.size == weight.shape[0]:
+            w_fp32 = np.multiply(np.divide(weight, self._s8_max).T, scales.T).T
+        elif len(weight.shape) > 1 and scales.size == weight.shape[1]:
+            w_fp32 = np.multiply(np.divide(weight, self._s8_max), scales)
+        else:
+            raise ValueError(
+                f"The size of weight scales vector ({scales.size}) does not match the dimensions ({weight.shape}) of the weights tensor {weight_var_name}."
+            )
+        w_fp32 = w_fp32.reshape(weight.shape).astype(np.float32)
+        self._restore_var(weight_var_name, w_fp32)
+
+    def _restore_var(self, name, array):
+        tensor = self._scope.find_var(name).get_tensor()
+        tensor.set(array, self._place)
+
+    def _update_activations(self, graph):
+        for op in graph.all_op_nodes():
+            if op.name() in self._conv_ops and not op.op().has_attr(
+                "fuse_activation"
+            ):
+                activation = ""
+                if op.op().has_attr("fuse_relu") and op.op().attr("fuse_relu"):
+                    activation = "relu"
+                op.set_attr("fuse_activation", activation)
+        return graph
+
+    def _remove_ctrl_vars(self, graph):
+        remove_ctr_vars = set()
+        for node in graph.all_var_nodes():
+            if node.is_ctrl_var():
+                remove_ctr_vars.add(node)
+        graph.safe_remove_nodes(remove_ctr_vars)
+        return graph
+
+    def _optimize_fp32_graph(self, graph):
+        graph = self._update_activations(graph)
+        graph = self._remove_ctrl_vars(graph)
+        graph = self._apply_pass(
+            graph, 'onednn_placement_pass', ['onednn_enabled_op_types'], [set()]
+        )
+        # remove dropout ops
+        graph = self._apply_pass(graph, 'simplify_with_basic_ops_pass')
+        graph = self._apply_pass(graph, 'layer_norm_fuse_pass')
+        graph = self._apply_pass(graph, 'attention_lstm_fuse_pass')
+        graph = self._apply_pass(graph, 'seqconv_eltadd_relu_fuse_pass')
+        graph = self._apply_pass(graph, 'fc_lstm_fuse_pass')
+        graph = self._apply_pass(graph, 'mul_lstm_fuse_pass')
+        graph = self._apply_pass(graph, 'fc_gru_fuse_pass')
+        graph = self._apply_pass(graph, 'mul_gru_fuse_pass')
+        graph = self._apply_pass(graph, 'multi_gru_fuse_pass')
+        graph = self._apply_pass(graph, 'multi_gru_seq_fuse_pass')
+        graph = self._apply_pass(graph, 'seq_concat_fc_fuse_pass')
+        graph = self._apply_pass(graph, 'gpu_cpu_squeeze2_matmul_fuse_pass')
+        graph = self._apply_pass(graph, 'gpu_cpu_reshape2_matmul_fuse_pass')
+        graph = self._apply_pass(graph, 'gpu_cpu_flatten2_matmul_fuse_pass')
+        graph = self._apply_pass(graph, 'matmul_v2_scale_fuse_pass')
+        graph = self._apply_pass(graph, 'squared_mat_sub_fuse_pass')
+        graph = self._apply_pass(graph, 'is_test_pass')
+        graph = self._apply_pass(graph, 'gpu_cpu_map_matmul_v2_to_mul_pass')
+        graph = self._apply_pass(graph, 'gpu_cpu_map_matmul_v2_to_matmul_pass')
+        graph = self._apply_pass(graph, 'matmul_scale_fuse_pass')
+        graph = self._apply_pass(graph, 'gpu_cpu_map_matmul_to_mul_pass')
+        graph = self._apply_pass(graph, 'repeated_fc_relu_fuse_pass')
+        graph = self._apply_pass(graph, 'depthwise_conv_onednn_pass')
+        graph = self._apply_pass(graph, 'conv_bn_fuse_pass')
+        graph = self._apply_pass(graph, 'conv_eltwiseadd_bn_fuse_pass')
+        graph = self._apply_pass(graph, 'conv_affine_channel_onednn_fuse_pass')
+        graph = self._apply_pass(graph, 'conv_transpose_bn_fuse_pass')
+        graph = self._apply_pass(
+            graph, 'conv_transpose_eltwiseadd_bn_fuse_pass'
+        )
+        graph = self._apply_pass(graph, 'conv_bias_onednn_fuse_pass')
+        graph = self._apply_pass(graph, 'conv_transpose_bias_onednn_fuse_pass')
+        graph = self._apply_pass(graph, 'conv_elementwise_add_onednn_fuse_pass')
+        graph = self._apply_pass(graph, 'conv_activation_onednn_fuse_pass')
+        graph = self._apply_pass(
+            graph, 'fc_fuse_pass', ['use_gpu', 'use_fc_padding'], [False, False]
+        )
+        graph = self._apply_pass(graph, 'repeated_fc_relu_fuse_pass')
+        if self._is_fc_quantized(graph):
+            # Disabled due to topology-dependent speed-up
+            graph = self._apply_pass(graph, 'fc_onednn_pass')
+            graph = self._apply_pass(graph, 'fc_act_onednn_fuse_pass')
+        graph = self._apply_pass(
+            graph, 'matmul_transpose_reshape_onednn_fuse_pass'
+        )
+        graph = self._apply_pass(
+            graph, 'matmul_elementwise_add_onednn_fuse_pass'
+        )
+        graph = self._apply_pass(graph, 'matmul_activation_onednn_fuse_pass')
+        graph = self._apply_pass(graph, 'batch_norm_act_fuse_pass')
+        graph = self._apply_pass(graph, 'softplus_activation_onednn_fuse_pass')
+        graph = self._apply_pass(graph, 'scale_matmul_fuse_pass')
+        graph = self._apply_pass(
+            graph, 'reshape_transpose_matmul_onednn_fuse_pass'
+        )
+        # the following pass should be the last one since it will work on all fused ops.
+        graph = self._apply_pass(graph, 'runtime_context_cache_pass')
+        return graph
+
+    def _apply_pass(self, graph, pass_name, attrs=None, attr_values=None):
+        ir_pass = core.get_pass(pass_name)
+        cpp_graph = graph.graph
+        if not cpp_graph.has('__param_scope__'):
+            cpp_graph.set_not_owned('__param_scope__', self._scope)
+        if attrs:
+            assert attr_values and len(attrs) == len(attr_values), (
+                "Different number of pass attributes and their values."
+            )
+            for attr, value in zip(attrs, attr_values):
+                ir_pass.set(attr, value)
+        ir_pass.apply(cpp_graph)
+        if self._debug:
+            graph.draw(
+                '.',
+                f'{self._pass_group}_{self._pass_idx}_{pass_name}',
+                graph.all_op_nodes(),
+            )
+        self._remove_unused_var_nodes(graph)
+        self._pass_idx += 1
+        return graph
+
+    def _cleanup(self, graph):
+        graph = self._remove_unused_var_nodes(graph)
+        graph = self._set_op_role_forward(graph)
+        return graph
+
+    def _remove_unused_var_nodes(self, graph):
+        all_used_vars = set()
+        ops = graph.all_op_nodes()
+        for op_node in ops:
+            for input_node in op_node.inputs:
+                all_used_vars.add(input_node)
+            for output_node in op_node.outputs:
+                all_used_vars.add(output_node)
+
+        all_used_vars = {n.node for n in all_used_vars}
+        all_unused_vars = set(
+            filter(
+                lambda node: node.node not in all_used_vars,
+                graph.all_var_nodes(),
+            )
+        )
+        graph.safe_remove_nodes(all_unused_vars)
+        return graph
+
+    def _set_op_role_forward(self, graph):
+        ops = graph.all_op_nodes()
+        for op in ops:
+            op.set_attr("op_role", OpRole.Forward)
+        return graph
+
+    def _compute_weight_scales(self, graph):
+        def _compute_var_scales(ops, w_name, axis):
+            for op in graph.all_op_nodes():
+                if op.op().type() in ops:
+                    weight_var_name = op.input(w_name)[0]
+                    weights = np.array(
+                        self._load_param(self._scope, weight_var_name)
+                    )
+                    scales = 1.0 / np.amax(
+                        np.abs(weights.reshape(weights.shape[0], -1)).astype(
+                            np.float64
+                        ),
+                        axis=axis,
+                    )
+                    scales[scales == np.inf] = 0.0
+
+                    lod_tensor = self._convert_scale2tensor(scales)
+                    use_unsigned_int = False
+                    self._var_quant_scales[weight_var_name] = (
+                        use_unsigned_int,
+                        lod_tensor,
+                    )
+
+        def _compute_single_gru_weight_scales(wx_var_name, wh_var_name):
+            wx = np.array(self._load_param(self._scope, wx_var_name))
+            wh = np.array(self._load_param(self._scope, wh_var_name))
+            OC = wh.shape[0]
+            scale_ur = 1.0 / np.max(
+                np.abs(
+                    np.concatenate(
+                        [
+                            wx[:, : 2 * OC],
+                            wh.flatten()[: 2 * OC * OC].reshape(OC, 2 * OC),
+                        ],
+                        axis=0,
+                    )
+                ),
+                axis=0,
+            )
+            scale_o = 1.0 / np.max(
+                np.abs(
+                    np.concatenate(
+                        [
+                            wx[:, 2 * OC :],
+                            wh.flatten()[2 * OC * OC :].reshape(OC, OC),
+                        ],
+                        axis=0,
+                    )
+                ),
+                axis=0,
+            )
+
+            gru_weights_scale = np.concatenate([scale_ur, scale_o]).astype(
+                'float'
+            )
+
+            return self._convert_scale2tensor(gru_weights_scale)
+
+        def _compute_gru_weight_scales(wx_name, wh_name):
+            for op in graph.all_op_nodes():
+                if op.op().type() in self._gru_ops:
+                    assert len(op.input(wx_name)) == len(op.input(wh_name)), (
+                        f'Mismatch in number of weights inputs ({len(op.input(wx_name))} for WeightX vs. {len(op.input(wh_name))} for WeightH).'
+                    )
+                    for i, wx_var_name in enumerate(op.input(wx_name)):
+                        wh_var_name = op.input(wh_name)[i]
+                        use_unsigned_int = False
+                        lod_tensor = _compute_single_gru_weight_scales(
+                            wx_var_name, wh_var_name
+                        )
+                        self._var_quant_scales[wx_var_name] = (
+                            use_unsigned_int,
+                            lod_tensor,
+                        )
+
+        def _compute_single_lstm_weight_scales(wx_var_name, wh_var_name):
+            wx = np.array(self._load_param(self._scope, wx_var_name))
+            wh = np.array(self._load_param(self._scope, wh_var_name))
+
+            lstm_weights_scale = 1.0 / np.max(
+                np.abs(np.concatenate([wx[:, :], wh[:, :]], axis=0)), axis=0
+            )
+            lstm_weights_scale = lstm_weights_scale.astype('float')
+
+            return self._convert_scale2tensor(lstm_weights_scale)
+
+        def _compute_lstm_weight_scales(wx_name, wh_name):
+            for op in graph.all_op_nodes():
+                if op.op().type() in self._lstm_ops:
+                    assert len(op.input(wx_name)) == len(op.input(wh_name)), (
+                        f'Mismatch in number of weights inputs ({len(op.input(wx_name))} for WeightX vs. {len(op.input(wh_name))} for WeightH).'
+                    )
+                    for i, wx_var_name in enumerate(op.input(wx_name)):
+                        wh_var_name = op.input(wh_name)[i]
+                        use_unsigned_int = False
+                        lod_tensor = _compute_single_lstm_weight_scales(
+                            wx_var_name, wh_var_name
+                        )
+                        self._var_quant_scales[wx_var_name] = (
+                            use_unsigned_int,
+                            lod_tensor,
+                        )
+
+        _compute_var_scales(self._conv_ops, "Filter", axis=1)
+        _compute_var_scales(self._fc_ops, "W", axis=0)
+        _compute_var_scales(self._gru_ops, "WeightH", axis=0)
+        _compute_var_scales(self._lstm_ops, "WeightH", axis=0)
+        _compute_gru_weight_scales("WeightX", "WeightH")
+        _compute_lstm_weight_scales("WeightX", "WeightH")
+        return graph
+
+    def _update_relu_output_scales(self, graph):
+        def _set_unsigned_scale(graph, ops, op_out_name, predicate):
+            '''
+            Sets the type of an output scale of a passed op type(s) to 'unsigned int8' if the
+            predicate applied on op passes. Typically, the predicate checks if op's
+            activation is set to relu.
+            '''
+            for op in graph.all_op_nodes():
+                if op.name() in ops:
+                    out_name = op.output(op_out_name)[0]
+                    if out_name in self._var_quant_scales and predicate(
+                        op.op()
+                    ):
+                        is_unsigned, tensor = self._var_quant_scales[out_name]
+                        if is_unsigned is False:
+                            # If the variable is signed, it means that the scales for this var
+                            # were computed for signed data, so the scale must be multiplied by 2
+                            # to fill the entire range of uint8
+                            scale = np.array(tensor) * 2
+                            tensor = self._convert_scale2tensor(
+                                scale.astype(np.float64)
+                            )
+                        self._var_quant_scales[out_name] = (True, tensor)
+            return graph
+
+        def conv_predicate(op):
+            return op.attr("fuse_activation") in self._relu_ops
+
+        graph = _set_unsigned_scale(
+            graph, self._conv_ops, "Output", conv_predicate
+        )
+
+        def fc_predicate(op):
+            return op.attr("activation_type") in self._relu_ops
+
+        graph = _set_unsigned_scale(graph, self._fc_ops, "Out", fc_predicate)
+
+        graph = _set_unsigned_scale(
+            graph, self._relu_ops, 'Out', lambda op: True
+        )
+
+        return graph
+
+    def _get_data_layout(self, graph):
+        return 'NHWC' if self._is_conv_quantized(graph) else 'NCHW'
+
+    def _quantize_fp32_graph(self, graph):
+        graph = self._apply_pass(graph, 'scale_matmul_fuse_pass')
+        graph = self._apply_pass(
+            graph, 'reshape_transpose_matmul_onednn_fuse_pass'
+        )
+        graph = self._apply_pass(
+            graph,
+            'cpu_quantize_placement_pass',
+            ['quantize_enabled_op_types'],
+            [self._ops_to_quantize],
+        )
+        graph = self._apply_pass(
+            graph,
+            'cpu_quantize_pass',
+            ['quant_var_scales', 'data_layout'],
+            [self._var_quant_scales, self._get_data_layout(graph)],
+        )
+        graph = self._apply_pass(graph, 'cpu_quantize_squash_pass')
+        graph = self._apply_pass(graph, 'int8_scale_calculation_onednn_pass')
+        graph = self._apply_pass(graph, 'params_quantization_onednn_pass')
+        return graph
+
+
+class Quant2Int8MkldnnPass(Quant2Int8OnednnPass):
+    @deprecated(
+        since="3.1.0",
+        update_to="paddle.static.quantization.Quant2Int8OnednnPass",
+        level=1,
+        reason="Quant2Int8MkldnnPass will be removed in future",
+    )
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)

python/paddle/static/quantization/quant_config.py

@@ -0,0 +1,287 @@
+#   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.
+
+
+# A dict of operators that contain weights and support quantization,
+# including operator names, actual input and output names.
+SUPPORT_WEIGHT_QUANTIZATION_OP_DICT = {
+    "conv2d": [["Input", "Filter"], ["Output"]],
+    "depthwise_conv2d": [["Input", "Filter"], ["Output"]],
+    "conv2d_transpose": [["Input", "Filter"], ["Output"]],
+    "mul": [["X", "Y"], ["Out"]],
+    "matmul": [["X", "Y"], ["Out"]],
+    "matmul_v2": [["X", "Y"], ["Out"]],
+}
+
+# A dict of operators that supports quantization and has only activation inputs,
+# including operator names, actual input and output names.
+SUPPORT_ACT_QUANTIZATION_OP_DICT = {
+    "mul": [["X", "Y"], ["Out"]],
+    "matmul": [["X", "Y"], ["Out"]],
+    "matmul_v2": [["X", "Y"], ["Out"]],
+    "pool2d": [["X"], ["Out"]],
+    "elementwise_add": [["X", "Y"], ["Out"]],
+    "concat": [["X"], ["Out"]],
+    "softmax": [["X"], ["Out"]],
+    "argmax": [["X"], ["Out"]],
+    "transpose": [["X"], ["Out"]],
+    "equal": [["X", "Y"], ["Out"]],
+    "gather": [["X"], ["Out"]],
+    "greater_equal": [["X", "Y"], ["Out"]],
+    "greater_than": [["X", "Y"], ["Out"]],
+    "less_equal": [["X", "Y"], ["Out"]],
+    "less_than": [["X", "Y"], ["Out"]],
+    "mean": [["X"], ["Out"]],
+    "not_equal": [["X", "Y"], ["Out"]],
+    "reshape": [["X"], ["Out"]],
+    "reshape2": [["X"], ["Out"]],
+    "transpose2": [["X"], ["Out"]],
+    "nearest_interp": [["X"], ["Out"]],
+    "trilinear_interp": [["X"], ["Out"]],
+    "slice": [["Input"], ["Out"]],
+    "squeeze": [["X"], ["Out"]],
+    "elementwise_sub": [["X", "Y"], ["Out"]],
+    "relu": [["X"], ["Out"]],
+    "relu6": [["X"], ["Out"]],
+    "leaky_relu": [["X"], ["Out"]],
+    "prelu": [["X", "Alpha"], ["Out"]],
+    "tanh": [["X"], ["Out"]],
+    "swish": [["X"], ["Out"]],
+    "dropout": [["X"], ["Out"]],
+    "batch_norm": [["X"], ["Y"]],
+    "layer_norm": [["X"], ["Y"]],
+    "sigmoid": [["X"], ["Out"]],
+    "elementwise_mul": [["X", "Y"], ["Out"]],
+    "elementwise_pow": [["X", "Y"], ["Out"]],
+    "hard_swish": [["X"], ["Out"]],
+    "hard_sigmoid": [["X"], ["Out"]],
+    "gru": [["Input", "Weight"], ["Hidden"]],
+    "lstm": [["Input", "Weight"], ["Hidden"]],
+    "pad2d": [["X"], ["Out"]],
+    "pad3d": [["X"], ["Out"]],
+    "flatten": [["X"], ["Out"]],
+    "flatten2": [["X"], ["Out"]],
+    "unsqueeze2": [["X"], ["Out"]],
+    "flatten_contiguous_range": [["X"], ["Out"]],
+    "split": [["X"], ["Out"]],
+    "squeeze2": [["X"], ["Out"]],
+    "nearest_interp_v2": [["X"], ["Out"]],
+    "bilinear_interp": [["X"], ["Out"]],
+    "bilinear_interp_v2": [["X"], ["Out"]],
+    "fill_constant_batch_size_like": [["Input"], ["Out"]],
+    "arg_max": [["X"], ["Out"]],
+    "abs": [["X"], ["Out"]],
+    "assign": [["X"], ["Out"]],
+    "cast": [["X"], ["Out"]],
+    "clip": [["X"], ["Out"]],
+    "box_coder": [["PriorBox"], ["OutputBox"]],
+    "crop": [["X"], ["Out"]],
+    "cumsum": [["X"], ["Out"]],
+    "expand_v2": [["X"], ["Out"]],
+    "fill_any_like": [["X"], ["Out"]],
+    "fill_constant": [[], ["Out"]],
+    "gelu": [["X"], ["Out"]],
+    "instance_norm": [["X"], ["Y"]],
+    "lookup_table": [["W", "Ids"], ["Out"]],
+    "lookup_table_v2": [["W", "Ids"], ["Out"]],
+    "norm": [["X"], ["Norm"]],
+    "p_norm": [["X"], ["Out"]],
+    "pow": [["X"], ["Out"]],
+    "reduce_mean": [["X"], ["Out"]],
+    "stack": [["X"], ["Y"]],
+    "top_k_v2": [["X"], ["Out", "Indices"]],
+    "logical_and": [["X", "Y"], ["Out"]],
+    "logical_not": [["X"], ["Out"]],
+    "meshgrid": [["X"], ["Out"]],
+    "roi_align": [["X", "ROIs"], ["Out"]],
+    "strided_slice": [["Input"], ["Out"]],
+    "where": [["Condition", "X", "Y"], ["Out"]],
+    "grid_sampler": [["X", "Grid"], ["Output"]],
+    "tile": [["X"], ["Out"]],
+    "group_norm": [["X"], ["Y", "Mean", "Variance"]],
+    "reduce_sum": [["X"], ["Out"]],
+    "square": [["X"], ["Out"]],
+    "softplus": [["X"], ["Out"]],
+    "shuffle_channel": [["X"], ["Out"]],
+    "reduce_max": [["X"], ["Out"]],
+    "scale": [["X"], ["Out"]],
+}
+
+# A full dict of operators that supports quantization,
+# including operator names, actual input and output names.
+SUPPORT_QUANTIZATION_OP_DICT = SUPPORT_WEIGHT_QUANTIZATION_OP_DICT.copy()
+SUPPORT_QUANTIZATION_OP_DICT.update(SUPPORT_ACT_QUANTIZATION_OP_DICT)
+
+
+class BaseQuantizer:
+    """
+    Basic quantization configuration class, which configures some hyperparameters
+    required for quantization, including the list of op types to be quantized,
+    quantization bit number for weight and activation and the range of quantization values.
+    Args:
+        quantizable_op_type(list[str], optional): List the type of ops
+            that will be quantized. Default is []. If quantizable_op_type is [],
+            it will use the default quantization op type of the qunat config in
+            the current Quantizer.
+        quant_bits(int, optional): Quantization bit number for weight and activation.
+            Default is 8.
+    """
+
+    def __init__(
+        self,
+        quantizable_op_type=[],
+        quant_bits=8,
+    ):
+        self._quantizable_op_type = quantizable_op_type
+        self._quant_bits = quant_bits
+        self._quant_min = -128
+        self._quant_max = 127
+
+    @property
+    def weight_quant_operation_types(self):
+        """
+        Operation type list which should support weight quantization.
+        And before these ops, quant dequant nodes will be inserted.
+        """
+        base_weight_op_type_list = list(
+            SUPPORT_WEIGHT_QUANTIZATION_OP_DICT.keys()
+        )
+        if self._quantizable_op_type:
+            weight_list = []
+            for _op_type in self._quantizable_op_type:
+                if _op_type in base_weight_op_type_list:
+                    weight_list.append(_op_type)
+            return weight_list
+        else:
+            return base_weight_op_type_list
+
+    @property
+    def activation_quant_operation_types(self):
+        """
+        Operation type list which should support activation quantization.
+        And before these ops, quant dequant nodes will be inserted.
+        """
+        base_act_op_type_list = list(SUPPORT_ACT_QUANTIZATION_OP_DICT.keys())
+        act_quant_op_list = []
+        if self._quantizable_op_type:
+            for _op_type in self._quantizable_op_type:
+                if _op_type in base_act_op_type_list:
+                    act_quant_op_list.append(_op_type)
+        else:
+            act_quant_op_list = [
+                'mul',
+                'matmul',
+                'matmul_v2',
+            ]
+        return act_quant_op_list
+
+    @property
+    def observer_operation_types(self):
+        """
+        Operation type list for observer in quantization. These nodes only count the
+        calibration boundary scale and do not participate in the fake quantization.
+        In order to facilitate the deployment of the prediction engine, quant
+        and dequant nodes will be inserted after these ops when exporting the model.
+        """
+        return list(SUPPORT_ACT_QUANTIZATION_OP_DICT.keys())
+
+
+class TensorRTQuantizer(BaseQuantizer):
+    """
+    TensorRT quantization configuration class.
+    Args:
+        quantizable_op_type(list[str], optional): List the type of ops
+            that will be quantized. Default is []. If quantizable_op_type is [],
+            it will use the default quantization op type of the qunat config in
+            the current Quantizer.
+        quant_bits(int, optional): Quantization bit number for weight and activation.
+            Default is 8.
+    """
+
+    def __init__(
+        self,
+        quantizable_op_type=[],
+        quant_bits=8,
+    ):
+        super().__init__()
+        self._quantizable_op_type = quantizable_op_type
+        self._quant_bits = quant_bits
+        self._quant_min = -128
+        self._quant_max = 127
+
+    @property
+    def activation_quant_operation_types(self):
+        """
+        Operation type list which should support activation quantization.
+        And before these ops, quant dequant nodes will be inserted.
+        """
+        return [
+            "pool2d",
+            "elementwise_add",
+            "elementwise_sub",
+            "elementwise_mul",
+            "elementwise_pow",
+            "concat",
+            "softmax",
+            "argmax",
+            "mean",
+            "relu",
+            "relu6",
+            "leaky_relu",
+            "tanh",
+            "swish",
+            "softplus",
+            "gelu",
+            "hard_sigmoid",
+            "hard_swish",
+            "sigmoid",
+            "layer_norm",
+            "matmul_v2",
+            "split",
+            "bilinear_interp",
+            "nearest_interp",
+            "trilinear_interp",
+            "nearest_interp_v2",
+            "bilinear_interp",
+            "bilinear_interp_v2",
+            "clip",
+            "pow",
+            "reduce_mean",
+            "reduce_sum",
+            "reduce_max",
+        ]
+
+
+class ARMCPUQuantizer(BaseQuantizer):
+    """
+    ARM CPU with Paddle Lite quantization configuration class.
+    Args:
+        quantizable_op_type(list[str], optional): List the type of ops
+            that will be quantized. Default is []. If quantizable_op_type is [],
+            it will use the default quantization op type of the qunat config in
+            the current Quantizer.
+        quant_bits(int, optional): Quantization bit number for weight and activation.
+            Default is 8.
+    """
+
+    def __init__(
+        self,
+        quantizable_op_type=[],
+        quant_bits=8,
+    ):
+        super().__init__()
+        self._quantizable_op_type = quantizable_op_type
+        self._quant_bits = quant_bits
+        self._quant_min = -127
+        self._quant_max = 127

python/paddle/static/quantization/quant_int8_onednn_pass.py

@@ -0,0 +1,302 @@
+#   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
+
+from paddle.utils import deprecated
+
+from ...base.framework import IrGraph
+from ...framework import _get_paddle_place
+
+
+class QuantInt8OnednnPass:
+    """
+    Convert QuantizationFreezePass generated IrGraph to MKL-DNN supported INT8
+    IrGraph. Following transformations did in this pass:
+        1. Convert int8 range weights with float32 data type, which are generated by
+           the QuantizationFreezePass, to float32 range weights with float32 data type
+           by using the corresponding scales. This conversion is because MKL-DNN INT8
+           conv2d kernel and mul kernel now only support float32 weights input, hence
+           weights quantization will happen inside the conv2d and mul INT8 kernel.
+        2. Create the new conv2d or mul op with the converted weights and link its output
+           to fake_dequantize_abs_max op's output and set conv2d's attribute "force_fp32
+           _output" as true
+        3. Transform fake_quantize_xx op to quantize op
+        4. Remove fake_dequantize_abs_max op
+    """
+
+    def __init__(self, _scope=None, _place=None):
+        r"""
+        Args:
+            scope(static.Scope): scope is used to initialize the new parameters.
+            place(static.CPUPlace|str): place is used to initialize the new parameters.
+            When it is string, it can be only 'cpu'.
+
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> # The original graph will be rewrite.
+                >>> import paddle
+                >>> from paddle import static
+                >>> from paddle.static.quantization import QuantInt8OnednnPass
+                >>> from paddle.framework import IrGraph
+                >>> from paddle.framework import core
+
+                >>> graph = IrGraph(core.Graph(static.Program().desc), for_test=False)
+                >>> place = paddle.CPUPlace()
+                >>> onednn_pass = QuantInt8OnednnPass(static.global_scope(), place)
+                >>> onednn_pass.apply(graph)
+        """
+
+        self._scope = _scope
+        self._place = _get_paddle_place(_place)
+
+        self._quantize_type = [
+            'fake_quantize_moving_average_abs_max',
+            'fake_quantize_range_abs_max',
+        ]
+        self._dequantize_type = ['fake_dequantize_max_abs']
+        self._quantize_dequantize_type = [
+            'fake_quantize_dequantize_moving_average_abs_max'
+        ]
+
+        self._quantizable_ops = ['conv2d', 'depthwise_conv2d', 'mul']
+        self._conv_ops = ['conv2d', 'depthwise_conv2d']
+        self._pool_ops = ['pool2d']
+
+        self._in_scale = {}
+        self._max_range = {}
+        self._new_output = {}
+        self._s8_max = 127
+
+    def apply(self, graph):
+        """
+        Quantize the graph for running MKL-DNN INT8 inference. According
+        to activation quantization type, the graph will transform fake
+        quantize ops to quantize ops and remove the fake dequantize ops.
+
+        Args:
+            graph(IrGraph): the applied graph.
+        """
+
+        assert isinstance(graph, IrGraph), (
+            'graph must be the instance of IrGraph.'
+        )
+        ops = graph.all_op_nodes()
+
+        persistable_vars = [p.name() for p in graph.all_persistable_nodes()]
+        # Collect the _in_scales and _max_range to calculate the new scales for MKL-DNN
+        # INT8 conv2d and mul
+        for op_node in ops:
+            if op_node.name() in self._dequantize_type:
+                input_name = op_node.input("X")[0]
+                scale_name = op_node.input("Scale")[0]
+                self._in_scale[input_name] = self._load_param(
+                    self._scope, scale_name
+                )[0]
+                self._max_range[input_name] = op_node.op().attr("max_range")
+                self._new_output[input_name] = op_node.output("Out")[0]
+
+            if op_node.name() in self._quantize_dequantize_type:
+                inputs = op_node.op().input_names()
+                attrs = op_node.op().attr_names()
+                input_name = op_node.input("X")[0]
+                scale_name = op_node.input("InScale")[0]
+                self._in_scale[input_name] = self._load_param(
+                    self._scope, scale_name
+                )[0]
+                #  self._max_range[input_name] = op_node.op().attr("max_range")
+                self._new_output[input_name] = op_node.output("Out")[0]
+
+        for op_node in ops:
+            if op_node.name() in self._quantizable_ops:
+                if op_node.name() in self._conv_ops:
+                    self._transform_to_conv_onednn(graph, op_node)
+                elif op_node.name() in self._pool_ops:
+                    self._transform_to_pool_onednn(graph, op_node)
+                else:
+                    self._transform_to_mul_onednn(graph, op_node)
+            elif op_node.name() in self._quantize_type:
+                self._transform_to_quantize_onednn(graph, op_node)
+            elif op_node.name() in self._dequantize_type:
+                self._remove_fake_dequantize_op(graph, op_node)
+            self._remove_unused_var_nodes(graph)
+        return graph
+
+    def _transform_to_pool_onednn(self, graph, op):
+        output_name = op.output("Out")[0]
+        input_name = op.input("X")[0]
+
+    def _transform_to_conv_onednn(self, graph, op_node):
+        weight_name = op_node.input("Filter")[0]
+        output_name = op_node.output("Output")[0]
+        # Convert int8 range weights to fp32 range weights
+        weight = self._load_param(self._scope, weight_name)
+        w_fp32 = np.divide(
+            np.multiply(weight, self._s8_max), self._max_range[output_name]
+        )
+        w_fp32 = w_fp32.reshape(weight.shape)
+        self._restore_var(weight_name, w_fp32)
+        input_var_node = graph._find_node_by_name(
+            op_node.inputs, op_node.input("Input")[0]
+        )
+        weight_var_node = graph._find_node_by_name(op_node.inputs, weight_name)
+
+        # Set fake_dequantize_abs_max's output as new output of conv2d
+        output_var_node = graph._find_node_by_name(
+            graph.all_var_nodes(), self._new_output[output_name]
+        )
+        attrs = {
+            name: op_node.op().attr(name) for name in op_node.op().attr_names()
+        }
+
+        conv_op_node = graph.create_op_node(
+            op_type='fused_conv2d',
+            attrs=attrs,
+            inputs={'Input': input_var_node, 'Filter': weight_var_node},
+            outputs={'Output': output_var_node},
+        )
+
+        # Based on the Quant's scales to calculate the scales of MKL-DNN INT8 conv2d
+        scale_in = self._s8_max / self._in_scale[output_name]
+        scale_w = []
+        scale_w = [self._max_range[output_name] / self._s8_max]
+
+        conv_op_node.set_attr("Scale_weights", scale_w)
+        conv_op_node.set_attr("Scale_in", scale_in)
+        conv_op_node.set_attr("Scale_out", 1.0)
+        conv_op_node.set_attr("use_onednn", 1)
+        conv_op_node.set_attr("force_fp32_output", 1)
+        graph.link_to(input_var_node, conv_op_node)
+        graph.link_to(weight_var_node, conv_op_node)
+        graph.link_to(conv_op_node, output_var_node)
+        graph.safe_remove_nodes(op_node)
+
+    def _transform_to_mul_onednn(self, graph, op_node):
+        # For MKL-DNN INT8 mul, input Y should be the weights
+        weight_name = op_node.input("Y")[0]
+        output_name = op_node.output("Out")[0]
+        # Convert int8 range weights to fp32 range weights
+        weight = self._load_param(self._scope, weight_name)
+        w_fp32 = np.divide(
+            np.multiply(weight, self._s8_max), self._max_range[output_name]
+        )
+        w_fp32 = w_fp32.reshape(weight.shape)
+        self._restore_var(weight_name, w_fp32)
+        input_var_node = graph._find_node_by_name(
+            op_node.inputs, op_node.input("X")[0]
+        )
+        weight_var_node = graph._find_node_by_name(op_node.inputs, weight_name)
+
+        # Set fake_dequantize_abs_max's output as new output of mul
+        output_var_node = graph._find_node_by_name(
+            graph.all_var_nodes(), self._new_output[output_name]
+        )
+        attrs = {
+            name: op_node.op().attr(name) for name in op_node.op().attr_names()
+        }
+
+        mul_op_node = graph.create_op_node(
+            op_type='mul',
+            attrs=attrs,
+            inputs={'X': input_var_node, 'Y': weight_var_node},
+            outputs={'Out': output_var_node},
+        )
+
+        # Based on the Quant's scales to calculate MKL-DNN INT8 mul's scales
+        scale_in = self._s8_max / self._in_scale[output_name]
+        scale_w = []
+        scale_w = [self._max_range[output_name] / self._s8_max]
+
+        mul_op_node.set_attr("scale_y", scale_w)
+        mul_op_node.set_attr("scale_x", scale_in)
+        mul_op_node.set_attr("scale_out", 1.0)
+        mul_op_node.set_attr("use_onednn", 1)
+        mul_op_node.set_attr("force_fp32_output", 1)
+        graph.link_to(input_var_node, mul_op_node)
+        graph.link_to(weight_var_node, mul_op_node)
+        graph.link_to(mul_op_node, output_var_node)
+        graph.safe_remove_nodes(op_node)
+
+    def _transform_to_quantize_onednn(self, graph, op_node):
+        """
+        Transform fake_quantize_xx op to quantize onednn op in the graph.
+        """
+        input_var_node = graph._find_node_by_name(
+            op_node.inputs, op_node.input("X")[0]
+        )
+        output_var_node = graph._find_node_by_name(
+            op_node.outputs, op_node.output("Out")[0]
+        )
+        scale_in = (
+            self._s8_max
+            / self._load_param(self._scope, op_node.input("InScale")[0])[0]
+        )
+        quant_op_node = graph.create_op_node(
+            op_type='quantize',
+            attrs={
+                'data_format': 'ONEDNNLAYOUT',
+                'use_onednn': 1,
+                'Scale': scale_in,
+                'is_negative_input': 1,
+            },
+            inputs={'Input': input_var_node},
+            outputs={'Output': output_var_node},
+        )
+        graph.link_to(input_var_node, quant_op_node)
+        graph.link_to(quant_op_node, output_var_node)
+        graph.safe_remove_nodes(op_node)
+
+    def _remove_fake_dequantize_op(self, graph, op_node):
+        input_var_node = graph._find_node_by_name(
+            op_node.inputs, op_node.input("X")[0]
+        )
+        graph.safe_remove_nodes(op_node)
+
+    def _load_param(self, scope, param_name):
+        return np.array(scope.find_var(param_name).get_tensor())
+
+    def _restore_var(self, name, array):
+        tensor = self._scope.find_var(name).get_tensor()
+        tensor.set(array, self._place)
+
+    def _remove_unused_var_nodes(self, graph):
+        all_used_vars = set()
+        ops = graph.all_op_nodes()
+        for op_node in ops:
+            for input_node in op_node.inputs:
+                all_used_vars.add(input_node)
+            for output_node in op_node.outputs:
+                all_used_vars.add(output_node)
+
+        all_used_vars = {n.node for n in all_used_vars}
+        all_unused_vars = set(
+            filter(
+                lambda node: node.node not in all_used_vars,
+                graph.all_var_nodes(),
+            )
+        )
+        graph.safe_remove_nodes(all_unused_vars)
+
+
+class QuantInt8MkldnnPass(QuantInt8OnednnPass):
+    @deprecated(
+        since="3.1.0",
+        update_to="paddle.static.quantization.QuantInt8OnednnPass",
+        level=1,
+        reason="QuantInt8MkldnnPass will be removed in future",
+    )
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)

python/paddle/static/quantization/quanter.py

@@ -0,0 +1,534 @@
+# Copyright (c) 2023 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 json
+import logging
+import os
+
+import paddle
+
+from ...base.framework import IrGraph, core
+from ..log_helper import get_logger
+from .quantization_pass import (
+    AddQuantDequantForResidual,
+    AddQuantDequantPass,
+    ConvertToInt8Pass,
+    OutScaleForInferencePass,
+    OutScaleForTrainingPass,
+    QuantizationFreezePass,
+    QuantizationTransformPass,
+)
+
+_logger = get_logger(__name__, level=logging.INFO)
+
+from . import quant_config
+from .post_training_quantization import PostTrainingQuantizationProgram
+from .quantization_pass import (
+    AddQuantDequantForInferencePass,
+    AddQuantDequantPassV2,
+    QuantizationTransformPassV2,
+    QuantWeightPass,
+)
+
+WEIGHT_QUANTIZATION_TYPES = [
+    'abs_max',
+    'channel_wise_abs_max',
+    'range_abs_max',
+    'moving_average_abs_max',
+]
+WEIGHT_QUANTIZATION_TYPES_TENSORRT = ['channel_wise_abs_max']
+
+ACTIVATION_QUANTIZATION_TYPES = [
+    'abs_max',
+    'range_abs_max',
+    'moving_average_abs_max',
+]
+
+ACTIVATION_QUANTIZATION_TYPES_TENSORRT = [
+    'range_abs_max',
+    'moving_average_abs_max',
+]
+
+VALID_DTYPES = ['int8']
+
+TRANSFORM_PASS_OP_TYPES = list(
+    quant_config.SUPPORT_WEIGHT_QUANTIZATION_OP_DICT.keys()
+)
+QUANT_DEQUANT_PASS_OP_TYPES = list(
+    quant_config.SUPPORT_ACT_QUANTIZATION_OP_DICT.keys()
+)
+
+TENSORRT_OP_TYPES = [
+    'mul',
+    'conv2d',
+    'pool2d',
+    'depthwise_conv2d',
+    'elementwise_add',
+    'leaky_relu',
+]
+
+VARS_MAPPING_TABLE = './mapping_table_for_saving_inference_model'
+
+_quant_config_default = {
+    # weight quantize type, default is 'channel_wise_abs_max'
+    'weight_quantize_type': 'channel_wise_abs_max',
+    # activation quantize type, default is 'moving_average_abs_max'
+    'activation_quantize_type': 'moving_average_abs_max',
+    # weight quantize bit num, default is 8
+    'weight_bits': 8,
+    # activation quantize bit num, default is 8
+    'activation_bits': 8,
+    # ops of name_scope in not_quant_pattern list, will not be quantized
+    'not_quant_pattern': ['skip_quant'],
+    # ops of type in quantize_op_types, will be quantized
+    'quantize_op_types': ['conv2d', 'depthwise_conv2d', 'mul'],
+    # data type after quantization, such as 'uint8', 'int8', etc. default is 'int8'
+    'dtype': 'int8',
+    # window size for 'range_abs_max' quantization. default is 10000
+    'window_size': 10000,
+    # The decay coefficient of moving average, default is 0.9
+    'moving_rate': 0.9,
+    # if True, 'quantize_op_types' will be TENSORRT_OP_TYPES
+    'for_tensorrt': False,
+    # if True, 'quantize_op_types' will be TRANSFORM_PASS_OP_TYPES + QUANT_DEQUANT_PASS_OP_TYPES
+    'is_full_quantize': False,
+    # if True, use onnx format to quant.
+    'onnx_format': True,
+    # quant post to get initial scale for quant_aware
+    'quant_post_first': False,
+    # whether scale can be train
+    'scale_trainable': True,
+}
+
+
+def load_dict():
+    with open(VARS_MAPPING_TABLE, 'r') as file:
+        data = file.read()
+        data = json.loads(data)
+        return data
+
+
+def save_dict(table):
+    with open(VARS_MAPPING_TABLE, 'w') as file:
+        file.write(json.dumps(table))
+
+
+def _parse_configs(user_config):
+    """
+    check if user's configs are valid.
+    Args:
+        user_config(dict): user's config.
+    Return:
+        configs(dict): final configs will be used.
+    """
+
+    configs = copy.deepcopy(_quant_config_default)
+    configs.update(user_config)
+
+    assert isinstance(configs['for_tensorrt'], bool) and isinstance(
+        configs['is_full_quantize'], bool
+    ), "'for_tensorrt' and 'is_full_quantize' must both be bool'"
+
+    # check if configs is valid
+    if configs['for_tensorrt']:
+        weight_types = WEIGHT_QUANTIZATION_TYPES_TENSORRT
+        activation_types = ACTIVATION_QUANTIZATION_TYPES_TENSORRT
+        platform = 'TensorRT'
+    else:
+        weight_types = WEIGHT_QUANTIZATION_TYPES
+        activation_types = WEIGHT_QUANTIZATION_TYPES
+        platform = 'PaddleLite'
+    assert configs['weight_quantize_type'] in weight_types, (
+        "Unknown weight_quantize_type: {}. {} only supports {} ".format(
+            configs['weight_quantize_type'], platform, weight_types
+        )
+    )
+
+    assert configs['activation_quantize_type'] in activation_types, (
+        "Unknown activation_quantize_type: {}. {} only supports {}".format(
+            configs['activation_quantize_type'], platform, activation_types
+        )
+    )
+
+    assert isinstance(configs['weight_bits'], int), (
+        "weight_bits must be int value."
+    )
+
+    assert configs['weight_bits'] >= 1 and configs['weight_bits'] <= 16, (
+        "weight_bits should be between 1 and 16."
+    )
+
+    assert isinstance(configs['activation_bits'], int), (
+        "activation_bits must be int value."
+    )
+
+    assert (
+        configs['activation_bits'] >= 1 and configs['activation_bits'] <= 16
+    ), "activation_bits should be between 1 and 16."
+
+    assert isinstance(configs['not_quant_pattern'], (list, str)), (
+        "not_quant_pattern must be list or str"
+    )
+
+    assert isinstance(configs['quantize_op_types'], list), (
+        "quantize_op_types must be a list"
+    )
+
+    if configs['for_tensorrt']:
+        configs['quantize_op_types'] = TENSORRT_OP_TYPES
+    elif configs['is_full_quantize']:
+        configs['quantize_op_types'] = (
+            TRANSFORM_PASS_OP_TYPES + QUANT_DEQUANT_PASS_OP_TYPES
+        )
+    else:
+        for op_type in configs['quantize_op_types']:
+            assert (op_type in QUANT_DEQUANT_PASS_OP_TYPES) or (
+                op_type in TRANSFORM_PASS_OP_TYPES
+            ), (
+                f"{op_type} is not support, \
+                        now support op types are {TRANSFORM_PASS_OP_TYPES + QUANT_DEQUANT_PASS_OP_TYPES}"
+            )
+
+    assert isinstance(configs['dtype'], str), "dtype must be a str."
+
+    assert configs['dtype'] in VALID_DTYPES, "dtype can only be " + " ".join(
+        VALID_DTYPES
+    )
+
+    assert isinstance(configs['window_size'], int), (
+        "window_size must be int value, window size for 'range_abs_max' quantization, default is 10000."
+    )
+
+    assert isinstance(configs['moving_rate'], float), (
+        "moving_rate must be float value, The decay coefficient of moving average, default is 0.9."
+    )
+
+    return configs
+
+
+def quant_aware(
+    program,
+    place,
+    config=None,
+    scope=None,
+    for_test=False,
+    weight_quantize_func=None,
+    act_quantize_func=None,
+    weight_preprocess_func=None,
+    act_preprocess_func=None,
+    optimizer_func=None,
+    executor=None,
+    return_program=False,
+    calib_config={},
+    draw_graph=False,
+    return_scale_dict=False,
+    scale_dict=None,
+    model_type=None,
+    pattern_ops=None,
+):
+    """Add quantization  and dequantization operators to "program"
+    for quantization training or testing.
+    Args:
+        program(paddle.static.Program): training or testing ``program``.
+        place(paddle.CPUPlace or paddle.CUDAPlace): This parameter represents
+            the executor run on which device.
+        config(dict, optional): configs for quantization. if None, will use default config.
+            Default: None.
+        scope(paddle.static.Scope): Scope records the mapping between variable names and variables,
+            similar to brackets in programming languages. Usually users can use
+            `paddle.static.global_scope <https://www.paddlepaddle.org.cn/documentation/docs/zh/develop/api_cn/executor_cn/global_scope_cn.html>`_.
+            When ``None`` will use `paddle.static.global_scope() <https://www.paddlepaddle.org.cn/documentation/docs/zh/develop/api_cn/executor_cn/global_scope_cn.html>`_ .
+            Default: ``None``.
+        for_test(bool): If the 'program' parameter is a test program, this parameter should be set to ``True``.
+            Otherwise, set to ``False``.Default: False
+        weight_quantize_func(function): Function that defines how to quantize weight. Using this
+                can quickly test if user's quantization method works or not. In this function, user should
+                both define quantization function and dequantization function, that is, the function's input
+                is non-quantized weight and function returns dequantized weight. If None, will use
+                quantization op defined by 'weight_quantize_type'.
+                Default is None.
+        act_quantize_func(function): Function that defines how to quantize activation. Using this
+                can quickly test if user's quantization method works or not. In this function, user should
+                both define quantization and dequantization process, that is, the function's input
+                is non-quantized activation and function returns dequantized activation. If None, will use
+                quantization op defined by 'activation_quantize_type'.
+                Default is None.
+        weight_preprocess_func(function): Function that defines how to preprocess weight before quantization. Using this
+                can quickly test if user's preprocess method works or not. The function's 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_func(function): Function that defines how to preprocess activation before quantization. Using this
+                can quickly test if user's preprocess method works or not. The function's input
+                is non-quantized activation and function returns processed activation to be quantized. If None, the activation will
+                be quantized directly.
+                Default is None.
+        optimizer_func(function): Function return a optimizer. When 'is_test' is False and user want to use self-defined
+            quantization function and preprocess function, this function must be set. Default is None.
+        exe(paddle.static.Executor): If user want to use self-defined quantization function and preprocess function, exe must be set for
+                initialization. Default is None.
+        return_program(bool): If user want return value is a Program rather than Compiled Program, This argument should be set True.
+                Default is False.
+        draw_graph(bool): whether to draw graph when quantization is initialized. In order to prevent cycle,
+                the ERNIE model needs to be set to True. Default is False.
+        return_scale_dict(bool): If user want to return scale dict, model_type and pattern_ops, this argument should be set True.
+                Default is False.
+        scale_dict(dict): Use scale dict to initialize scales in program. Default is None.
+        model_type(str): Model type can be 'transformer' or 'non-transformer'. If model type is transformer, patterns will be analyzed.
+                Default is None.
+        pattern_ops(dict): Pattern_ops contain pattern name and corresponding ops. Default is None.
+    Returns:
+        paddle.static.CompiledProgram | paddle.static.Program: Program with quantization and dequantization ``operators``
+    """
+
+    scope = paddle.static.global_scope() if not scope else scope
+    if config is None:
+        config = _quant_config_default
+    else:
+        assert isinstance(config, dict), "config must be dict"
+        config = _parse_configs(config)
+    _logger.info(f"quant_aware config {config}")
+
+    skip_tensor_list = []
+    same_scale_tensor_list = []
+
+    is_test = True if for_test else not config['scale_trainable']
+    if config['quant_post_first'] and for_test:
+        if 'quantizable_op_type' not in calib_config:
+            calib_config['quantizable_op_type'] = config['quantize_op_types']
+        exe = paddle.static.Executor() if executor is None else executor
+        post_training_quantization = PostTrainingQuantizationProgram(
+            exe,
+            program,
+            freeze_model=False,
+            skip_tensor_list=skip_tensor_list,
+            same_scale_tensor_list=same_scale_tensor_list,
+            batch_nums=10,
+            scale_dict=scale_dict,
+            return_graph=True,
+            **calib_config,
+        )
+        main_graph = post_training_quantization.quantize()
+        scale_dict = post_training_quantization._scale_dict
+        sub_graphs = list(main_graph.all_sub_graphs())
+    else:
+        main_graph = IrGraph(core.Graph(program.desc), for_test=for_test)
+        sub_graphs = list(main_graph.all_sub_graphs())
+        transform_pass_ops = []
+        quant_dequant_ops = []
+        if config.get('quant_config'):
+            transform_pass_ops = config[
+                'quant_config'
+            ].weight_quant_operation_types
+            quant_dequant_ops = config[
+                'quant_config'
+            ].activation_quant_operation_types
+        else:
+            for op_type in config['quantize_op_types']:
+                if op_type in TRANSFORM_PASS_OP_TYPES:
+                    transform_pass_ops.append(op_type)
+                elif op_type in QUANT_DEQUANT_PASS_OP_TYPES:
+                    quant_dequant_ops.append(op_type)
+        if len(transform_pass_ops) > 0:
+            transform_func = (
+                QuantizationTransformPassV2
+                if config['onnx_format']
+                else QuantizationTransformPass
+            )
+            transform_pass = transform_func(
+                scope=scope,
+                place=place,
+                weight_bits=config['weight_bits'],
+                activation_bits=config['activation_bits'],
+                activation_quantize_type=config['activation_quantize_type'],
+                weight_quantize_type=config['weight_quantize_type'],
+                window_size=config['window_size'],
+                moving_rate=config['moving_rate'],
+                quantizable_op_type=transform_pass_ops,
+                skip_pattern=config['not_quant_pattern'],
+                weight_quantize_func=weight_quantize_func,
+                act_quantize_func=act_quantize_func,
+                weight_preprocess_func=weight_preprocess_func,
+                act_preprocess_func=act_preprocess_func,
+                optimizer_func=optimizer_func,
+                executor=executor,
+                is_test=is_test,
+            )
+
+            for sub_graph in sub_graphs:
+                transform_pass.apply(sub_graph)
+
+            residual_pass = AddQuantDequantForResidual(
+                scope=scope,
+                place=place,
+                quant_bits=config['activation_bits'],
+                is_test=is_test,
+            )
+
+            for subgraph in sub_graphs:
+                residual_pass.apply(sub_graph)
+
+        if len(quant_dequant_ops) > 0:
+            qdq_func = (
+                AddQuantDequantPassV2
+                if config['onnx_format']
+                else AddQuantDequantPass
+            )
+            quant_dequant_pass = qdq_func(
+                scope=scope,
+                place=place,
+                moving_rate=config['moving_rate'],
+                quant_bits=config['activation_bits'],
+                skip_pattern=config['not_quant_pattern'],
+                quantizable_op_type=quant_dequant_ops,
+                is_test=is_test,
+            )
+
+            for sub_graph in sub_graphs:
+                quant_dequant_pass.apply(sub_graph)
+
+    out_scale_training_pass = OutScaleForTrainingPass(
+        scope=scope,
+        place=place,
+        moving_rate=config['moving_rate'],
+        is_test=is_test,
+        scale_dict=scale_dict,
+    )
+
+    for sub_graph in sub_graphs:
+        out_scale_training_pass.apply(sub_graph)
+
+    if (
+        (weight_preprocess_func is not None or act_preprocess_func is not None)
+        and not for_test
+        and not config['onnx_format']
+    ):
+        _logger.info(
+            "When a preprocess_func is used in quant_aware, Need to save a mapping table to match variable names in the convert phase."
+        )
+        _logger.info(f"The mapping table is saved as '{VARS_MAPPING_TABLE}'.")
+        for sub_graph in sub_graphs:
+            save_dict(sub_graph.out_node_mapping_table)
+
+    # TODO: remove it.
+    if draw_graph:
+        main_graph.draw('./', 'graph.pdf')
+
+    if for_test or return_program:
+        quant_program = main_graph.to_program()
+    else:
+        quant_program = paddle.static.CompiledProgram(main_graph.graph)
+
+    if return_scale_dict:
+        return quant_program, scale_dict, model_type, pattern_ops
+    else:
+        return quant_program
+
+
+def convert(program, place, config=None, scope=None, save_int8=False):
+    """
+    convert quantized and well-trained ``program`` to final  quantized
+    ``program``that can be used to  save ``inference model``.
+
+    Args:
+        program(paddle.static.Program): quantized and well-trained ``test program``.
+        place(paddle.CPUPlace or paddle.CUDAPlace): This parameter represents
+                the executor run on which device.
+        config(dict, optional): configs for convert. if set None, will use
+                default config. It must be same with config that used in
+                'quant_aware'. Default is None.
+        scope(paddle.static.Scope, optional):  Scope records the mapping between
+                variable names and variables, similar to brackets in
+                programming languages. Usually users can use
+                `paddle.static.global_scope <https://www.paddlepaddle.org.cn/documentation/docs/zh/develop/api_cn/executor_cn/global_scope_cn.html>`_.
+                When ``None`` will use
+                `paddle.static.global_scope() <https://www.paddlepaddle.org.cn/documentation/docs/zh/develop/api_cn/executor_cn/global_scope_cn.html>`_
+                . Default: ``None``.
+        save_int8: Whether to return ``program`` which model parameters'
+                dtype is ``int8``. This parameter can only be used to
+                get model size. Default: ``False``.
+    Returns:
+        Tuple : freezed program which can be used for inference.
+                when ``save_int8`` is False, return ``freezed_program(paddle.static.Program)``.
+                when ``save_int8`` is True, return ``freezed_program(paddle.static.Program)``
+                and ``freezed_program_int8(paddle.static.Program)``
+    """
+    scope = paddle.static.global_scope() if not scope else scope
+
+    if config is None:
+        config = _quant_config_default
+    else:
+        assert isinstance(config, dict), "config must be dict"
+        config = _parse_configs(config)
+    _logger.info(f"convert config {config}")
+    test_graph = IrGraph(core.Graph(program.desc), for_test=True)
+
+    if config['onnx_format']:
+        quant_weight_pass = QuantWeightPass(scope, place)
+        for sub_graph in test_graph.all_sub_graphs():
+            quant_weight_pass.apply(sub_graph)
+        out_scale_infer_pass = AddQuantDequantForInferencePass(
+            scope=scope, place=place, quant_bits=config['activation_bits']
+        )
+        for sub_graph in test_graph.all_sub_graphs():
+            out_scale_infer_pass.apply(sub_graph)
+    else:
+        out_scale_infer_pass = OutScaleForInferencePass(scope=scope)
+        for sub_graph in test_graph.all_sub_graphs():
+            out_scale_infer_pass.apply(sub_graph)
+        # Freeze the graph after training by adjusting the quantize
+        # operators' order for the inference.
+        freeze_pass = QuantizationFreezePass(
+            scope=scope,
+            place=place,
+            weight_bits=config['weight_bits'],
+            activation_bits=config['activation_bits'],
+            weight_quantize_type=config['weight_quantize_type'],
+        )
+        if os.path.exists(VARS_MAPPING_TABLE):
+            test_graph.out_node_mapping_table = load_dict()
+        for sub_graph in test_graph.all_sub_graphs():
+            freeze_pass.apply(sub_graph)
+
+    freezed_program = test_graph.to_program()
+
+    # Move sub blocks persistable var to global block
+    global_block = freezed_program.global_block()
+    for _op in global_block.ops:
+        if _op.type == "while":
+            _block_id = _op.attr("sub_block").id
+            _block = freezed_program.block(_block_id)
+            persistables = []
+            for _name, _var in _block.vars.items():
+                if _var.persistable:
+                    global_block._clone_variable(_var)
+                    persistables.append(_name)
+            for _name in persistables:
+                _block._remove_var(_name)
+            persistables.extend(_op.input('X'))
+            _op.desc.set_input("X", persistables)
+
+    assert not (save_int8 and config['onnx_format']), (
+        "When onnx_format=True, already saved int8 weight,so you can't set save_int8=True."
+    )
+    if save_int8:
+        convert_int8_pass = ConvertToInt8Pass(scope=scope, place=place)
+        for sub_graph in test_graph.all_sub_graphs():
+            convert_int8_pass.apply(sub_graph)
+        freezed_program_int8 = test_graph.to_program()
+        return freezed_program, freezed_program_int8
+    else:
+        return freezed_program

python/paddle/static/quantization/utils.py

@@ -0,0 +1,293 @@
+#   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 sys
+
+import numpy as np
+
+from ...base.framework import IrNode, Operator
+from .quant_config import SUPPORT_QUANTIZATION_OP_DICT
+
+_channelwise_quant_axis1_ops = [
+    'conv2d_transpose',
+    'mul',
+    'matmul',
+    'matmul_v2',
+]
+
+
+def _get_op_input_var_names(op):
+    """
+    Get the input var names of the op.
+    Args:
+        op(IrNode, Operator): the input op.
+    Returns:
+        input_var_names or None.
+    """
+    assert isinstance(op, (IrNode, Operator)), (
+        "The input op should be IrNode or Operator."
+    )
+    var_names = []
+    op_name = op.name() if isinstance(op, IrNode) else op.type
+    if op_name not in SUPPORT_QUANTIZATION_OP_DICT:
+        return []
+
+    name_list = SUPPORT_QUANTIZATION_OP_DICT[op_name][0]
+    for name in name_list:
+        var_name = op.input(name)
+        if isinstance(var_name, list):
+            var_names.extend(var_name)
+        else:
+            var_names.append(var_name)
+    return var_names
+
+
+def _get_op_output_var_names(op):
+    """ """
+    assert isinstance(op, (IrNode, Operator)), (
+        "The input op should be IrNode or Operator."
+    )
+    var_names = []
+    op_name = op.name() if isinstance(op, IrNode) else op.type
+    if op_name not in SUPPORT_QUANTIZATION_OP_DICT:
+        return []
+
+    name_list = SUPPORT_QUANTIZATION_OP_DICT[op_name][1]
+    for name in name_list:
+        var_name = op.output(name)
+        if isinstance(var_name, list):
+            var_names.extend(var_name)
+        else:
+            var_names.append(var_name)
+    return var_names
+
+
+def _get_input_name_index(op, input_var_name):
+    """Get the input name and index of the var_name in the op"""
+    assert isinstance(op, (IrNode, Operator)), (
+        "The input op should be IrNode or Operator."
+    )
+    op_name = op.name() if isinstance(op, IrNode) else op.type
+    if op_name not in SUPPORT_QUANTIZATION_OP_DICT:
+        return None
+
+    res = None
+    for argname in SUPPORT_QUANTIZATION_OP_DICT[op_name][0]:
+        var_names = op.input(argname)
+        for index, name in enumerate(var_names):
+            if name == input_var_name:
+                res = (argname, index)
+    return res
+
+
+def _get_output_name_index(op, output_var_name):
+    """Get the output name and index of the var_name in the op"""
+    assert isinstance(op, (IrNode, Operator)), (
+        "The input op should be IrNode or Operator."
+    )
+    op_name = op.name() if isinstance(op, IrNode) else op.type
+    if op_name not in SUPPORT_QUANTIZATION_OP_DICT:
+        return None
+
+    name_list = SUPPORT_QUANTIZATION_OP_DICT[op_name][1]
+    res = None
+    for name in name_list:
+        var_name = op.output(name)
+        for index, val in enumerate(var_name):
+            if val == output_var_name:
+                res = (name, index)
+    return res
+
+
+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, "Cannot find " + var_name + " in scope."
+    tensor = np.array(var_node.get_tensor())
+    if tensor.shape == ():
+        return tensor.reshape(1)
+    else:
+        return tensor
+
+
+def set_variable_data(scope, place, var_name, np_value):
+    '''
+    Set the value of var node by name, if the node exits,
+    '''
+    assert isinstance(np_value, np.ndarray), (
+        'The type of value should be numpy array.'
+    )
+    var_node = scope.find_var(var_name)
+    if var_node is not None:
+        tensor = var_node.get_tensor()
+        tensor.set(np_value, place)
+
+
+def quant_tensor(x, scale, quant_axis=0, weight_bits=8, onnx_format=False):
+    # symmetry quant
+    def _clip(x, scale):
+        x[x > scale] = scale
+        x[x < -scale] = -scale
+        return x
+
+    bnt = (1 << (weight_bits - 1)) - 1
+    if isinstance(scale, list) and len(scale) == 1:
+        scale = scale[0]
+    if isinstance(scale, list):
+        assert quant_axis in [-1, 0, 1], 'quant_axis should be 0 or 1 for now.'
+        for i, s in enumerate(scale):
+            if s == 0.0:
+                s = 1e-8
+            if quant_axis == 0:
+                if onnx_format:
+                    x[i] = np.round(x[i] / s * bnt)
+                    x[i] = np.clip(x[i], -bnt - 1, bnt)
+                else:
+                    x[i] = _clip(x[i], s)
+                    x[i] = x[i] / s * bnt
+            else:
+                if onnx_format:
+                    x[:, i] = np.round(x[:, i] / s * bnt)
+                    x[:, i] = np.clip(x[:, i], -bnt - 1, bnt)
+                else:
+                    x[:, i] = _clip(x[:, i], s)
+                    x[:, i] = x[:, i] / s * bnt
+    else:
+        scale = 1e-8 if scale == 0.0 else scale
+        if onnx_format:
+            x = np.round(x / scale * bnt)
+            x = np.clip(x, -bnt - 1, bnt)
+        else:
+            x = _clip(x, scale)
+            x = x / scale * bnt
+    return x
+
+
+def dequant_tensor(x, scale, quant_axis=0, weight_bits=8):
+    assert quant_axis in [0, 1], 'quant_axis should be 0 or 1 for now.'
+    bnt = (1 << (weight_bits - 1)) - 1
+    if isinstance(scale, list):
+        for i, s in enumerate(scale):
+            if s == 0.0:
+                s = 1e-8
+            if quant_axis == 0:
+                x[i] = x[i] * s / bnt
+            else:
+                x[:, i] = x[:, i] * s / bnt
+    else:
+        scale = 1e-8 if scale == 0.0 else scale
+        x = x * scale / bnt
+    return x
+
+
+def bias_correction_w(x, x_quant, scale_v, quant_axis, weight_bits=8):
+    '''
+    Bias correction for weight
+    '''
+    eps = 1e-8
+    bnt = (1 << (weight_bits - 1)) - 1
+    x_dequant = x_quant.copy()
+    if isinstance(scale_v, list):
+        if quant_axis == 0:
+            for i, s in enumerate(scale_v):
+                x_dequant[i] = x_dequant[i] * s / bnt
+            quant_bias = x - x_dequant
+            mean_bias = quant_bias.reshape(quant_bias.shape[0], -1).mean(-1)
+            std_orig = x.reshape(x.shape[0], -1).std(-1)
+            std_quant = x_dequant.reshape(x_dequant.shape[0], -1).std(-1)
+            std_bias = std_orig / (std_quant + eps)
+        else:
+            for i, s in enumerate(scale_v):
+                x_dequant[:, i] = x_quant[:, i] * s / bnt
+            quant_bias = x - x_dequant
+            mean_bias = np.array(
+                [quant_bias[:, i].mean() for i in range(quant_bias.shape[1])]
+            )
+            std_orig = np.array([x[:, i].std() for i in range(x.shape[1])])
+            std_quant = np.array(
+                [x_dequant[:, i].std() for i in range(x_dequant.shape[1])]
+            )
+            std_bias = std_orig / (std_quant + eps)
+    else:
+        x_dequant = x_quant * scale_v / bnt
+        mean_bias = (x - x_dequant).mean()
+        std_bias = x.std() / (x_dequant.std() + eps)
+    if mean_bias.ndim == 1:
+        std_bias = np.resize(std_bias, x.shape)
+        mean_bias = np.resize(mean_bias, x.shape)
+
+    x_dequant = (mean_bias + x_dequant) * std_bias
+    quantized_param_v = quant_tensor(
+        x_dequant, scale_v, quant_axis, weight_bits
+    )
+    return quantized_param_v
+
+
+def stable_sigmoid(x):
+    sig = np.where(x < 0, np.exp(x) / (1 + np.exp(x)), 1 / (1 + np.exp(-x)))
+    return sig
+
+
+def calculate_quant_cos_error(orig_tensor, qdq_tensor):
+    cos_sim = np.inner(orig_tensor.flatten(), qdq_tensor.flatten()) / (
+        np.linalg.norm(orig_tensor.flatten())
+        * np.linalg.norm(qdq_tensor.flatten())
+    )
+    return cos_sim
+
+
+def move_persistable_var_to_global_block(program):
+    # Move sub blocks persistable var to global block
+    global_block = program.global_block()
+    for _op in global_block.ops:
+        if _op.type == "while":
+            _block_id = _op.attr("sub_block").id
+            _block = program.block(_block_id)
+            persistables = []
+            for _name, _var in _block.vars.items():
+                if _var.persistable:
+                    global_block._clone_variable(_var)
+                    persistables.append(_name)
+            for _name in persistables:
+                _block._remove_var(_name)
+            persistables.extend(_op.input('X'))
+            _op.desc.set_input("X", persistables)
+
+
+def l2_loss(gt, pred):
+    return ((gt - pred) ** 2).mean()
+
+
+class tqdm:
+    def __init__(self, total, bar_format='Loading|{bar}', ncols=80):
+        self.total = total
+        self.bar_format = bar_format
+        self.ncols = ncols
+        self.n = 0
+
+    def update(self, n=1):
+        self.n += n
+        a = "=" * round((self.n / self.total) * self.ncols)
+        b = " " * (self.ncols - len(a))
+        prefix = self.bar_format.split('|')[0]
+        sys.stderr.write(f"\r{prefix}|{a}=>{b}| {self.n}/{self.total}")
+        sys.stderr.flush()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        sys.stderr.write('\n')