# Source code for tensorplay.ao.nn.intrinsic Source: https://www.tensorplay.cn/docs/_modules/tensorplay/ao/nn/intrinsic.html ``` """Fused float building blocks. These modules bundle an operator with the activation that immediately follows it. They are the fusion targets of the eager workflow and the float counterparts used during quantization-aware training. """ from __future__ import annotations from tensorplay import nn __all__ = [ "LinearReLU", "ConvReLU1d", "ConvReLU2d", "ConvReLU3d", "BNReLU2d", "BNReLU3d", ] class _FusedActivation(nn.Module): """Base for operator + relu bundles; subclasses set the operator class.""" _OP_CLASS = None def __init__(self, *args, **kwargs): super().__init__() self.op = self._OP_CLASS(*args, **kwargs) def forward(self, input): return self.op(input).relu() @classmethod def from_float(cls, mod): """Build the fused module from an unfused float operator module.""" fused = cls( mod.in_channels if hasattr(mod, "in_channels") else mod.in_features, mod.out_channels if hasattr(mod, "out_channels") else mod.out_features, *( mod.kernel_size if hasattr(mod, "kernel_size") else () ), **( { "stride": mod.stride, "padding": mod.padding, "dilation": mod.dilation, "groups": mod.groups, "bias": mod.bias is not None, } if hasattr(mod, "stride") else {"bias": mod.bias is not None} ), ) fused.op.weight = mod.weight if mod.bias is not None: fused.op.bias = mod.bias fused.train(mod.training) return fused [docs] class LinearReLU(_FusedActivation): """A linear transformation followed by a rectifier.""" _OP_CLASS = nn.Linear [docs] class ConvReLU1d(_FusedActivation): """A 1d convolution followed by a rectifier.""" _OP_CLASS = nn.Conv1d [docs] class ConvReLU2d(_FusedActivation): """A 2d convolution followed by a rectifier.""" _OP_CLASS = nn.Conv2d [docs] class ConvReLU3d(_FusedActivation): """A 3d convolution followed by a rectifier.""" _OP_CLASS = nn.Conv3d [docs] class BNReLU2d(nn.Module): """A 2d batch norm followed by a rectifier.""" def __init__(self, num_features, eps=1e-5, momentum=0.1): super().__init__() self.bn = nn.BatchNorm2d(num_features, eps, momentum) self.relu = nn.ReLU() def forward(self, input): return self.relu(self.bn(input)) [docs] class BNReLU3d(nn.Module): """A 3d batch norm followed by a rectifier.""" def __init__(self, num_features, eps=1e-5, momentum=0.1): super().__init__() self.bn = nn.BatchNorm3d(num_features, eps, momentum) self.relu = nn.ReLU() def forward(self, input): return self.relu(self.bn(input)) ```