# Source code for tensorplay.ao.nn.quantized.linear Source: https://www.tensorplay.cn/docs/_modules/tensorplay/ao/nn/quantized/linear.html ``` """Linear inference with reusable integer weight storage and fused output transforms.""" import tensorplay as tp from tensorplay import nn, _C from .packed import PackedWeight, QuantizedState, quantize_module_weight, conversion_source __all__ = ['QuantizedLinear'] [docs] class QuantizedLinear(QuantizedState, nn.Module): _FUSED_RELU = False def __init__(self, in_features, out_features, input_scale=1., input_zero_point=0, qweight=None, weight_scales=None, weight_zero_points=None, bias=None, out_scale=1., out_zero_point=0): super().__init__() self.in_features, self.out_features = int(in_features), int(out_features) self.input_scale, self.input_zero_point = float(input_scale), int(input_zero_point) self.out_scale, self.out_zero_point = float(out_scale), int(out_zero_point) if qweight is None: qweight = _C.quantize_per_tensor(tp.zeros(out_features, in_features), 1., 0, tp.qint8) elif qweight.dtype == tp.int8: qweight = _C._make_per_channel_quantized_tensor(qweight, weight_scales.float(), weight_zero_points.long(), 0) self._packed_params = PackedWeight(qweight, bias) @property def qweight(self): return self._packed_params.weight @property def bias(self): return self._packed_params.bias @property def weight_scales(self): return self._packed_params.scales @property def weight_zero_points(self): return self._packed_params.zero_points @property def scale(self): return self.out_scale @scale.setter def scale(self, value): self.out_scale = float(value) @property def zero_point(self): return self.out_zero_point @zero_point.setter def zero_point(self, value): self.out_zero_point = int(value) def weight(self): return self.qweight def set_weight_bias(self, weight, bias): if tuple(weight.shape) != (self.out_features, self.in_features): raise ValueError('Linear weight shape does not match its declared dimensions') self._packed_params.set_weight_bias(weight, bias) def forward(self, x): if x.dtype == tp.int8: x = _C._make_per_tensor_quantized_tensor(x, self.input_scale, self.input_zero_point) if x.dtype not in (tp.qint8, tp.quint8): raise TypeError('Quantized linear requires signed or unsigned 8-bit quantized input') return self._packed_params.run(x, self.out_scale, self.out_zero_point, relu=self._FUSED_RELU) [docs] @classmethod def from_reference(cls, ref_qlinear, output_scale, output_zero_point): """Build a native quantized linear from a reference quantized module.""" qlinear = cls(ref_qlinear.in_features, ref_qlinear.out_features) qweight = ref_qlinear.get_quantized_weight() qlinear.set_weight_bias(qweight, ref_qlinear.bias) qlinear.out_scale = float(output_scale) qlinear.out_zero_point = int(output_zero_point) return qlinear def extra_repr(self): return f'in_features={self.in_features}, out_features={self.out_features}, scale={self.out_scale}, zero_point={self.out_zero_point}' @classmethod def from_float(cls, float_module, input_scale=None, input_zero_point=None, out_scale=None, out_zero_point=None, use_precomputed_fake_quant=False): source = conversion_source(float_module) if not isinstance(source, nn.Linear): raise TypeError('Linear conversion requires a floating linear module') if input_scale is None or input_zero_point is None: observer = getattr(float_module, 'input_activation_post_process', None) if observer is None: raise ValueError('Linear conversion requires input calibration or explicit parameters') input_scale, input_zero_point = observer.calculate_qparams() if out_scale is None or out_zero_point is None: observer = getattr(float_module, 'activation_post_process', None) if observer is None: raise ValueError('Linear conversion requires output calibration or explicit parameters') out_scale, out_zero_point = observer.calculate_qparams() owner = float_module if hasattr(float_module, 'weight_fake_quant') else source qweight = quantize_module_weight(owner, use_precomputed_fake_quant, weight=source.weight) return cls(source.in_features, source.out_features, input_scale, input_zero_point, qweight, bias=source.bias.detach() if source.bias is not None else None, out_scale=out_scale, out_zero_point=out_zero_point) ```