mindspore 1.10.0__cp38-cp38-win_amd64.whl → 2.0.0rc1__cp38-cp38-win_amd64.whl

mindspore/compression/quant/qat.py

@@ -1,636 +0,0 @@
-# Copyright 2020-2022 Huawei Technologies Co., Ltd
-#
-# 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.
-# ============================================================================
-"""
-Quantization aware training
-
-User can use quantization aware to train a model. MindSpore supports quantization aware training,
-which models quantization errors in both the forward and backward passes using fake-quantization
-operations. Note that the entire computation is carried out in floating point. At the end of quantization
-aware training, MindSpore provides conversion functions to convert the trained model into lower precision.
-
-Note: This is an experimental interface that is subject to change and/or deletion.
-"""
-
-from __future__ import absolute_import
-import re
-import numpy as np
-import mindspore.context as context
-from mindspore import log as logger
-from mindspore import nn, ops
-from mindspore._checkparam import Validator, Rel
-from mindspore.nn.layer import quant
-from mindspore.ops import functional as F
-from ..common import QuantDtype
-from .quantizer import Quantizer, OptimizeOption
-from .quant_utils import compute_kl_threshold
-
-
-__all__ = ["QuantizationAwareTraining", "create_quant_config"]
-
-
-def create_quant_config(quant_observer=(nn.FakeQuantWithMinMaxObserver, nn.FakeQuantWithMinMaxObserver),
-                        quant_delay=(0, 0),
-                        quant_dtype=(QuantDtype.INT8, QuantDtype.INT8),
-                        per_channel=(False, False),
-                        symmetric=(False, False),
-                        narrow_range=(False, False),
-                        mode="DEFAULT"):
-    r"""
-    Config the observer type of weights and data flow with quant parameters.
-
-    Args:
-        quant_observer (Union[Observer, list, tuple]): The types of observer for quantization. The first element
-            applies to weights and the second applies to data flow. Currently, only
-            :class:`FakeQuantWithMinMaxObserver` supported.
-            Default: (nn.FakeQuantWithMinMaxObserver, nn.FakeQuantWithMinMaxObserver).
-        quant_delay (Union[int, list, tuple]): Number of steps after which weights and activations are quantized
-            during train and eval. The first element represents weights and the second element represents data flow.
-            Default: (0, 0).
-        quant_dtype (Union[QuantDtype, list, tuple]): Datatype used to quantize weights and activations. The first
-            element represents weights and the second element represents data flow.
-            Default: (QuantDtype.INT8, QuantDtype.INT8).
-        per_channel (Union[bool, list, tuple]):  Quantization granularity based on layer or on channel. If `True`
-            then base on per channel, otherwise base on per layer. The first element represents weights
-            and the second element represents data flow, and the second element must be `False` now.
-            Default: (False, False).
-        symmetric (Union[bool, list, tuple]): Whether the quantization algorithm is symmetric or not. If `True` then
-            base on symmetric, otherwise base on asymmetric. The first element represents weights and the second
-            element represents data flow. Default: (False, False).
-        narrow_range (Union[bool, list, tuple]): Whether the quantization algorithm uses narrow range or not.
-            The first element represents weights and the second element represents data flow.
-            Default: (False, False).
-        mode (str): Optional quantization mode, currently only `DEFAULT`(QAT) and `LEARNED_SCALE` are supported.
-            Default: "DEFAULT".
-
-    Returns:
-        QuantConfig, contains the observer type of weight and activation.
-
-    Raises:
-        ValueError: If the second element of `per_channel` is not `False`.
-    """
-    if per_channel[-1]:
-        raise ValueError("Arg 'per_channel' second element must be 'False'.")
-    weight_observer = quant_observer[0].partial_init(quant_delay=quant_delay[0], quant_dtype=quant_dtype[0],
-                                                     per_channel=per_channel[0], symmetric=symmetric[0],
-                                                     narrow_range=narrow_range[0], mode=mode)
-    act_observer = quant_observer[-1].partial_init(quant_delay=quant_delay[-1], quant_dtype=quant_dtype[-1],
-                                                   per_channel=per_channel[-1], symmetric=symmetric[-1],
-                                                   narrow_range=narrow_range[-1], mode=mode)
-    return quant.QuantConfig(weight=weight_observer, activation=act_observer)
-
-
-class _AddFakeQuantInput(nn.Cell):
-    """
-    Add FakeQuant OP at input of the network. Only support one input case.
-    """
-
-    def __init__(self, network, quant_delay=0):
-        super(_AddFakeQuantInput, self).__init__(auto_prefix=False)
-        self.fake_quant_input = quant.FakeQuantWithMinMaxObserver(min_init=-6, max_init=6,
-                                                                  quant_delay=quant_delay, ema=True)
-        self.fake_quant_input.update_parameters_name('fake_quant_input.')
-        self.network = network
-
-    def construct(self, data):
-        data = self.fake_quant_input(data)
-        output = self.network(data)
-        return output
-
-
-class _AddFakeQuantAfterSubCell(nn.Cell):
-    """
-    Add FakeQuant OP after of the sub Cell.
-    """
-
-    def __init__(self, subcell, **kwargs):
-        super(_AddFakeQuantAfterSubCell, self).__init__(auto_prefix=False)
-        self.subcell = subcell
-        self.mode = "DEFAULT"
-        self.max_init = 6
-        self.min_init = -6
-
-        if kwargs.get("optimize_option") is not None and OptimizeOption.LEARNED_SCALE in kwargs["optimize_option"]:
-            self.mode = "LEARNED_SCALE"
-            self.max_init = 16
-            self.min_init = -16
-
-        self.fake_quant_act = quant.FakeQuantWithMinMaxObserver(min_init=self.min_init,
-                                                                max_init=self.max_init,
-                                                                ema=True,
-                                                                quant_dtype=kwargs.get("quant_dtype"),
-                                                                quant_delay=kwargs.get("quant_delay"),
-                                                                per_channel=kwargs.get("per_channel"),
-                                                                symmetric=kwargs.get("symmetric"),
-                                                                narrow_range=kwargs.get("narrow_range"),
-                                                                mode=self.mode)
-
-    def construct(self, *data):
-        output = self.subcell(*data)
-        output = self.fake_quant_act(output)
-        return output
-
-
-class QuantizationAwareTraining(Quantizer):
-    r"""
-    Quantizer for quantization aware training.
-
-    Args:
-        bn_fold (bool): Whether to use bn fold ops for simulation inference operation. Default: True.
-        freeze_bn (int): Number of steps after which BatchNorm OP parameters fixed to global mean and variance.
-            Default: 1e7.
-        quant_delay (Union[int, list, tuple]): Number of steps after which weights and activations are quantized
-            during train and eval. The first element represents weights and the second element represents data flow.
-            Default: (0, 0).
-        quant_dtype (Union[QuantDtype, list, tuple]): Datatype used to quantize weights and activations. The first
-            element represents weights and the second element represents data flow. It is necessary to consider the
-            precision support of hardware devices in the practical quantization infer scenario.
-            Default: (QuantDtype.INT8, QuantDtype.INT8).
-        per_channel (Union[bool, list, tuple]):  Quantization granularity based on layer or on channel. If `True`
-            then base on per channel, otherwise base on per layer. The first element represents weights and the
-            second element represents data flow, and the second element must be `False` now. Default: (False, False).
-        symmetric (Union[bool, list, tuple]): Whether the quantization algorithm is symmetric or not. If `True` then
-            base on symmetric, otherwise base on asymmetric. The first element represents weights and the second
-            element represents data flow. Default: (False, False).
-        narrow_range (Union[bool, list, tuple]): Whether the quantization algorithm uses narrow range or not.
-            The first element represents weights and the second element represents data flow.
-            Default: (False, False).
-        optimize_option (Union[OptimizeOption, list, tuple]): Specifies the quant algorithm and options, currently
-            only support `QAT` and `LEARNED_SCALE` (Note that, if both `QAT` and `LEARNED_SCALE` are configured,
-            `LEARNED_SCALE` has a higher priority. `LEARNED_SCALE` currently only work under some constraints, which
-            includes: freeze_bn=0, quant_delay=0, symmetric=True, narrow_range=True, More specifically, for operators
-            such as Relu and Relu6, which only have positive values, we add a negative truncation to optimize this
-            scenario, and narrow_range will automatically match to False). Default: OptimizeOption.QAT.
-        one_conv_fold (bool): Whether to use one conv bn fold ops for simulation inference operation. Default: True.
-
-    Supported Platforms:
-        ``Ascend`` ``GPU``
-
-    Raises:
-        TypeError: If the element of `quant_delay` or `freeze_bn` is not int.
-        TypeError: If `bn_fold`, `one_conv_fold` or the element of `per_channel`, `symmetric`, `narrow_range`
-            is not bool.
-        TypeError: If the element of `quant_dtype` is not `QuantDtype`.
-        ValueError: If the length of `quant_delay`, `quant_dtype`, `per_channel`, `symmetric` or `narrow_range` is
-            not less than 2.
-        ValueError: If the `optimize_option` is `LEARNED_SCALE` and `freeze_bn` is not equal to 0.
-        ValueError: If the `optimize_option` is `LEARNED_SCALE` and `symmetric` is not (True, True).
-        ValueError: If the `optimize_option` is `LEARNED_SCALE` and `narrow_range` is not (True, True).
-        ValueError: If the `optimize_option` is `LEARNED_SCALE` and `quant_delay` is not (0, 0).
-
-    Examples:
-        >>> from mindspore.compression.quant import QuantizationAwareTraining
-        >>> from mindspore import nn
-        >>> class LeNet5(nn.Cell):
-        ...     def __init__(self, num_class=10, channel=1):
-        ...         super(LeNet5, self).__init__()
-        ...         self.type = "fusion"
-        ...         self.num_class = num_class
-        ...
-        ...         # change `nn.Conv2d` to `nn.Conv2dBnAct`
-        ...         self.conv1 = nn.Conv2dBnAct(channel, 6, 5, pad_mode='valid', activation='relu')
-        ...         self.conv2 = nn.Conv2dBnAct(6, 16, 5, pad_mode='valid', activation='relu')
-        ...         # change `nn.Dense` to `nn.DenseBnAct`
-        ...         self.fc1 = nn.DenseBnAct(16 * 5 * 5, 120, activation='relu')
-        ...         self.fc2 = nn.DenseBnAct(120, 84, activation='relu')
-        ...         self.fc3 = nn.DenseBnAct(84, self.num_class)
-        ...
-        ...         self.max_pool2d = nn.MaxPool2d(kernel_size=2, stride=2)
-        ...         self.flatten = nn.Flatten()
-        ...
-        ...     def construct(self, x):
-        ...         x = self.conv1(x)
-        ...         x = self.max_pool2d(x)
-        ...         x = self.conv2(x)
-        ...         x = self.max_pool2d(x)
-        ...         x = self.flatten(x)
-        ...         x = self.fc1(x)
-        ...         x = self.fc2(x)
-        ...         x = self.fc3(x)
-        ...         return x
-        ...
-        >>> net = LeNet5()
-        >>> quantizer = QuantizationAwareTraining(bn_fold=False, per_channel=[True, False], symmetric=[True, False])
-        >>> net_qat = quantizer.quantize(net)
-    """
-    __quant_op_name__ = ["Add", "Sub", "Mul", "RealDiv", "ReduceMean"]
-
-    def __init__(self,
-                 bn_fold=True,
-                 freeze_bn=10000000,
-                 quant_delay=(0, 0),
-                 quant_dtype=(QuantDtype.INT8, QuantDtype.INT8),
-                 per_channel=(False, False),
-                 symmetric=(False, False),
-                 narrow_range=(False, False),
-                 optimize_option=OptimizeOption.QAT,
-                 one_conv_fold=True):
-        """Init for QuantizationAwareTraining quantizer"""
-        super(QuantizationAwareTraining, self).__init__(optimize_option=optimize_option)
-
-        def convert2list(name, value):
-            if not isinstance(value, list) and not isinstance(value, tuple):
-                value = [value]
-            elif len(value) > 2:
-                raise ValueError("input `{}` len should less then 2".format(name))
-            return value
-
-        quant_delay_list = convert2list("quant delay", quant_delay)
-        quant_dtype_list = convert2list("quant dtype", quant_dtype)
-        per_channel_list = convert2list("per channel", per_channel)
-        symmetric_list = convert2list("symmetric", symmetric)
-        narrow_range_list = convert2list("narrow range", narrow_range)
-
-        self.weight_qdelay = Validator.check_non_negative_int(quant_delay_list[0], "quant delay")
-        self.act_qdelay = Validator.check_int(quant_delay_list[-1], 0, Rel.GE, "quant delay")
-        self.bn_fold = Validator.check_bool(bn_fold, "bn fold")
-        self.freeze_bn = Validator.check_non_negative_int(freeze_bn, "freeze bn")
-        self.weight_dtype = Validator.check_isinstance("weights dtype", quant_dtype_list[0], QuantDtype)
-        self.act_dtype = Validator.check_isinstance("activations dtype", quant_dtype_list[-1], QuantDtype)
-        self.weight_channel = Validator.check_bool(per_channel_list[0], "per channel")
-        self.act_channel = Validator.check_bool(per_channel_list[-1], "per channel")
-        self.weight_symmetric = Validator.check_bool(symmetric_list[0], "symmetric")
-        self.act_symmetric = Validator.check_bool(symmetric_list[-1], "symmetric")
-        self.weight_range = Validator.check_bool(narrow_range_list[0], "narrow range")
-        self.act_range = Validator.check_bool(narrow_range_list[-1], "narrow range")
-        self.one_conv_fold = Validator.check_bool(one_conv_fold, "one conv fold")
-        self._convert_method_map = {nn.Conv2dBnAct: self._convert_conv,
-                                    nn.DenseBnAct: self._convert_dense}
-        self.mode = "DEFAULT"
-        if OptimizeOption.LEARNED_SCALE in self.optimize_option:
-            self.mode = "LEARNED_SCALE"
-            if not self.weight_symmetric or not self.act_symmetric:
-                raise ValueError("OptimizeOption.LEARNED_SCALE currently only support "
-                                 "symmetric=(True, True) for quant")
-            if not self.weight_range or not self.act_range:
-                raise ValueError("OptimizeOption.LEARNED_SCALE currently only support narrow_range=(True, True) "
-                                 "for quant")
-            if self.freeze_bn != 0:
-                raise ValueError("OptimizeOption.LEARNED_SCALE currently only support freeze_bn equal to 0, "
-                                 "but get freeze_bn={}".format(self.freeze_bn))
-            if self.weight_qdelay != 0 or self.act_qdelay != 0:
-                raise ValueError("OptimizeOption.LEARNED_SCALE currently only support quant_delay=(0, 0)")
-        self.quant_config = create_quant_config(quant_delay=quant_delay_list,
-                                                quant_dtype=quant_dtype_list,
-                                                per_channel=per_channel_list,
-                                                symmetric=symmetric_list,
-                                                narrow_range=narrow_range_list,
-                                                mode=self.mode)
-        self.eps = 1e-5
-
-    @staticmethod
-    def _convert_op_name(name):
-        pattern = re.compile(r'([A-Z]{1})')
-        name_new = re.sub(pattern, r'_\1', name).lower()
-        if name_new[0] == '_':
-            name_new = name_new[1:]
-        return name_new
-
-    def quantize(self, network):
-        """
-        Quant API to convert input network to a quantization aware training network.
-
-        Note:
-            Please refer to the Examples of class: `mindspore.compression.quant.QuantizationAwareTraining`.
-
-        Args:
-            network (Cell): network to be quantized.
-
-        Returns:
-            Cell, a quantization aware training network.
-
-        Raises:
-            KeyError: If the `device_target` set in context is not in `support_device`.
-        """
-
-        logger.warning("The compression module is deprecated and may not be supported in later version, please use "
-                       "MindSpore Golden Stick(https://gitee.com/mindspore/golden-stick) instead.")
-        support_device = ["Ascend", "GPU"]
-        if context.get_context('device_target') not in support_device:
-            raise KeyError("Unsupported {} device target.".format(context.get_context('device_target')))
-
-        if OptimizeOption.QAT in self.optimize_option or OptimizeOption.LEARNED_SCALE in self.optimize_option:
-            network.update_cell_prefix()
-            network = self._convert_subcells2quant(network)
-            network.update_cell_type("quant")
-        return network
-
-    def _convert_subcells2quant(self, network):
-        """
-        convert sub cell like `Conv2dBnAct` and `DenseBnAct` to quant cell
-        """
-        cells = network.name_cells()
-        change = False
-        for name in cells:
-            subcell = cells[name]
-            if subcell == network:
-                continue
-            if isinstance(subcell, (nn.Conv2dBnAct, nn.DenseBnAct)):
-                prefix = subcell.param_prefix
-                new_subcell = self._convert_method_map[type(subcell)](subcell)
-                new_subcell.update_parameters_name(prefix + '.')
-                network.insert_child_to_cell(name, new_subcell)
-                change = True
-            else:
-                self._convert_subcells2quant(subcell)
-        if isinstance(network, nn.SequentialCell) and change:
-            network.cell_list = list(network.cells())
-
-        # add FakeQuant OP after OP in white list, but not including those wrapped in the below quantization cell.
-        if isinstance(network, (nn.FakeQuantWithMinMaxObserver,
-                                nn.Conv2dBnFoldQuantOneConv,
-                                nn.Conv2dBnFoldQuant,
-                                nn.Conv2dBnWithoutFoldQuant,
-                                nn.Conv2dQuant,
-                                nn.DenseQuant,
-                                nn.ActQuant,
-                                nn.TensorAddQuant,
-                                nn.MulQuant)):
-            return network
-
-        add_list = []
-        for name in network.__dict__:
-            if name[0] == '_':
-                continue
-            attr = network.__dict__[name]
-            if isinstance(attr, ops.Primitive) and attr.name in self.__quant_op_name__:
-                add_list.append((name, attr))
-        for name, prim_op in add_list:
-            prefix = name
-            add_quant = _AddFakeQuantAfterSubCell(prim_op,
-                                                  quant_dtype=self.act_dtype,
-                                                  quant_delay=self.act_qdelay,
-                                                  per_channel=self.act_channel,
-                                                  symmetric=self.act_symmetric,
-                                                  narrow_range=self.act_range,
-                                                  optimize_option=self.optimize_option)
-            if network.param_prefix:
-                prefix = '.'.join([network.param_prefix, prefix])
-            add_quant.update_parameters_name(prefix + '.')
-            del network.__dict__[name]
-            network.insert_child_to_cell(name, add_quant)
-        return network
-
-    def _convert_conv(self, subcell):
-        """
-        convert Conv2d cell to quant cell
-        """
-        min_init = -6
-        max_init = 6
-        if self.eps == 0:
-            raise ValueError("`epsilon` is zero may lead to divide zero error")
-        if OptimizeOption.LEARNED_SCALE in self.optimize_option:
-            subcell_weight_para = subcell.conv.weight.data.asnumpy()
-            if subcell.has_bn:
-                scale_factor = (subcell.batchnorm.gamma.data.asnumpy() /
-                                np.sqrt(subcell.batchnorm.moving_variance.data.asnumpy() + self.eps))
-                subcell_weight_para = subcell_weight_para * scale_factor.reshape(-1, 1, 1, 1)
-            min_init, max_init = self._kl_init(subcell_weight_para, self.weight_dtype)
-        self.quant_config = self.quant_config._replace(
-            weight=self.quant_config.weight.partial_init(min_init=min_init, max_init=max_init))
-
-        conv_inner = subcell.conv
-        if subcell.has_bn:
-            bn_inner = subcell.batchnorm
-            if self.bn_fold:
-                if self.one_conv_fold:
-                    conv_inner = quant.Conv2dBnFoldQuantOneConv(conv_inner.in_channels,
-                                                                conv_inner.out_channels,
-                                                                kernel_size=conv_inner.kernel_size,
-                                                                stride=conv_inner.stride,
-                                                                pad_mode=conv_inner.pad_mode,
-                                                                padding=conv_inner.padding,
-                                                                dilation=conv_inner.dilation,
-                                                                group=conv_inner.group,
-                                                                eps=bn_inner.eps,
-                                                                momentum=1 - bn_inner.momentum,
-                                                                has_bias=conv_inner.has_bias,
-                                                                bias_init=conv_inner.bias_init,
-                                                                quant_config=self.quant_config,
-                                                                quant_dtype=self.weight_dtype,
-                                                                fake=True)
-                else:
-                    conv_inner = quant.Conv2dBnFoldQuant(conv_inner.in_channels,
-                                                         conv_inner.out_channels,
-                                                         kernel_size=conv_inner.kernel_size,
-                                                         stride=conv_inner.stride,
-                                                         pad_mode=conv_inner.pad_mode,
-                                                         padding=conv_inner.padding,
-                                                         dilation=conv_inner.dilation,
-                                                         group=conv_inner.group,
-                                                         eps=bn_inner.eps,
-                                                         momentum=1 - bn_inner.momentum,
-                                                         has_bias=conv_inner.has_bias,
-                                                         bias_init=conv_inner.bias_init,
-                                                         freeze_bn=self.freeze_bn,
-                                                         quant_config=self.quant_config,
-                                                         quant_dtype=self.weight_dtype,
-                                                         fake=True)
-                # change original network Batch Normalization OP parameters to quant network
-                conv_inner.gamma = subcell.batchnorm.gamma
-                conv_inner.beta = subcell.batchnorm.beta
-                conv_inner.moving_mean = subcell.batchnorm.moving_mean
-                conv_inner.moving_variance = subcell.batchnorm.moving_variance
-            else:
-                conv_inner = quant.Conv2dBnWithoutFoldQuant(conv_inner.in_channels,
-                                                            conv_inner.out_channels,
-                                                            kernel_size=conv_inner.kernel_size,
-                                                            stride=conv_inner.stride,
-                                                            pad_mode=conv_inner.pad_mode,
-                                                            padding=conv_inner.padding,
-                                                            dilation=conv_inner.dilation,
-                                                            group=conv_inner.group,
-                                                            eps=bn_inner.eps,
-                                                            momentum=1 - bn_inner.momentum,
-                                                            has_bias=conv_inner.has_bias,
-                                                            bias_init=conv_inner.bias_init,
-                                                            quant_config=self.quant_config)
-                # change original network Batch Normalization OP parameters to quant network
-                conv_inner.batchnorm.gamma = subcell.batchnorm.gamma
-                conv_inner.batchnorm.beta = subcell.batchnorm.beta
-                conv_inner.batchnorm.moving_mean = subcell.batchnorm.moving_mean
-                conv_inner.batchnorm.moving_variance = subcell.batchnorm.moving_variance
-            del subcell.batchnorm
-            subcell.batchnorm = None
-            subcell.has_bn = False
-        else:
-            conv_inner = quant.Conv2dQuant(conv_inner.in_channels, conv_inner.out_channels,
-                                           kernel_size=conv_inner.kernel_size, stride=conv_inner.stride,
-                                           pad_mode=conv_inner.pad_mode, padding=conv_inner.padding,
-                                           dilation=conv_inner.dilation, group=conv_inner.group,
-                                           has_bias=conv_inner.has_bias, quant_config=self.quant_config,
-                                           quant_dtype=self.weight_dtype)
-        # change original network Conv2D OP parameters to quant network
-        conv_inner.weight = subcell.conv.weight
-        if subcell.conv.has_bias:
-            conv_inner.bias = subcell.conv.bias
-        subcell.conv = conv_inner
-        if subcell.has_act and subcell.activation is not None:
-            subcell.activation = self._convert_activation(subcell.activation)
-        elif subcell.after_fake:
-            subcell.has_act = True
-            subcell.activation = _AddFakeQuantAfterSubCell(F.identity, quant_dtype=self.act_dtype,
-                                                           quant_delay=self.act_qdelay, per_channel=self.act_channel,
-                                                           symmetric=self.act_symmetric, narrow_range=self.act_range,
-                                                           optimize_option=self.optimize_option)
-        return subcell
-
-    def _convert_dense(self, subcell):
-        """
-        convert dense cell to quant cell
-        """
-        min_init = -6
-        max_init = 6
-        if self.eps == 0:
-            raise ValueError("`epsilon` is zero may lead to divide zero error")
-        if OptimizeOption.LEARNED_SCALE in self.optimize_option:
-            subcell_weight_para = subcell.dense.weight.data.asnumpy()
-            if subcell.has_bn:
-                scale_factor = (subcell.batchnorm.gamma.data.asnumpy() /
-                                np.sqrt(subcell.batchnorm.moving_variance.data.asnumpy() + self.eps))
-                subcell_weight_para = subcell_weight_para * scale_factor.reshape(-1, 1, 1, 1)
-            min_init, max_init = self._kl_init(subcell_weight_para, self.weight_dtype)
-        self.quant_config = self.quant_config._replace(
-            weight=self.quant_config.weight.partial_init(min_init=min_init, max_init=max_init))
-
-        dense_inner = subcell.dense
-        dense_inner = quant.DenseQuant(dense_inner.in_channels,
-                                       dense_inner.out_channels,
-                                       has_bias=dense_inner.has_bias,
-                                       quant_config=self.quant_config,
-                                       quant_dtype=self.weight_dtype)
-        # change original network Dense OP parameters to quant network
-        dense_inner.weight = subcell.dense.weight
-        if subcell.dense.has_bias:
-            dense_inner.bias = subcell.dense.bias
-        subcell.dense = dense_inner
-        if subcell.has_act and subcell.activation is not None:
-            subcell.activation = self._convert_activation(subcell.activation)
-        elif subcell.after_fake:
-            subcell.has_act = True
-            subcell.activation = _AddFakeQuantAfterSubCell(F.identity,
-                                                           quant_dtype=self.act_dtype,
-                                                           quant_delay=self.act_qdelay,
-                                                           per_channel=self.act_channel,
-                                                           symmetric=self.act_symmetric,
-                                                           narrow_range=self.act_range,
-                                                           optimize_option=self.optimize_option)
-        return subcell
-
-    def _convert_activation(self, activation):
-        """
-        convert activation cell to quant cell
-        """
-        act_class = activation.__class__
-        act_list = [nn.ReLU, nn.ReLU6, nn.Sigmoid]
-        act_list_with_fake_before = [nn.LeakyReLU, nn.HSigmoid, nn.HSwish]
-
-        if act_class in act_list:
-            return quant.ActQuant(activation=activation,
-                                  quant_config=self.quant_config,
-                                  quant_dtype=self.act_dtype)
-        if act_class in act_list_with_fake_before:
-            return quant.ActQuant(activation=activation,
-                                  ema=True,
-                                  fake_before=True,
-                                  quant_config=self.quant_config,
-                                  quant_dtype=self.act_dtype)
-        raise ValueError("Unsupported activation in auto quant: ", act_class)
-
-    def _kl_init(self, subcell_weight_para, weight_dtype):
-        """
-        Calculate the value of max_init and min_init with compute_kl_threshold.
-        """
-        if self.weight_channel:
-            max_init = [compute_kl_threshold(weight_para_each, weight_dtype)
-                        for weight_para_each in subcell_weight_para]
-            min_init = [-x for x in max_init]
-        else:
-            max_init = [compute_kl_threshold(subcell_weight_para, weight_dtype)]
-            min_init = [-x for x in max_init]
-        return min_init, max_init
-
-    def _set_mixed_bits(self, network, strategy):
-        r"""
-        Set network's quantization strategy, this function is currently only valid for `LEARNED_SCALE`
-        optimize_option.
-
-        Args:
-            network (Cell): Input network.
-            strategy (list): The quantization strategy for layers that need to be quantified (eg. [[8], [8],
-                ..., [6], [4], [8]]), currently only the quant_dtype for weights of the dense layer and the
-                convolution layer is supported.
-
-        Returns:
-            Cell, a network with mixed bit strategy configured.
-
-        Raises:
-            ValueError: If `OptimizeOption.LEARNED_SCALE` is not in `self.optimize_option`.
-        """
-        if OptimizeOption.LEARNED_SCALE not in self.optimize_option:
-            raise ValueError("The `_set_mixed_bits` function is currently only valid for `LEARNED_SCALE` "
-                             "optimize_option.")
-
-        quantizable_idx = []
-        pass_cell = None
-        for i, cell_and_name in enumerate(network.cells_and_names()):
-            cell = cell_and_name[1]
-            if isinstance(cell, (nn.Conv2dBnAct, nn.DenseBnAct)) and cell is not pass_cell:
-                quantizable_idx.append(i)
-
-        if len(quantizable_idx) != len(strategy):
-            raise ValueError("The dimension of quantifiable layers is not consistent with that of strategy.")
-
-        quantizable_layer_bit_dict = {idx: bit for idx, bit in zip(quantizable_idx, strategy)}
-        type_map = {
-            QuantDtype.INT2.num_bits: QuantDtype.INT2,
-            QuantDtype.INT3.num_bits: QuantDtype.INT3,
-            QuantDtype.INT4.num_bits: QuantDtype.INT4,
-            QuantDtype.INT5.num_bits: QuantDtype.INT5,
-            QuantDtype.INT6.num_bits: QuantDtype.INT6,
-            QuantDtype.INT7.num_bits: QuantDtype.INT7,
-            QuantDtype.INT8.num_bits: QuantDtype.INT8
-        }
-        if self.eps == 0:
-            raise ValueError("`epsilon` is zero may lead to divide zero error")
-        for i, cell_and_name in enumerate(network.cells_and_names()):
-            cell = cell_and_name[1]
-            if i not in quantizable_idx:
-                continue
-            else:
-                if isinstance(cell, (nn.Conv2dBnAct, nn.DenseBnAct)):
-                    cell.weight_dtype = type_map.get(quantizable_layer_bit_dict[i][0])
-                    if cell.weight_dtype is None:
-                        raise ValueError("Input strategy is invalid: ", quantizable_layer_bit_dict[i][0])
-                    if isinstance(cell, nn.Conv2dBnAct):
-                        subcell_weight_para = cell.conv.weight.data.asnumpy()
-                        if hasattr(cell.conv, 'gamma'):
-                            scale_factor = (cell.conv.gamma.data.asnumpy() /
-                                            np.sqrt(cell.conv.moving_variance.data.asnumpy() + self.eps))
-                            subcell_weight_para = subcell_weight_para * scale_factor.reshape(-1, 1, 1, 1)
-                        min_init, max_init = self._kl_init(subcell_weight_para, cell.weight_dtype)
-                        cell.conv.fake_quant_weight.reset(quant_dtype=cell.weight_dtype,
-                                                          min_init=min_init,
-                                                          max_init=max_init)
-                    elif isinstance(cell, nn.DenseBnAct):
-                        subcell_weight_para = cell.dense.weight.data.asnumpy()
-                        if hasattr(cell.dense, 'gamma'):
-                            scale_factor = (cell.dense.gamma.data.asnumpy() /
-                                            np.sqrt(cell.dense.moving_variance.data.asnumpy() + self.eps))
-                            subcell_weight_para = subcell_weight_para * scale_factor.reshape(-1, 1, 1, 1)
-                        min_init, max_init = self._kl_init(subcell_weight_para, cell.weight_dtype)
-                        cell.dense.fake_quant_weight.reset(quant_dtype=cell.weight_dtype,
-                                                           min_init=min_init,
-                                                           max_init=max_init)
-        return network