Function fake_quant_bias

imperative/python/megengine/quantization/utils.py:189–234

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

(bias: Tensor, inp: Tensor, w_qat: Tensor)

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