Shortcuts

Source code for torch.ao.nn.intrinsic.quantized.modules.linear_relu

# mypy: allow-untyped-defs
import torch
import torch.ao.nn.quantized as nnq
import torch.ao.nn.intrinsic as nni
from torch.ao.nn.quantized.modules.utils import _quantize_weight

__all__ = [
    "LinearReLU",
    "LinearLeakyReLU",
    "LinearTanh",
]

[docs]class LinearReLU(nnq.Linear): r""" A LinearReLU module fused from Linear and ReLU modules We adopt the same interface as :class:`torch.ao.nn.quantized.Linear`. Attributes: Same as torch.ao.nn.quantized.Linear Examples:: >>> # xdoctest: +SKIP >>> m = nn.intrinsic.LinearReLU(20, 30) >>> input = torch.randn(128, 20) >>> output = m(input) >>> print(output.size()) torch.Size([128, 30]) """ _FLOAT_MODULE = nni.LinearReLU # type: ignore[assignment] def __init__(self, in_features, out_features, bias=True, dtype=torch.qint8): super().__init__(in_features, out_features, bias, dtype) def forward(self, x: torch.Tensor) -> torch.Tensor: return torch.ops.quantized.linear_relu( x, self._packed_params._packed_params, self.scale, self.zero_point) def _get_name(self): return 'QuantizedLinearReLU' @classmethod def from_float(cls, mod, use_precomputed_fake_quant=False): return super().from_float(mod, use_precomputed_fake_quant) @classmethod def from_reference(cls, ref_linear_relu, output_scale, output_zero_point): return super().from_reference(ref_linear_relu[0], output_scale, output_zero_point)
class LinearLeakyReLU(nnq.Linear): r""" For onednn backend only A LinearLeakyReLU module fused from Linear and LeakyReLU modules We adopt the same interface as :class:`torch.ao.nn.quantized.Linear`. Attributes: Same as torch.ao.nn.quantized.Linear + negative_slope Examples:: >>> # xdoctest: +SKIP >>> m = nn.intrinsic.LinearLeakyReLU(20, 30, 0.01) >>> input = torch.randn(128, 20) >>> output = m(input) >>> print(output.size()) torch.Size([128, 30]) """ _FLOAT_MODULE = nni.LinearLeakyReLU # type: ignore[assignment] def __init__(self, in_features, out_features, negative_slope, bias=True, dtype=torch.qint8): super().__init__(in_features, out_features, bias, dtype) self.negative_slope = negative_slope def forward(self, x: torch.Tensor) -> torch.Tensor: return torch.ops.quantized.linear_leaky_relu( x, self._packed_params._packed_params, self.scale, self.zero_point, self.negative_slope) def _get_name(self): return 'QuantizedLinearLeakyReLU' @classmethod def from_float(cls, mod, use_precomputed_fake_quant=False): assert type(mod) == nni.LinearLeakyReLU, 'Input float module should be LinearLeakyReLU' assert hasattr(mod, 'qconfig'), 'Input float module must have qconfig defined' activation_post_process = mod.activation_post_process leaky_relu = mod[1] mod = mod[0] weight_post_process = mod.qconfig.weight() weight_post_process(mod.weight) dtype = weight_post_process.dtype act_scale, act_zp = activation_post_process.calculate_qparams() # type: ignore[union-attr,operator] assert dtype == torch.qint8, 'Weight observer must have dtype torch.qint8' qweight = _quantize_weight(mod.weight.float(), weight_post_process) qlinear_leaky_relu = cls( mod.in_features, mod.out_features, leaky_relu.negative_slope, dtype=dtype) qlinear_leaky_relu.set_weight_bias(qweight, mod.bias) qlinear_leaky_relu.scale = float(act_scale) qlinear_leaky_relu.zero_point = int(act_zp) return qlinear_leaky_relu @classmethod def from_reference(cls, ref_mod, output_scale, output_zero_point): linear = ref_mod[0] leaky_relu = ref_mod[1] qlinear_leaky_relu = cls( linear.in_features, linear.out_features, leaky_relu.negative_slope) qweight = linear.get_quantized_weight() qlinear_leaky_relu.set_weight_bias(qweight, linear.bias) qlinear_leaky_relu.scale = float(output_scale) qlinear_leaky_relu.zero_point = int(output_zero_point) return qlinear_leaky_relu class LinearTanh(nnq.Linear): r""" A LinearTanh module fused from Linear and Tanh modules We adopt the same interface as :class:`torch.ao.nn.quantized.Linear`. Attributes: Same as torch.ao.nn.quantized.Linear Examples:: >>> # xdoctest: +SKIP >>> m = nn.intrinsic.LinearTanh(20, 30) >>> input = torch.randn(128, 20) >>> output = m(input) >>> print(output.size()) torch.Size([128, 30]) """ _FLOAT_MODULE = nni.LinearTanh # type: ignore[assignment] def __init__(self, in_features, out_features, bias=True, dtype=torch.qint8): super().__init__(in_features, out_features, bias, dtype) def forward(self, x: torch.Tensor) -> torch.Tensor: return torch.ops.quantized.linear_tanh( x, self._packed_params._packed_params, self.scale, self.zero_point) def _get_name(self): return 'QuantizedLinearTanh' @classmethod def from_float(cls, mod, use_precomputed_fake_quant=False): assert type(mod) == nni.LinearTanh, 'Input float module should be LinearTanh' assert hasattr(mod, 'qconfig'), 'Input float module must have qconfig defined' activation_post_process = mod.activation_post_process mod = mod[0] weight_post_process = mod.qconfig.weight() weight_post_process(mod.weight) dtype = weight_post_process.dtype act_scale, act_zp = activation_post_process.calculate_qparams() # type: ignore[union-attr,operator] assert dtype == torch.qint8, 'Weight observer must have dtype torch.qint8' qweight = _quantize_weight(mod.weight.float(), weight_post_process) qlinear_tanh = cls( mod.in_features, mod.out_features, dtype=dtype) qlinear_tanh.set_weight_bias(qweight, mod.bias) qlinear_tanh.scale = float(act_scale) qlinear_tanh.zero_point = int(act_zp) return qlinear_tanh @classmethod def from_reference(cls, ref_mod, output_scale, output_zero_point): linear = ref_mod[0] qlinear_tanh = cls( linear.in_features, linear.out_features) qweight = linear.get_quantized_weight() qlinear_tanh.set_weight_bias(qweight, linear.bias) qlinear_tanh.scale = float(output_scale) qlinear_tanh.zero_point = int(output_zero_point) return qlinear_tanh

Docs

Access comprehensive developer documentation for PyTorch

View Docs

Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

Resources

Find development resources and get your questions answered

View Resources