megengine.quantization.utils 源代码

import abc
from enum import Enum
from functools import partial, update_wrapper, wraps
from typing import Union

import numpy as np

from .. import functional as F
from ..autodiff import Function
from ..core._imperative_rt.core2 import apply
from ..core.ops import builtin
from ..core.tensor.dtype import (
from ..tensor import Tensor

class Round(Function):
    r"""The functional round have no grad and can not use for quantization-aware-training.
    We use Function and STE(Straight-Through Estimator) to implement backward propagation.

    def forward(self, x):
        return F.round(x)

    def backward(self, output_grads):
        return output_grads

def tqt_forward(qmin, qmax, inp, scale):
    op = builtin.TQT(qmin=qmin, qmax=qmax)
    (output,) = apply(op, inp, scale)
    return output

def lsq_forward(qmin, qmax, inp, step_size, zero_point=None, scale_grad=None):
    if zero_point is None:
        zero_point = Tensor([0.0], dtype=np.float32)
    if scale_grad is None:
        scale_grad = Tensor([1.0], dtype=np.float32)
    op = builtin.LSQ(qmin=qmin, qmax=qmax)
    (output,) = apply(op, inp, step_size, zero_point, scale_grad)
    return output

def register_method_to_class(cls):
    def decorator(func):
        def wrapper(self, *args, **kwargs):
            return func(self, *args, **kwargs)

        if isinstance(func, partial):
            update_wrapper(func, func.func)
        setattr(cls, func.__name__, wrapper)
        return func

    return decorator

[文档]class QuantMode(Enum): r"""Quantization mode enumerate class.""" SYMMERTIC = 1 ASYMMERTIC = 2
[文档]class QParams: r"""To standardize FakeQuant, Observer and Tensor's qparams format. If custom qparams is needed, inherit this class and add custom ``__slots__``. """ __slots__ = "mode", "dtype_meta", "scale", "zero_point" def __init__( self, mode: QuantMode, dtype_meta: QuantDtypeMeta, scale: Tensor, zero_point: Tensor, ): self.mode = mode self.dtype_meta = dtype_meta self.scale = scale self.zero_point = zero_point def update(self, qparams: "QParams"): for key in self.__slots__: setattr(self, key, getattr(qparams, key)) def __eq__(self, other): if len(self.__slots__) != len(other.__slots__): return False for key in self.__slots__: if not hasattr(other, key) or getattr(self, key) != getattr(other, key): return False return True def __repr__(self): content = ", ".join( ["{}={}".format(key, getattr(self, key)) for key in self.__slots__] ) return "QParams({})".format(content)
class LSQParams(QParams): r"""LSQ qparams with extra grad_scale slot.""" __slots__ = "mode", "dtype_meta", "scale", "zero_point", "grad_scale" def __init__( self, mode: QuantMode, dtype_meta: QuantDtypeMeta, scale: Tensor, zero_point: Tensor, grad_scale: Tensor, ): super().__init__(mode, dtype_meta, scale, zero_point) self.grad_scale = grad_scale class QParamsModuleMixin(abc.ABC): def get_quantized_dtype(self): qparams = self.get_qparams() dtype = qparams.dtype_meta scale = float(qparams.scale.numpy()) if qparams.scale is not None else None zero_point = ( int(qparams.zero_point.numpy()) if qparams.zero_point is not None else None ) return create_quantized_dtype(dtype, scale, zero_point) @abc.abstractmethod def get_qparams(self) -> QParams: pass _builtin_qparams = { QuantMode.SYMMERTIC: partial(QParams, mode=QuantMode.SYMMERTIC), QuantMode.ASYMMERTIC: partial(QParams, mode=QuantMode.ASYMMERTIC), }
[文档]def create_qparams( mode: QuantMode = QuantMode.SYMMERTIC, dtype_meta: Union[str, QuantDtypeMeta] = None, scale: Tensor = None, zero_point: Tensor = None, ): r""" Args: mode: QuantMode: dtype_meta: Union[str: QuantDtypeMeta]: scale: Tensor: zero_point: Tensor: """ if isinstance(dtype_meta, str): dtype_meta = _builtin_quant_dtypes[dtype_meta] if mode is None: return QParams(mode, dtype_meta, scale, zero_point) assert isinstance(mode, QuantMode) return _builtin_qparams[mode]( dtype_meta=dtype_meta, scale=scale, zero_point=zero_point )
[文档]def fake_quant_tensor(inp: Tensor, qparams: QParams) -> Tensor: """Apply fake quantization to the inp tensor. Args: inp: the input tensor which need to be faked. qparams: to get mode, qmin, qmax, scale and zero_point from. """ scale = qparams.scale if qparams.mode == QuantMode.ASYMMERTIC: zero_point = qparams.zero_point else: zero_point = Tensor([0.0], dtype=np.float32) qmin = qparams.dtype_meta.qmin qmax = qparams.dtype_meta.qmax op = builtin.FakeQuant(qmin=qmin, qmax=qmax) return apply(op, inp, scale, zero_point)[0]
[文档]def fake_quant_bias(bias: Tensor, inp: Tensor, w_qat: Tensor) -> Tensor: """Apply fake quantization to bias, with the special scale from input tensor and weight tensor, the quantized type set to qint32 also. Args: bias: the bias tensor which need to be faked. inp: the input tensor which contain the quantization parameters. w_qat: the weight tensor which contain the quantization parameters. Warning: Only work for symmetric quantization method now. """ b_qat = bias if ( getattr(inp, "qparams", None) is not None and getattr(w_qat, "qparams", None) is not None and bias is not None ): inp_params = inp.qparams w_params = w_qat.qparams if inp_params.scale is None or w_params.scale is None: return b_qat if inp_params.mode != QuantMode.SYMMERTIC: return b_qat # TODO: support different mode if inp_params.mode != w_params.mode: return b_qat # TODO: support quint8 dtype. if inp_params.dtype_meta.np_dtype_str != "int8": return b_qat if w_params.dtype_meta.np_dtype_str != "int8": return b_qat # use the same mode with weight. # TODO: avoid hardcode b_dtype = _builtin_quant_dtypes["qint32"] b_param = create_qparams( w_params.mode, b_dtype, scale=inp_params.scale * w_params.scale ) b_qat = fake_quant_tensor(bias, b_param) b_qat.qparams.update(b_param) return b_qat