from typing import Tuple, Union
import numpy as np
from ... import module as Float
from ...core.tensor import dtype
from ...functional import expand_dims, squeeze
from ...functional.quantized import batch_conv_bias_activation
from ...tensor import Parameter
from ..qat import batch_matmul_activation as QAT
from .module import QuantizedModule
[文档]class BatchMatMulActivation(Float.BatchMatMulActivation, QuantizedModule):
r"""Quantized version of :class:`~.qat.BatchMatMulActivation`."""
def __init__(
self,
batch: int,
in_features: int,
out_features: int,
bias: bool = True,
nonlinear_mode="identity",
dtype=None,
**kwargs
):
super().__init__(batch, in_features, out_features, bias, **kwargs)
self.output_dtype = dtype
def calc_bmm_quantized(self, inp):
inp_scale = dtype.get_scale(inp.dtype)
w_scale = dtype.get_scale(self.weight.dtype)
bias_scale = inp_scale * w_scale
inp = expand_dims(inp, [-1])
res = batch_conv_bias_activation(
inp,
self.weight,
self.bias.astype(dtype.qint32(bias_scale)),
dtype=self.output_dtype,
stride=1,
padding=0,
dilation=1,
groups=1,
nonlinear_mode=self.nonlinear_mode,
)
return squeeze(res, -1)
[文档] @classmethod
def from_qat_module(cls, qat_module: QAT.BatchMatMulActivation):
output_dtype = qat_module.get_activation_dtype()
qbmm = cls(
qat_module.batch,
qat_module.in_features,
qat_module.out_features,
qat_module.bias is not None,
dtype=output_dtype,
name=qat_module.name,
)
weight = qat_module.weight.astype(qat_module.get_weight_dtype())
weight = expand_dims(weight, [-1, -2])
qbmm.weight = Parameter(weight.numpy(), name=qat_module.weight.name)
if qat_module.bias is not None:
bias = qat_module.bias.reshape((1, qbmm.out_features, 1, 1))
qbmm.bias = Parameter(bias.numpy(), name=qat_module.bias.name)
else:
qbmm.bias = Parameter(
np.zeros((1, qbmm.out_features, 1, 1), dtype=np.float32)
)
return qbmm
def forward(self, inp):
return self.calc_bmm_quantized(inp)