megengine.module.quantized.linear 源代码

import numpy as np

from ... import functional as F
from ... import module as Float
from ...core.tensor import dtype
from ...tensor import Parameter
from ..qat import linear as QAT
from .module import QuantizedModule

[文档]class Linear(QuantizedModule): r"""Quantized version of :class:`~.qat.Linear`.""" def __init__(self, dtype: np.dtype = None, **kwargs): super().__init__(**kwargs) self.weight = None self.bias = None self.output_dtype = dtype def calc_linear_quantized(self, inp, nonlinear_mode="identity"): if raise ValueError("quantized module only support inference.") assert nonlinear_mode in ["identity", "relu"] inp_scale = dtype.get_scale(inp.dtype) w_scale = dtype.get_scale(self.weight.dtype) bias_dtype = dtype.qint32(inp_scale * w_scale) ret = F.linear( inp, self.weight, None if self.bias is None else self.bias.astype(bias_dtype), ) ret = ret if self.output_dtype is None else ret.astype(self.output_dtype) if nonlinear_mode == "relu": ret = F.relu(ret) return ret def forward(self, inp): return self.calc_linear_quantized(inp)
[文档] @classmethod def from_qat_module(cls, qat_module: QAT.Linear): r""" Return a :class:`~.QuantizedModule` instance converted from a :class:`~.QATModule` instance. """ output_dtype = qat_module.get_activation_dtype() qmod = cls(dtype=output_dtype, = weight = qat_module.weight.astype(qat_module.get_weight_dtype()) qmod.weight = Parameter(weight.numpy(), if qat_module.bias is not None: qmod.bias = Parameter(qat_module.bias.numpy(), return qmod
class LinearRelu(Linear): r"""Quantized version of :class:`~.qat.LinearRelu`.""" def forward(self, inp): return self.calc_linear_quantized(inp, nonlinear_mode="relu")