Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,41 @@
|
||||
import types
|
||||
|
||||
from .modules import * # noqa: F403
|
||||
from .modules.fused import _FusedModule # noqa: F403
|
||||
|
||||
|
||||
# # Subpackages
|
||||
# from . import qat # noqa: F403
|
||||
# from . import quantized # noqa: F403
|
||||
|
||||
__all__ = [
|
||||
"ConvBn1d",
|
||||
"ConvBn2d",
|
||||
"ConvBn3d",
|
||||
"ConvBnReLU1d",
|
||||
"ConvBnReLU2d",
|
||||
"ConvBnReLU3d",
|
||||
"ConvReLU1d",
|
||||
"ConvReLU2d",
|
||||
"ConvReLU3d",
|
||||
"LinearReLU",
|
||||
"BNReLU2d",
|
||||
"BNReLU3d",
|
||||
"LinearBn1d",
|
||||
"LinearLeakyReLU",
|
||||
"LinearTanh",
|
||||
"ConvAdd2d",
|
||||
"ConvAddReLU2d",
|
||||
]
|
||||
|
||||
|
||||
# We are exposing all subpackages to the end-user.
|
||||
# Because of possible inter-dependency, we want to avoid
|
||||
# the cyclic imports, thus implementing lazy version
|
||||
# as per https://peps.python.org/pep-0562/
|
||||
def __getattr__(name: str) -> types.ModuleType:
|
||||
if name in __all__:
|
||||
import importlib
|
||||
|
||||
return importlib.import_module("." + name, __name__)
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -0,0 +1,41 @@
|
||||
from .fused import ( # noqa: F401
|
||||
_FusedModule,
|
||||
BNReLU2d,
|
||||
BNReLU3d,
|
||||
ConvAdd2d,
|
||||
ConvAddReLU2d,
|
||||
ConvBn1d,
|
||||
ConvBn2d,
|
||||
ConvBn3d,
|
||||
ConvBnReLU1d,
|
||||
ConvBnReLU2d,
|
||||
ConvBnReLU3d,
|
||||
ConvReLU1d,
|
||||
ConvReLU2d,
|
||||
ConvReLU3d,
|
||||
LinearBn1d,
|
||||
LinearLeakyReLU,
|
||||
LinearReLU,
|
||||
LinearTanh,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ConvBn1d",
|
||||
"ConvBn2d",
|
||||
"ConvBn3d",
|
||||
"ConvBnReLU1d",
|
||||
"ConvBnReLU2d",
|
||||
"ConvBnReLU3d",
|
||||
"ConvReLU1d",
|
||||
"ConvReLU2d",
|
||||
"ConvReLU3d",
|
||||
"LinearReLU",
|
||||
"BNReLU2d",
|
||||
"BNReLU3d",
|
||||
"LinearBn1d",
|
||||
"LinearLeakyReLU",
|
||||
"LinearTanh",
|
||||
"ConvAdd2d",
|
||||
"ConvAddReLU2d",
|
||||
]
|
||||
@@ -0,0 +1,319 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import torch
|
||||
from torch.nn import (
|
||||
BatchNorm1d,
|
||||
BatchNorm2d,
|
||||
BatchNorm3d,
|
||||
Conv1d,
|
||||
Conv2d,
|
||||
Conv3d,
|
||||
Linear,
|
||||
ReLU,
|
||||
)
|
||||
from torch.nn.utils.parametrize import type_before_parametrizations
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ConvReLU1d",
|
||||
"ConvReLU2d",
|
||||
"ConvReLU3d",
|
||||
"LinearReLU",
|
||||
"ConvBn1d",
|
||||
"ConvBn2d",
|
||||
"ConvBnReLU1d",
|
||||
"ConvBnReLU2d",
|
||||
"ConvBn3d",
|
||||
"ConvBnReLU3d",
|
||||
"BNReLU2d",
|
||||
"BNReLU3d",
|
||||
"LinearBn1d",
|
||||
"LinearLeakyReLU",
|
||||
"LinearTanh",
|
||||
"ConvAdd2d",
|
||||
"ConvAddReLU2d",
|
||||
]
|
||||
|
||||
|
||||
# Used for identifying intrinsic modules used in quantization
|
||||
class _FusedModule(torch.nn.Sequential):
|
||||
pass
|
||||
|
||||
|
||||
class ConvReLU1d(_FusedModule):
|
||||
r"""This is a sequential container which calls the Conv1d and ReLU modules.
|
||||
During quantization this will be replaced with the corresponding fused module."""
|
||||
|
||||
def __init__(self, conv, relu):
|
||||
if not (
|
||||
type_before_parametrizations(conv) == Conv1d
|
||||
and type_before_parametrizations(relu) == ReLU
|
||||
):
|
||||
raise AssertionError(
|
||||
f"Incorrect types for input modules: "
|
||||
f"{type_before_parametrizations(conv).__name__} and "
|
||||
f"{type_before_parametrizations(relu).__name__}"
|
||||
)
|
||||
super().__init__(conv, relu)
|
||||
|
||||
|
||||
class ConvReLU2d(_FusedModule):
|
||||
r"""This is a sequential container which calls the Conv2d and ReLU modules.
|
||||
During quantization this will be replaced with the corresponding fused module."""
|
||||
|
||||
def __init__(self, conv, relu):
|
||||
if not (
|
||||
type_before_parametrizations(conv) == Conv2d
|
||||
and type_before_parametrizations(relu) == ReLU
|
||||
):
|
||||
raise AssertionError(
|
||||
f"Incorrect types for input modules: "
|
||||
f"{type_before_parametrizations(conv).__name__} and "
|
||||
f"{type_before_parametrizations(relu).__name__}"
|
||||
)
|
||||
super().__init__(conv, relu)
|
||||
|
||||
|
||||
class ConvReLU3d(_FusedModule):
|
||||
r"""This is a sequential container which calls the Conv3d and ReLU modules.
|
||||
During quantization this will be replaced with the corresponding fused module."""
|
||||
|
||||
def __init__(self, conv, relu):
|
||||
if not (
|
||||
type_before_parametrizations(conv) == Conv3d
|
||||
and type_before_parametrizations(relu) == ReLU
|
||||
):
|
||||
raise AssertionError(
|
||||
f"Incorrect types for input modules: "
|
||||
f"{type_before_parametrizations(conv).__name__} and "
|
||||
f"{type_before_parametrizations(relu).__name__}"
|
||||
)
|
||||
super().__init__(conv, relu)
|
||||
|
||||
|
||||
class LinearReLU(_FusedModule):
|
||||
r"""This is a sequential container which calls the Linear and ReLU modules.
|
||||
During quantization this will be replaced with the corresponding fused module."""
|
||||
|
||||
def __init__(self, linear, relu):
|
||||
if not (
|
||||
type_before_parametrizations(linear) == Linear
|
||||
and type_before_parametrizations(relu) == ReLU
|
||||
):
|
||||
raise AssertionError(
|
||||
f"Incorrect types for input modules: "
|
||||
f"{type_before_parametrizations(linear).__name__} and "
|
||||
f"{type_before_parametrizations(relu).__name__}"
|
||||
)
|
||||
super().__init__(linear, relu)
|
||||
|
||||
|
||||
class ConvBn1d(_FusedModule):
|
||||
r"""This is a sequential container which calls the Conv 1d and Batch Norm 1d modules.
|
||||
During quantization this will be replaced with the corresponding fused module."""
|
||||
|
||||
def __init__(self, conv, bn):
|
||||
if not (
|
||||
type_before_parametrizations(conv) == Conv1d
|
||||
and type_before_parametrizations(bn) == BatchNorm1d
|
||||
):
|
||||
raise AssertionError(
|
||||
f"Incorrect types for input modules: "
|
||||
f"{type_before_parametrizations(conv).__name__} and "
|
||||
f"{type_before_parametrizations(bn).__name__}"
|
||||
)
|
||||
super().__init__(conv, bn)
|
||||
|
||||
|
||||
class ConvBn2d(_FusedModule):
|
||||
r"""This is a sequential container which calls the Conv 2d and Batch Norm 2d modules.
|
||||
During quantization this will be replaced with the corresponding fused module."""
|
||||
|
||||
def __init__(self, conv, bn):
|
||||
if not (
|
||||
type_before_parametrizations(conv) == Conv2d
|
||||
and type_before_parametrizations(bn) == BatchNorm2d
|
||||
):
|
||||
raise AssertionError(
|
||||
f"Incorrect types for input modules: "
|
||||
f"{type_before_parametrizations(conv).__name__} and "
|
||||
f"{type_before_parametrizations(bn).__name__}"
|
||||
)
|
||||
super().__init__(conv, bn)
|
||||
|
||||
|
||||
class ConvBnReLU1d(_FusedModule):
|
||||
r"""This is a sequential container which calls the Conv 1d, Batch Norm 1d, and ReLU modules.
|
||||
During quantization this will be replaced with the corresponding fused module."""
|
||||
|
||||
def __init__(self, conv, bn, relu):
|
||||
if not (
|
||||
type_before_parametrizations(conv) == Conv1d
|
||||
and type_before_parametrizations(bn) == BatchNorm1d
|
||||
and type_before_parametrizations(relu) == ReLU
|
||||
):
|
||||
raise AssertionError(
|
||||
f"Incorrect types for input modules: "
|
||||
f"{type_before_parametrizations(conv).__name__}, "
|
||||
f"{type_before_parametrizations(bn).__name__}, and "
|
||||
f"{type_before_parametrizations(relu).__name__}"
|
||||
)
|
||||
super().__init__(conv, bn, relu)
|
||||
|
||||
|
||||
class ConvBnReLU2d(_FusedModule):
|
||||
r"""This is a sequential container which calls the Conv 2d, Batch Norm 2d, and ReLU modules.
|
||||
During quantization this will be replaced with the corresponding fused module."""
|
||||
|
||||
def __init__(self, conv, bn, relu):
|
||||
if not (
|
||||
type_before_parametrizations(conv) == Conv2d
|
||||
and type_before_parametrizations(bn) == BatchNorm2d
|
||||
and type_before_parametrizations(relu) == ReLU
|
||||
):
|
||||
raise AssertionError(
|
||||
f"Incorrect types for input modules: "
|
||||
f"{type_before_parametrizations(conv).__name__}, "
|
||||
f"{type_before_parametrizations(bn).__name__}, and "
|
||||
f"{type_before_parametrizations(relu).__name__}"
|
||||
)
|
||||
super().__init__(conv, bn, relu)
|
||||
|
||||
|
||||
class ConvBn3d(_FusedModule):
|
||||
r"""This is a sequential container which calls the Conv 3d and Batch Norm 3d modules.
|
||||
During quantization this will be replaced with the corresponding fused module."""
|
||||
|
||||
def __init__(self, conv, bn):
|
||||
if not (
|
||||
type_before_parametrizations(conv) == Conv3d
|
||||
and type_before_parametrizations(bn) == BatchNorm3d
|
||||
):
|
||||
raise AssertionError(
|
||||
f"Incorrect types for input modules: "
|
||||
f"{type_before_parametrizations(conv).__name__} and "
|
||||
f"{type_before_parametrizations(bn).__name__}"
|
||||
)
|
||||
super().__init__(conv, bn)
|
||||
|
||||
|
||||
class ConvBnReLU3d(_FusedModule):
|
||||
r"""This is a sequential container which calls the Conv 3d, Batch Norm 3d, and ReLU modules.
|
||||
During quantization this will be replaced with the corresponding fused module."""
|
||||
|
||||
def __init__(self, conv, bn, relu):
|
||||
if not (
|
||||
type_before_parametrizations(conv) == Conv3d
|
||||
and type_before_parametrizations(bn) == BatchNorm3d
|
||||
and type_before_parametrizations(relu) == ReLU
|
||||
):
|
||||
raise AssertionError(
|
||||
f"Incorrect types for input modules: "
|
||||
f"{type_before_parametrizations(conv).__name__}, "
|
||||
f"{type_before_parametrizations(bn).__name__}, and "
|
||||
f"{type_before_parametrizations(relu).__name__}"
|
||||
)
|
||||
super().__init__(conv, bn, relu)
|
||||
|
||||
|
||||
class BNReLU2d(_FusedModule):
|
||||
r"""This is a sequential container which calls the BatchNorm 2d and ReLU modules.
|
||||
During quantization this will be replaced with the corresponding fused module."""
|
||||
|
||||
def __init__(self, batch_norm, relu):
|
||||
if not (
|
||||
type_before_parametrizations(batch_norm) == BatchNorm2d
|
||||
and type_before_parametrizations(relu) == ReLU
|
||||
):
|
||||
raise AssertionError(
|
||||
f"Incorrect types for input modules: "
|
||||
f"{type_before_parametrizations(batch_norm).__name__} and "
|
||||
f"{type_before_parametrizations(relu).__name__}"
|
||||
)
|
||||
super().__init__(batch_norm, relu)
|
||||
|
||||
|
||||
class BNReLU3d(_FusedModule):
|
||||
r"""This is a sequential container which calls the BatchNorm 3d and ReLU modules.
|
||||
During quantization this will be replaced with the corresponding fused module."""
|
||||
|
||||
def __init__(self, batch_norm, relu):
|
||||
if not (
|
||||
type_before_parametrizations(batch_norm) == BatchNorm3d
|
||||
and type_before_parametrizations(relu) == ReLU
|
||||
):
|
||||
raise AssertionError(
|
||||
f"Incorrect types for input modules: "
|
||||
f"{type_before_parametrizations(batch_norm).__name__} and "
|
||||
f"{type_before_parametrizations(relu).__name__}"
|
||||
)
|
||||
super().__init__(batch_norm, relu)
|
||||
|
||||
|
||||
class LinearBn1d(_FusedModule):
|
||||
r"""This is a sequential container which calls the Linear and BatchNorm1d modules.
|
||||
During quantization this will be replaced with the corresponding fused module."""
|
||||
|
||||
def __init__(self, linear, bn):
|
||||
if not (
|
||||
type_before_parametrizations(linear) == Linear
|
||||
and type_before_parametrizations(bn) == BatchNorm1d
|
||||
):
|
||||
raise AssertionError(
|
||||
f"Incorrect types for input modules: "
|
||||
f"{type_before_parametrizations(linear).__name__} and "
|
||||
f"{type_before_parametrizations(bn).__name__}"
|
||||
)
|
||||
super().__init__(linear, bn)
|
||||
|
||||
|
||||
class LinearLeakyReLU(_FusedModule):
|
||||
r"""This is a sequential container which calls the Linear and LeakyReLU modules.
|
||||
During quantization this will be replaced with the corresponding fused module."""
|
||||
|
||||
def __init__(self, linear, leaky_relu):
|
||||
if not (type(linear) is Linear and type(leaky_relu) is torch.nn.LeakyReLU):
|
||||
raise AssertionError(
|
||||
f"Incorrect types for input modules: "
|
||||
f"{type(linear).__name__} and {type(leaky_relu).__name__}"
|
||||
)
|
||||
super().__init__(linear, leaky_relu)
|
||||
|
||||
|
||||
class LinearTanh(_FusedModule):
|
||||
r"""This is a sequential container which calls the Linear and Tanh modules.
|
||||
During quantization this will be replaced with the corresponding fused module."""
|
||||
|
||||
def __init__(self, linear, tanh):
|
||||
if not (type(linear) is Linear and type(tanh) is torch.nn.Tanh):
|
||||
raise AssertionError(
|
||||
f"Incorrect types for input modules: "
|
||||
f"{type(linear).__name__} and {type(tanh).__name__}"
|
||||
)
|
||||
super().__init__(linear, tanh)
|
||||
|
||||
|
||||
class ConvAdd2d(_FusedModule):
|
||||
r"""This is a sequential container which calls the Conv2d modules with extra Add.
|
||||
During quantization this will be replaced with the corresponding fused module."""
|
||||
|
||||
def __init__(self, conv, add):
|
||||
super().__init__(conv)
|
||||
self.add = add
|
||||
|
||||
def forward(self, x1, x2): # type: ignore[override]
|
||||
r"""Applies convolution to x1 and adds the result to x2."""
|
||||
return self.add(self[0](x1), x2)
|
||||
|
||||
|
||||
class ConvAddReLU2d(_FusedModule):
|
||||
r"""This is a sequential container which calls the Conv2d, add, Relu.
|
||||
During quantization this will be replaced with the corresponding fused module."""
|
||||
|
||||
def __init__(self, conv, add, relu):
|
||||
super().__init__(conv)
|
||||
self.add = add
|
||||
self.relu = relu
|
||||
|
||||
def forward(self, x1, x2): # type: ignore[override]
|
||||
r"""Applies convolution to x1, adds the result to x2, and applies ReLU."""
|
||||
return self.relu(self.add(self[0](x1), x2))
|
||||
@@ -0,0 +1 @@
|
||||
from .modules import * # noqa: F403
|
||||
+32
@@ -0,0 +1,32 @@
|
||||
from .conv_fused import (
|
||||
ConvBn1d,
|
||||
ConvBn2d,
|
||||
ConvBn3d,
|
||||
ConvBnReLU1d,
|
||||
ConvBnReLU2d,
|
||||
ConvBnReLU3d,
|
||||
ConvReLU1d,
|
||||
ConvReLU2d,
|
||||
ConvReLU3d,
|
||||
freeze_bn_stats,
|
||||
update_bn_stats,
|
||||
)
|
||||
from .linear_fused import LinearBn1d
|
||||
from .linear_relu import LinearReLU
|
||||
|
||||
|
||||
__all__ = [
|
||||
"LinearReLU",
|
||||
"LinearBn1d",
|
||||
"ConvReLU1d",
|
||||
"ConvReLU2d",
|
||||
"ConvReLU3d",
|
||||
"ConvBn1d",
|
||||
"ConvBn2d",
|
||||
"ConvBn3d",
|
||||
"ConvBnReLU1d",
|
||||
"ConvBnReLU2d",
|
||||
"ConvBnReLU3d",
|
||||
"update_bn_stats",
|
||||
"freeze_bn_stats",
|
||||
]
|
||||
+971
@@ -0,0 +1,971 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import math
|
||||
from typing import ClassVar
|
||||
|
||||
import torch
|
||||
import torch.ao.nn.intrinsic as nni
|
||||
import torch.ao.nn.qat as nnqat
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.nn import init
|
||||
from torch.nn.modules.utils import _pair, _single, _triple
|
||||
from torch.nn.parameter import Parameter
|
||||
from torch.nn.utils import fuse_conv_bn_weights
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ConvBn1d",
|
||||
"ConvBnReLU1d",
|
||||
"ConvReLU1d",
|
||||
"ConvBn2d",
|
||||
"ConvBnReLU2d",
|
||||
"ConvReLU2d",
|
||||
"ConvBn3d",
|
||||
"ConvBnReLU3d",
|
||||
"ConvReLU3d",
|
||||
"update_bn_stats",
|
||||
"freeze_bn_stats",
|
||||
]
|
||||
_BN_CLASS_MAP = {
|
||||
1: nn.BatchNorm1d,
|
||||
2: nn.BatchNorm2d,
|
||||
3: nn.BatchNorm3d,
|
||||
}
|
||||
|
||||
|
||||
class _ConvBnNd(nn.modules.conv._ConvNd, nni._FusedModule):
|
||||
_version = 2
|
||||
_FLOAT_MODULE: ClassVar[type[nn.modules.conv._ConvNd]]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
# ConvNd args
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
dilation,
|
||||
transposed,
|
||||
output_padding,
|
||||
groups,
|
||||
bias,
|
||||
padding_mode,
|
||||
# BatchNormNd args
|
||||
# num_features: out_channels
|
||||
eps=1e-05,
|
||||
momentum=0.1,
|
||||
# affine: True
|
||||
# track_running_stats: True
|
||||
# Args for this module
|
||||
freeze_bn=False,
|
||||
qconfig=None,
|
||||
dim=2,
|
||||
):
|
||||
nn.modules.conv._ConvNd.__init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
dilation,
|
||||
transposed,
|
||||
output_padding,
|
||||
groups,
|
||||
False,
|
||||
padding_mode,
|
||||
)
|
||||
if not qconfig:
|
||||
raise AssertionError("qconfig must be provided for QAT module")
|
||||
self.qconfig = qconfig
|
||||
self.freeze_bn = freeze_bn if self.training else True
|
||||
self.bn = _BN_CLASS_MAP[dim](out_channels, eps, momentum, True, True)
|
||||
self.weight_fake_quant = self.qconfig.weight()
|
||||
if bias:
|
||||
self.bias = Parameter(torch.empty(out_channels))
|
||||
else:
|
||||
self.register_parameter("bias", None)
|
||||
self.reset_bn_parameters()
|
||||
|
||||
# this needs to be called after reset_bn_parameters,
|
||||
# as they modify the same state
|
||||
if self.training:
|
||||
if freeze_bn:
|
||||
self.freeze_bn_stats()
|
||||
else:
|
||||
self.update_bn_stats()
|
||||
else:
|
||||
self.freeze_bn_stats()
|
||||
|
||||
self._enable_slow_path_for_better_numerical_stability = False
|
||||
|
||||
def reset_running_stats(self):
|
||||
self.bn.reset_running_stats()
|
||||
|
||||
def reset_bn_parameters(self):
|
||||
self.bn.reset_running_stats()
|
||||
init.uniform_(self.bn.weight)
|
||||
init.zeros_(self.bn.bias)
|
||||
# note: below is actually for conv, not BN
|
||||
if self.bias is not None:
|
||||
fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
|
||||
bound = 1 / math.sqrt(fan_in)
|
||||
init.uniform_(self.bias, -bound, bound)
|
||||
|
||||
def update_bn_stats(self):
|
||||
self.freeze_bn = False
|
||||
self.bn.training = True
|
||||
return self
|
||||
|
||||
def freeze_bn_stats(self):
|
||||
self.freeze_bn = True
|
||||
self.bn.training = False
|
||||
return self
|
||||
|
||||
def _forward(self, input):
|
||||
if self._enable_slow_path_for_better_numerical_stability:
|
||||
return self._forward_slow(input)
|
||||
return self._forward_approximate(input)
|
||||
|
||||
def _forward_approximate(self, input):
|
||||
"""Approximated method to fuse conv and bn. It requires only one forward pass.
|
||||
conv_orig = conv / scale_factor where scale_factor = bn.weight / running_std
|
||||
"""
|
||||
if self.bn.running_var is None:
|
||||
raise AssertionError("self.bn.running_var must not be None")
|
||||
running_std = torch.sqrt(self.bn.running_var + self.bn.eps)
|
||||
scale_factor = self.bn.weight / running_std
|
||||
weight_shape = [1] * len(self.weight.shape)
|
||||
weight_shape[0] = -1
|
||||
bias_shape = [1] * len(self.weight.shape)
|
||||
bias_shape[1] = -1
|
||||
scaled_weight = self.weight_fake_quant(
|
||||
self.weight * scale_factor.reshape(weight_shape)
|
||||
)
|
||||
# using zero bias here since the bias for original conv
|
||||
# will be added later
|
||||
if self.bias is not None:
|
||||
zero_bias = torch.zeros_like(self.bias, dtype=input.dtype)
|
||||
else:
|
||||
zero_bias = torch.zeros(
|
||||
self.out_channels, device=scaled_weight.device, dtype=input.dtype
|
||||
)
|
||||
conv = self._conv_forward(input, scaled_weight, zero_bias)
|
||||
conv_orig = conv / scale_factor.reshape(bias_shape)
|
||||
if self.bias is not None:
|
||||
conv_orig = conv_orig + self.bias.reshape(bias_shape)
|
||||
conv = self.bn(conv_orig)
|
||||
return conv
|
||||
|
||||
def _forward_slow(self, input):
|
||||
"""
|
||||
A more accurate but slow method to compute conv bn fusion, following https://arxiv.org/pdf/1806.08342.pdf
|
||||
It requires two forward passes but handles the case bn.weight == 0
|
||||
|
||||
Conv: Y = WX + B_c
|
||||
Conv without bias: Y0 = WX = Y - B_c, Y = Y0 + B_c
|
||||
|
||||
Batch statistics:
|
||||
mean_Y = Y.mean()
|
||||
= Y0.mean() + B_c
|
||||
var_Y = (Y - mean_Y)^2.mean()
|
||||
= (Y0 - Y0.mean())^2.mean()
|
||||
BN (r: bn.weight, beta: bn.bias):
|
||||
Z = r * (Y - mean_Y) / sqrt(var_Y + eps) + beta
|
||||
= r * (Y0 - Y0.mean()) / sqrt(var_Y + eps) + beta
|
||||
|
||||
Fused Conv BN training (std_Y = sqrt(var_Y + eps)):
|
||||
Z = (r * W / std_Y) * X + r * (B_c - mean_Y) / std_Y + beta
|
||||
= (r * W / std_Y) * X - r * Y0.mean() / std_Y + beta
|
||||
|
||||
Fused Conv BN inference (running_std = sqrt(running_var + eps)):
|
||||
Z = (r * W / running_std) * X - r * (running_mean - B_c) / running_std + beta
|
||||
|
||||
QAT with fused conv bn:
|
||||
Z_train = fake_quant(r * W / running_std) * X * (running_std / std_Y) - r * Y0.mean() / std_Y + beta
|
||||
= conv(X, fake_quant(r * W / running_std)) * (running_std / std_Y) - r * Y0.mean() / std_Y + beta
|
||||
Z_inference = conv(X, fake_quant(r * W / running_std)) - r * (running_mean - B_c) / running_std + beta
|
||||
"""
|
||||
|
||||
if self.bn.running_var is None:
|
||||
raise AssertionError("self.bn.running_var must not be None")
|
||||
if self.bn.running_mean is None:
|
||||
raise AssertionError("self.bn.running_mean must not be None")
|
||||
|
||||
# using zero bias here since the bias for original conv
|
||||
# will be added later
|
||||
zero_bias = torch.zeros(
|
||||
self.out_channels, device=self.weight.device, dtype=input.dtype
|
||||
)
|
||||
|
||||
weight_shape = [1] * len(self.weight.shape)
|
||||
weight_shape[0] = -1
|
||||
bias_shape = [1] * len(self.weight.shape)
|
||||
bias_shape[1] = -1
|
||||
|
||||
if self.bn.training:
|
||||
# needed to compute batch mean/std
|
||||
conv_out = self._conv_forward(input, self.weight, zero_bias)
|
||||
# update bn statistics
|
||||
with torch.no_grad():
|
||||
conv_out_bias = (
|
||||
conv_out
|
||||
if self.bias is None
|
||||
else conv_out + self.bias.reshape(bias_shape)
|
||||
)
|
||||
self.bn(conv_out_bias)
|
||||
|
||||
# fused conv + bn without bias using bn running statistics
|
||||
running_std = torch.sqrt(self.bn.running_var + self.bn.eps)
|
||||
scale_factor = self.bn.weight / running_std
|
||||
scaled_weight = self.weight_fake_quant(
|
||||
self.weight * scale_factor.reshape(weight_shape)
|
||||
)
|
||||
# fused conv without bias for inference: (r * W / running_std) * X
|
||||
conv_bn = self._conv_forward(input, scaled_weight, zero_bias)
|
||||
|
||||
avg_dims = [0] + list(range(2, len(self.weight.shape)))
|
||||
batch_mean = conv_out.mean(avg_dims)
|
||||
batch_var = torch.square(conv_out - batch_mean.reshape(bias_shape)).mean(
|
||||
avg_dims
|
||||
)
|
||||
batch_std = torch.sqrt(batch_var + self.bn.eps)
|
||||
|
||||
# scale to use batch std in training mode
|
||||
# conv(X, r * W / std_Y) = conv(X, r * W / running_std) * (running_std / std_Y)
|
||||
unscale_factor = running_std / batch_std
|
||||
conv_bn *= unscale_factor.reshape(bias_shape)
|
||||
|
||||
fused_mean = batch_mean
|
||||
fused_std = batch_std
|
||||
else:
|
||||
# fused conv + bn without bias using bn running statistics
|
||||
running_std = torch.sqrt(self.bn.running_var + self.bn.eps)
|
||||
scale_factor = self.bn.weight / running_std
|
||||
scaled_weight = self.weight_fake_quant(
|
||||
self.weight * scale_factor.reshape(weight_shape)
|
||||
)
|
||||
# fused conv without bias for inference: (r * W / running_std) * X
|
||||
conv_bn = self._conv_forward(input, scaled_weight, zero_bias)
|
||||
|
||||
fused_mean = self.bn.running_mean - (
|
||||
self.bias if self.bias is not None else 0
|
||||
)
|
||||
fused_std = running_std
|
||||
|
||||
# fused bias = beta - r * mean / std
|
||||
fused_bias = self.bn.bias - self.bn.weight * fused_mean / fused_std
|
||||
conv_bn += fused_bias.reshape(bias_shape)
|
||||
|
||||
# HACK to let conv bias participate in loss to avoid DDP error (parameters
|
||||
# were not used in producing loss)
|
||||
if self.bias is not None:
|
||||
conv_bn += (self.bias - self.bias).reshape(bias_shape)
|
||||
|
||||
return conv_bn
|
||||
|
||||
def forward(self, input):
|
||||
return self._forward(input)
|
||||
|
||||
def train(self, mode=True):
|
||||
"""
|
||||
Batchnorm's training behavior is using the self.training flag. Prevent
|
||||
changing it if BN is frozen. This makes sure that calling `model.train()`
|
||||
on a model with a frozen BN will behave properly.
|
||||
"""
|
||||
self.training = mode
|
||||
if not self.freeze_bn:
|
||||
for module in self.children():
|
||||
module.train(mode)
|
||||
return self
|
||||
|
||||
# ===== Serialization version history =====
|
||||
#
|
||||
# Version 1/None
|
||||
# self
|
||||
# |--- weight : Tensor
|
||||
# |--- bias : Tensor
|
||||
# |--- gamma : Tensor
|
||||
# |--- beta : Tensor
|
||||
# |--- running_mean : Tensor
|
||||
# |--- running_var : Tensor
|
||||
# |--- num_batches_tracked : Tensor
|
||||
#
|
||||
# Version 2
|
||||
# self
|
||||
# |--- weight : Tensor
|
||||
# |--- bias : Tensor
|
||||
# |--- bn : Module
|
||||
# |--- weight : Tensor (moved from v1.self.gamma)
|
||||
# |--- bias : Tensor (moved from v1.self.beta)
|
||||
# |--- running_mean : Tensor (moved from v1.self.running_mean)
|
||||
# |--- running_var : Tensor (moved from v1.self.running_var)
|
||||
# |--- num_batches_tracked : Tensor (moved from v1.self.num_batches_tracked)
|
||||
def _load_from_state_dict(
|
||||
self,
|
||||
state_dict,
|
||||
prefix,
|
||||
local_metadata,
|
||||
strict,
|
||||
missing_keys,
|
||||
unexpected_keys,
|
||||
error_msgs,
|
||||
):
|
||||
version = local_metadata.get("version", None)
|
||||
if version is None or version == 1:
|
||||
# BN related parameters and buffers were moved into the BN module for v2
|
||||
v2_to_v1_names = {
|
||||
"bn.weight": "gamma",
|
||||
"bn.bias": "beta",
|
||||
"bn.running_mean": "running_mean",
|
||||
"bn.running_var": "running_var",
|
||||
"bn.num_batches_tracked": "num_batches_tracked",
|
||||
}
|
||||
for v2_name, v1_name in v2_to_v1_names.items():
|
||||
if prefix + v1_name in state_dict:
|
||||
state_dict[prefix + v2_name] = state_dict[prefix + v1_name]
|
||||
state_dict.pop(prefix + v1_name)
|
||||
elif prefix + v2_name in state_dict:
|
||||
# there was a brief period where forward compatibility
|
||||
# for this module was broken (between
|
||||
# https://github.com/pytorch/pytorch/pull/38478
|
||||
# and https://github.com/pytorch/pytorch/pull/38820)
|
||||
# and modules emitted the v2 state_dict format while
|
||||
# specifying that version == 1. This patches the forward
|
||||
# compatibility issue by allowing the v2 style entries to
|
||||
# be used.
|
||||
pass
|
||||
elif strict:
|
||||
missing_keys.append(prefix + v2_name)
|
||||
|
||||
super()._load_from_state_dict(
|
||||
state_dict,
|
||||
prefix,
|
||||
local_metadata,
|
||||
strict,
|
||||
missing_keys,
|
||||
unexpected_keys,
|
||||
error_msgs,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False):
|
||||
r"""Create a qat module from a float module or qparams_dict
|
||||
|
||||
Args: `mod` a float module, either produced by torch.ao.quantization utilities
|
||||
or directly from user
|
||||
"""
|
||||
# The ignore is because _FLOAT_MODULE is a TypeVar here where the bound
|
||||
# has no __name__ (code is fine though)
|
||||
if type(mod) is not cls._FLOAT_MODULE:
|
||||
raise AssertionError(
|
||||
"qat."
|
||||
+ cls.__name__
|
||||
+ ".from_float only works for "
|
||||
+ cls._FLOAT_MODULE.__name__
|
||||
)
|
||||
if not hasattr(mod, "qconfig"):
|
||||
raise AssertionError("Input float module must have qconfig defined")
|
||||
if not mod.qconfig:
|
||||
raise AssertionError("Input float module must have a valid qconfig")
|
||||
qconfig = mod.qconfig
|
||||
conv, bn = mod[0], mod[1] # type: ignore[index]
|
||||
qat_convbn = cls(
|
||||
conv.in_channels,
|
||||
conv.out_channels,
|
||||
conv.kernel_size,
|
||||
conv.stride,
|
||||
conv.padding,
|
||||
conv.dilation,
|
||||
conv.groups,
|
||||
conv.bias is not None,
|
||||
conv.padding_mode,
|
||||
bn.eps,
|
||||
bn.momentum,
|
||||
False,
|
||||
qconfig,
|
||||
)
|
||||
qat_convbn.weight = conv.weight
|
||||
qat_convbn.bias = conv.bias
|
||||
qat_convbn.bn.weight = bn.weight
|
||||
qat_convbn.bn.bias = bn.bias
|
||||
qat_convbn.bn.running_mean = bn.running_mean
|
||||
qat_convbn.bn.running_var = bn.running_var
|
||||
# mypy error: Cannot determine type of 'num_batches_tracked'
|
||||
qat_convbn.bn.num_batches_tracked = bn.num_batches_tracked
|
||||
return qat_convbn
|
||||
|
||||
def to_float(self):
|
||||
cls = type(self)
|
||||
conv = cls._FLOAT_CONV_MODULE( # type: ignore[attr-defined]
|
||||
self.in_channels,
|
||||
self.out_channels,
|
||||
self.kernel_size,
|
||||
self.stride,
|
||||
self.padding,
|
||||
self.dilation,
|
||||
self.groups,
|
||||
self.bias is not None,
|
||||
self.padding_mode,
|
||||
)
|
||||
conv.weight = torch.nn.Parameter(self.weight.detach())
|
||||
if self.bias is not None:
|
||||
conv.bias = torch.nn.Parameter(self.bias.detach())
|
||||
|
||||
if cls._FLOAT_BN_MODULE: # type: ignore[attr-defined]
|
||||
# fuse bn into conv
|
||||
if self.bn.running_var is None or self.bn.running_mean is None:
|
||||
raise AssertionError(
|
||||
"self.bn.running_var and self.bn.running_mean must not be None"
|
||||
)
|
||||
conv.weight, conv.bias = fuse_conv_bn_weights(
|
||||
conv.weight,
|
||||
conv.bias,
|
||||
self.bn.running_mean,
|
||||
self.bn.running_var,
|
||||
self.bn.eps,
|
||||
self.bn.weight,
|
||||
self.bn.bias,
|
||||
)
|
||||
|
||||
if cls._FLOAT_RELU_MODULE: # type: ignore[attr-defined]
|
||||
modules = []
|
||||
modules.append(conv)
|
||||
relu = cls._FLOAT_RELU_MODULE() # type: ignore[attr-defined]
|
||||
modules.append(relu)
|
||||
conv_relu = cls._FUSED_FLOAT_MODULE(*modules) # type: ignore[attr-defined]
|
||||
conv_relu.train(self.training)
|
||||
return conv_relu
|
||||
else:
|
||||
conv.train(self.training)
|
||||
return conv
|
||||
|
||||
|
||||
class ConvBn1d(_ConvBnNd, nn.Conv1d):
|
||||
r"""
|
||||
A ConvBn1d module is a module fused from Conv1d and BatchNorm1d,
|
||||
attached with FakeQuantize modules for weight,
|
||||
used in quantization aware training.
|
||||
|
||||
We combined the interface of :class:`torch.nn.Conv1d` and
|
||||
:class:`torch.nn.BatchNorm1d`.
|
||||
|
||||
Similar to :class:`torch.nn.Conv1d`, with FakeQuantize modules initialized
|
||||
to default.
|
||||
|
||||
Attributes:
|
||||
freeze_bn:
|
||||
weight_fake_quant: fake quant module for weight
|
||||
|
||||
"""
|
||||
|
||||
_FLOAT_BN_MODULE: ClassVar[type[nn.BatchNorm1d]] = nn.BatchNorm1d
|
||||
_FLOAT_RELU_MODULE: ClassVar[type[nn.Module] | None] = None
|
||||
_FLOAT_MODULE: ClassVar[type[nn.Module]] = nni.ConvBn1d # type: ignore[assignment]
|
||||
_FLOAT_CONV_MODULE: ClassVar[type[nn.Conv1d]] = nn.Conv1d
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
# Conv1d args
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
bias=None,
|
||||
padding_mode="zeros",
|
||||
# BatchNorm1d args
|
||||
# num_features: out_channels
|
||||
eps=1e-05,
|
||||
momentum=0.1,
|
||||
# affine: True
|
||||
# track_running_stats: True
|
||||
# Args for this module
|
||||
freeze_bn=False,
|
||||
qconfig=None,
|
||||
):
|
||||
kernel_size = _single(kernel_size)
|
||||
stride = _single(stride)
|
||||
padding = _single(padding)
|
||||
dilation = _single(dilation)
|
||||
_ConvBnNd.__init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
dilation,
|
||||
False,
|
||||
_single(0),
|
||||
groups,
|
||||
bias,
|
||||
padding_mode,
|
||||
eps,
|
||||
momentum,
|
||||
freeze_bn,
|
||||
qconfig,
|
||||
dim=1,
|
||||
)
|
||||
|
||||
|
||||
class ConvBnReLU1d(ConvBn1d):
|
||||
r"""
|
||||
A ConvBnReLU1d module is a module fused from Conv1d, BatchNorm1d and ReLU,
|
||||
attached with FakeQuantize modules for weight,
|
||||
used in quantization aware training.
|
||||
|
||||
We combined the interface of :class:`torch.nn.Conv1d` and
|
||||
:class:`torch.nn.BatchNorm1d` and :class:`torch.nn.ReLU`.
|
||||
|
||||
Similar to `torch.nn.Conv1d`, with FakeQuantize modules initialized to
|
||||
default.
|
||||
|
||||
Attributes:
|
||||
weight_fake_quant: fake quant module for weight
|
||||
|
||||
"""
|
||||
|
||||
# base class defines _FLOAT_MODULE as "ConvBn1d"
|
||||
_FLOAT_MODULE: ClassVar[type[nn.Module]] = nni.ConvBnReLU1d
|
||||
_FLOAT_CONV_MODULE: ClassVar[type[nn.Conv1d]] = nn.Conv1d
|
||||
_FLOAT_BN_MODULE: ClassVar[type[nn.BatchNorm1d]] = nn.BatchNorm1d
|
||||
_FLOAT_RELU_MODULE: ClassVar[type[nn.Module] | None] = nn.ReLU
|
||||
# module class after fusing bn into conv
|
||||
_FUSED_FLOAT_MODULE: ClassVar[type[nn.Module] | None] = nni.ConvReLU1d
|
||||
|
||||
def forward(self, input):
|
||||
r"""Performs forward pass through fused Conv1d, BatchNorm1d, and ReLU."""
|
||||
return F.relu(self._forward(input))
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False):
|
||||
r"""Creates a QAT module from a floating point module."""
|
||||
return super().from_float(mod, use_precomputed_fake_quant)
|
||||
|
||||
|
||||
class ConvReLU1d(nnqat.Conv1d, nni._FusedModule):
|
||||
r"""A ConvReLU1d module is a fused module of Conv1d and ReLU, attached with
|
||||
FakeQuantize modules for weight for
|
||||
quantization aware training.
|
||||
|
||||
We combined the interface of :class:`~torch.nn.Conv1d` and
|
||||
:class:`~torch.nn.BatchNorm1d`.
|
||||
|
||||
Attributes:
|
||||
weight_fake_quant: fake quant module for weight
|
||||
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE: ClassVar[type[nni.ConvReLU1d]] = nni.ConvReLU1d # type: ignore[assignment]
|
||||
_FLOAT_CONV_MODULE: ClassVar[type[nn.Conv1d]] = nn.Conv1d
|
||||
_FLOAT_BN_MODULE: ClassVar[type[nn.Module] | None] = None
|
||||
_FLOAT_RELU_MODULE: ClassVar[type[nn.Module] | None] = nn.ReLU
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
bias=True,
|
||||
padding_mode="zeros",
|
||||
qconfig=None,
|
||||
):
|
||||
super().__init__(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
padding_mode=padding_mode,
|
||||
qconfig=qconfig,
|
||||
)
|
||||
if not qconfig:
|
||||
raise AssertionError("qconfig must be provided for QAT module")
|
||||
self.qconfig = qconfig
|
||||
self.weight_fake_quant = self.qconfig.weight()
|
||||
|
||||
def forward(self, input):
|
||||
r"""Performs forward pass through fused Conv1d and ReLU."""
|
||||
return F.relu(
|
||||
self._conv_forward(input, self.weight_fake_quant(self.weight), self.bias)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False): # type: ignore[override]
|
||||
r"""Creates a QAT module from a floating point module."""
|
||||
return super().from_float(
|
||||
mod, use_precomputed_fake_quant=use_precomputed_fake_quant
|
||||
)
|
||||
|
||||
|
||||
class ConvBn2d(_ConvBnNd, nn.Conv2d):
|
||||
r"""
|
||||
A ConvBn2d module is a module fused from Conv2d and BatchNorm2d,
|
||||
attached with FakeQuantize modules for weight,
|
||||
used in quantization aware training.
|
||||
|
||||
We combined the interface of :class:`torch.nn.Conv2d` and
|
||||
:class:`torch.nn.BatchNorm2d`.
|
||||
|
||||
Similar to :class:`torch.nn.Conv2d`, with FakeQuantize modules initialized
|
||||
to default.
|
||||
|
||||
Attributes:
|
||||
freeze_bn:
|
||||
weight_fake_quant: fake quant module for weight
|
||||
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE: ClassVar[type[nni.ConvBn2d]] = nni.ConvBn2d # type: ignore[assignment]
|
||||
_FLOAT_CONV_MODULE: ClassVar[type[nn.Conv2d]] = nn.Conv2d
|
||||
_FLOAT_BN_MODULE: ClassVar[type[nn.Module] | None] = nn.BatchNorm2d
|
||||
_FLOAT_RELU_MODULE: ClassVar[type[nn.Module] | None] = None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
# ConvNd args
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
bias=None,
|
||||
padding_mode="zeros",
|
||||
# BatchNorm2d args
|
||||
# num_features: out_channels
|
||||
eps=1e-05,
|
||||
momentum=0.1,
|
||||
# affine: True
|
||||
# track_running_stats: True
|
||||
# Args for this module
|
||||
freeze_bn=False,
|
||||
qconfig=None,
|
||||
):
|
||||
kernel_size = _pair(kernel_size)
|
||||
stride = _pair(stride)
|
||||
padding = _pair(padding)
|
||||
dilation = _pair(dilation)
|
||||
_ConvBnNd.__init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
dilation,
|
||||
False,
|
||||
_pair(0),
|
||||
groups,
|
||||
bias,
|
||||
padding_mode,
|
||||
eps,
|
||||
momentum,
|
||||
freeze_bn,
|
||||
qconfig,
|
||||
dim=2,
|
||||
)
|
||||
|
||||
|
||||
class ConvBnReLU2d(ConvBn2d):
|
||||
r"""
|
||||
A ConvBnReLU2d module is a module fused from Conv2d, BatchNorm2d and ReLU,
|
||||
attached with FakeQuantize modules for weight,
|
||||
used in quantization aware training.
|
||||
|
||||
We combined the interface of :class:`torch.nn.Conv2d` and
|
||||
:class:`torch.nn.BatchNorm2d` and :class:`torch.nn.ReLU`.
|
||||
|
||||
Similar to `torch.nn.Conv2d`, with FakeQuantize modules initialized to
|
||||
default.
|
||||
|
||||
Attributes:
|
||||
weight_fake_quant: fake quant module for weight
|
||||
|
||||
"""
|
||||
|
||||
# base class defines _FLOAT_MODULE as "ConvBn2d"
|
||||
_FLOAT_MODULE: ClassVar[type[nni.ConvBnReLU2d]] = nni.ConvBnReLU2d # type: ignore[assignment]
|
||||
_FLOAT_CONV_MODULE: ClassVar[type[nn.Conv2d]] = nn.Conv2d
|
||||
_FLOAT_BN_MODULE: ClassVar[type[nn.BatchNorm2d]] = nn.BatchNorm2d
|
||||
_FLOAT_RELU_MODULE: ClassVar[type[nn.Module] | None] = nn.ReLU
|
||||
# module class after fusing bn into conv
|
||||
_FUSED_FLOAT_MODULE: ClassVar[type[nni.ConvReLU2d] | None] = nni.ConvReLU2d
|
||||
|
||||
def forward(self, input):
|
||||
r"""Performs forward pass through fused Conv2d, BatchNorm2d, and ReLU."""
|
||||
return F.relu(self._forward(input))
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False):
|
||||
r"""Creates a QAT module from a floating point module."""
|
||||
return super().from_float(mod, use_precomputed_fake_quant)
|
||||
|
||||
|
||||
class ConvReLU2d(nnqat.Conv2d, nni._FusedModule):
|
||||
r"""A ConvReLU2d module is a fused module of Conv2d and ReLU, attached with
|
||||
FakeQuantize modules for weight for
|
||||
quantization aware training.
|
||||
|
||||
We combined the interface of :class:`~torch.nn.Conv2d` and
|
||||
:class:`~torch.nn.BatchNorm2d`.
|
||||
|
||||
Attributes:
|
||||
weight_fake_quant: fake quant module for weight
|
||||
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE: ClassVar[type[nn.Module]] = nni.ConvReLU2d # type: ignore[assignment]
|
||||
_FLOAT_CONV_MODULE: ClassVar[type[nn.Conv2d]] = nn.Conv2d
|
||||
_FLOAT_BN_MODULE: ClassVar[type[nn.Module] | None] = None
|
||||
_FLOAT_RELU_MODULE: ClassVar[type[nn.Module] | None] = nn.ReLU
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
bias=True,
|
||||
padding_mode="zeros",
|
||||
qconfig=None,
|
||||
):
|
||||
super().__init__(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
padding_mode=padding_mode,
|
||||
qconfig=qconfig,
|
||||
)
|
||||
if not qconfig:
|
||||
raise AssertionError("qconfig must be provided for QAT module")
|
||||
self.qconfig = qconfig
|
||||
self.weight_fake_quant = self.qconfig.weight()
|
||||
|
||||
def forward(self, input):
|
||||
r"""Performs forward pass through fused Conv2d and ReLU."""
|
||||
return F.relu(
|
||||
self._conv_forward(input, self.weight_fake_quant(self.weight), self.bias)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False): # type: ignore[override]
|
||||
r"""Creates a QAT module from a floating point module."""
|
||||
return super().from_float(
|
||||
mod, use_precomputed_fake_quant=use_precomputed_fake_quant
|
||||
)
|
||||
|
||||
|
||||
class ConvBn3d(_ConvBnNd, nn.Conv3d):
|
||||
r"""
|
||||
A ConvBn3d module is a module fused from Conv3d and BatchNorm3d,
|
||||
attached with FakeQuantize modules for weight,
|
||||
used in quantization aware training.
|
||||
|
||||
We combined the interface of :class:`torch.nn.Conv3d` and
|
||||
:class:`torch.nn.BatchNorm3d`.
|
||||
|
||||
Similar to :class:`torch.nn.Conv3d`, with FakeQuantize modules initialized
|
||||
to default.
|
||||
|
||||
Attributes:
|
||||
freeze_bn:
|
||||
weight_fake_quant: fake quant module for weight
|
||||
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE: ClassVar[type[nni.ConvBn3d]] = nni.ConvBn3d # type: ignore[assignment]
|
||||
_FLOAT_CONV_MODULE: ClassVar[type[nn.Conv3d]] = nn.Conv3d
|
||||
_FLOAT_BN_MODULE: ClassVar[type[nn.Module] | None] = nn.BatchNorm3d
|
||||
_FLOAT_RELU_MODULE: ClassVar[type[nn.Module] | None] = None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
# ConvNd args
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
bias=None,
|
||||
padding_mode="zeros",
|
||||
# BatchNorm3d args
|
||||
# num_features: out_channels
|
||||
eps=1e-05,
|
||||
momentum=0.1,
|
||||
# affine: True
|
||||
# track_running_stats: True
|
||||
# Args for this module
|
||||
freeze_bn=False,
|
||||
qconfig=None,
|
||||
):
|
||||
kernel_size = _triple(kernel_size)
|
||||
stride = _triple(stride)
|
||||
padding = _triple(padding)
|
||||
dilation = _triple(dilation)
|
||||
_ConvBnNd.__init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
dilation,
|
||||
False,
|
||||
_triple(0),
|
||||
groups,
|
||||
bias,
|
||||
padding_mode,
|
||||
eps,
|
||||
momentum,
|
||||
freeze_bn,
|
||||
qconfig,
|
||||
dim=3,
|
||||
)
|
||||
|
||||
|
||||
class ConvBnReLU3d(ConvBn3d):
|
||||
r"""
|
||||
A ConvBnReLU3d module is a module fused from Conv3d, BatchNorm3d and ReLU,
|
||||
attached with FakeQuantize modules for weight,
|
||||
used in quantization aware training.
|
||||
|
||||
We combined the interface of :class:`torch.nn.Conv3d` and
|
||||
:class:`torch.nn.BatchNorm3d` and :class:`torch.nn.ReLU`.
|
||||
|
||||
Similar to `torch.nn.Conv3d`, with FakeQuantize modules initialized to
|
||||
default.
|
||||
|
||||
Attributes:
|
||||
weight_fake_quant: fake quant module for weight
|
||||
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE: ClassVar[type[nni.ConvBnReLU3d]] = nni.ConvBnReLU3d # type: ignore[assignment]
|
||||
_FLOAT_CONV_MODULE: ClassVar[type[nn.Conv3d]] = nn.Conv3d
|
||||
_FLOAT_BN_MODULE: ClassVar[type[nn.BatchNorm3d]] = nn.BatchNorm3d
|
||||
_FLOAT_RELU_MODULE: ClassVar[type[nn.ReLU] | None] = nn.ReLU
|
||||
# module class after fusing bn into conv
|
||||
_FUSED_FLOAT_MODULE: ClassVar[type[nni.ConvReLU3d] | None] = nni.ConvReLU3d
|
||||
|
||||
def forward(self, input):
|
||||
r"""Performs forward pass through fused Conv3d, BatchNorm3d, and ReLU."""
|
||||
return F.relu(ConvBn3d._forward(self, input))
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False):
|
||||
r"""Creates a QAT module from a floating point module."""
|
||||
return super().from_float(
|
||||
mod, use_precomputed_fake_quant=use_precomputed_fake_quant
|
||||
)
|
||||
|
||||
|
||||
class ConvReLU3d(nnqat.Conv3d, nni._FusedModule):
|
||||
r"""A ConvReLU3d module is a fused module of Conv3d and ReLU, attached with
|
||||
FakeQuantize modules for weight for
|
||||
quantization aware training.
|
||||
|
||||
We combined the interface of :class:`~torch.nn.Conv3d` and
|
||||
:class:`~torch.nn.BatchNorm3d`.
|
||||
|
||||
Attributes:
|
||||
weight_fake_quant: fake quant module for weight
|
||||
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE: ClassVar[type[nni.ConvReLU3d]] = nni.ConvReLU3d # type: ignore[assignment]
|
||||
_FLOAT_CONV_MODULE: ClassVar[type[nn.Conv3d]] = nn.Conv3d
|
||||
_FLOAT_BN_MODULE: ClassVar[type[nn.Module] | None] = None
|
||||
_FLOAT_RELU_MODULE: ClassVar[type[nn.Module] | None] = nn.ReLU
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
bias=True,
|
||||
padding_mode="zeros",
|
||||
qconfig=None,
|
||||
):
|
||||
super().__init__(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
padding_mode=padding_mode,
|
||||
qconfig=qconfig,
|
||||
)
|
||||
if not qconfig:
|
||||
raise AssertionError("qconfig must be provided for QAT module")
|
||||
self.qconfig = qconfig
|
||||
self.weight_fake_quant = self.qconfig.weight()
|
||||
|
||||
def forward(self, input):
|
||||
r"""Performs forward pass through fused Conv3d and ReLU."""
|
||||
return F.relu(
|
||||
self._conv_forward(input, self.weight_fake_quant(self.weight), self.bias)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False): # type: ignore[override]
|
||||
r"""Creates a QAT module from a floating point module."""
|
||||
return super().from_float(
|
||||
mod, use_precomputed_fake_quant=use_precomputed_fake_quant
|
||||
)
|
||||
|
||||
|
||||
def update_bn_stats(mod):
|
||||
if type(mod) in {
|
||||
ConvBnReLU1d,
|
||||
ConvBnReLU2d,
|
||||
ConvBnReLU3d,
|
||||
ConvBn1d,
|
||||
ConvBn2d,
|
||||
ConvBn3d,
|
||||
}:
|
||||
mod.update_bn_stats()
|
||||
|
||||
|
||||
def freeze_bn_stats(mod):
|
||||
if type(mod) in {
|
||||
ConvBnReLU1d,
|
||||
ConvBnReLU2d,
|
||||
ConvBnReLU3d,
|
||||
ConvBn1d,
|
||||
ConvBn2d,
|
||||
ConvBn3d,
|
||||
}:
|
||||
mod.freeze_bn_stats()
|
||||
+199
@@ -0,0 +1,199 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import torch
|
||||
import torch.ao.nn.intrinsic as nni
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.nn import init
|
||||
from torch.nn.parameter import Parameter
|
||||
from torch.nn.utils.fusion import fuse_linear_bn_weights
|
||||
|
||||
|
||||
__all__ = [
|
||||
"LinearBn1d",
|
||||
]
|
||||
|
||||
|
||||
class LinearBn1d(nn.modules.linear.Linear, nni._FusedModule):
|
||||
r"""
|
||||
A LinearBn1d module is a module fused from Linear and BatchNorm1d, attached
|
||||
with FakeQuantize modules for weight, used in quantization aware training.
|
||||
|
||||
We combined the interface of :class:`torch.nn.Linear` and
|
||||
:class:torch.nn.BatchNorm1d`.
|
||||
|
||||
Similar to :class:`torch.nn.Linear`, with FakeQuantize modules initialized
|
||||
to default.
|
||||
|
||||
Attributes:
|
||||
freeze_bn:
|
||||
weight_fake_quant: fake quant module for weight
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
# Linear args
|
||||
in_features,
|
||||
out_features,
|
||||
bias=True,
|
||||
# BatchNorm1d args
|
||||
# num_features: out_features
|
||||
eps=1e-05,
|
||||
momentum=0.1,
|
||||
# affine: True
|
||||
# track_running_stats: True
|
||||
# Args for this module
|
||||
freeze_bn=False,
|
||||
qconfig=None,
|
||||
):
|
||||
nn.modules.linear.Linear.__init__(self, in_features, out_features, bias)
|
||||
if not qconfig:
|
||||
raise AssertionError("qconfig must be provided for QAT module")
|
||||
self.qconfig = qconfig
|
||||
self.freeze_bn = freeze_bn if self.training else True
|
||||
self.bn = nn.BatchNorm1d(out_features, eps, momentum, True, True)
|
||||
self.weight_fake_quant = self.qconfig.weight()
|
||||
if bias:
|
||||
self.bias = Parameter(torch.empty(out_features))
|
||||
else:
|
||||
self.register_parameter("bias", None)
|
||||
self.reset_bn_parameters()
|
||||
|
||||
# this needs to be called after reset_bn_parameters,
|
||||
# as they modify the same state
|
||||
if self.training:
|
||||
if freeze_bn:
|
||||
self.freeze_bn_stats()
|
||||
else:
|
||||
self.update_bn_stats()
|
||||
else:
|
||||
self.freeze_bn_stats()
|
||||
|
||||
def reset_running_stats(self):
|
||||
self.bn.reset_running_stats()
|
||||
|
||||
def reset_bn_parameters(self):
|
||||
self.bn.reset_running_stats()
|
||||
init.uniform_(self.bn.weight)
|
||||
init.zeros_(self.bn.bias)
|
||||
|
||||
def update_bn_stats(self):
|
||||
self.freeze_bn = False
|
||||
self.bn.training = True
|
||||
return self
|
||||
|
||||
def freeze_bn_stats(self):
|
||||
self.freeze_bn = True
|
||||
self.bn.training = False
|
||||
return self
|
||||
|
||||
def forward(self, input):
|
||||
if self.bn.running_var is None:
|
||||
raise AssertionError("self.bn.running_var must not be None")
|
||||
|
||||
# Scale the linear weights by BN's running statistics to reduce
|
||||
# weight jitter, see https://arxiv.org/pdf/1806.08342.pdf, page 18
|
||||
# for motivation.
|
||||
#
|
||||
# Instead of
|
||||
#
|
||||
# x1 = F.linear(x0, fq(w), b)
|
||||
# x2 = self.bn(x1)
|
||||
#
|
||||
# We have
|
||||
#
|
||||
# # scale the weight by previous batch's running statistics
|
||||
# scale_factor = bn.w / bn.running_std_from_prev_batch
|
||||
# # do the linear transformation without bias
|
||||
# x1_scaled = F.linear(x0, fq(w * scale_factor), 0)
|
||||
# # reverse the scaling and add original bias
|
||||
# x1_orig = x1_scaled / scale_factor + b
|
||||
# x2 = self.bn(x1_orig)
|
||||
|
||||
running_std = torch.sqrt(self.bn.running_var + self.bn.eps)
|
||||
scale_factor = self.bn.weight / running_std
|
||||
weight_shape = [1] * len(self.weight.shape)
|
||||
weight_shape[0] = -1
|
||||
bias_shape = [1] * len(self.weight.shape)
|
||||
bias_shape[1] = -1
|
||||
scaled_weight = self.weight_fake_quant(
|
||||
self.weight * scale_factor.reshape(weight_shape)
|
||||
)
|
||||
if self.bias is not None:
|
||||
zero_bias = torch.zeros_like(self.bias)
|
||||
else:
|
||||
zero_bias = torch.zeros(self.out_features, device=scaled_weight.device)
|
||||
linear_out = F.linear(input, scaled_weight, zero_bias)
|
||||
linear_out_orig = linear_out / scale_factor.reshape(bias_shape)
|
||||
if self.bias is not None:
|
||||
linear_out_orig = linear_out_orig + self.bias.reshape(bias_shape)
|
||||
bn_out = self.bn(linear_out_orig)
|
||||
return bn_out
|
||||
|
||||
def train(self, mode=True):
|
||||
"""
|
||||
Batchnorm's training behavior is using the self.training flag. Prevent
|
||||
changing it if BN is frozen. This makes sure that calling `model.train()`
|
||||
on a model with a frozen BN will behave properly.
|
||||
"""
|
||||
self.training = mode
|
||||
if not self.freeze_bn:
|
||||
for module in self.children():
|
||||
module.train(mode)
|
||||
return self
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False):
|
||||
r"""Create a qat module from a float module or qparams_dict
|
||||
|
||||
Args:
|
||||
mod: A float module, either produced by torch.ao.quantization
|
||||
utilities or directly from the user.
|
||||
"""
|
||||
if type(mod) is not nni.LinearBn1d:
|
||||
raise AssertionError(
|
||||
"qat."
|
||||
+ cls.__name__
|
||||
+ ".from_float only works for "
|
||||
+ nni.LinearBn1d.__name__
|
||||
)
|
||||
if not hasattr(mod, "qconfig"):
|
||||
raise AssertionError("Input float module must have qconfig defined")
|
||||
if not mod.qconfig:
|
||||
raise AssertionError("Input float module must have a valid config")
|
||||
qconfig = mod.qconfig
|
||||
linear, bn = mod[0], mod[1]
|
||||
qat_linearbn = cls(
|
||||
linear.in_features,
|
||||
linear.out_features,
|
||||
linear.bias is not None,
|
||||
bn.eps,
|
||||
bn.momentum,
|
||||
False,
|
||||
qconfig,
|
||||
)
|
||||
qat_linearbn.weight = linear.weight # type: ignore[assignment]
|
||||
qat_linearbn.bias = linear.bias # type: ignore[assignment]
|
||||
qat_linearbn.bn.weight = bn.weight # type: ignore[assignment]
|
||||
qat_linearbn.bn.bias = bn.bias # type: ignore[assignment]
|
||||
qat_linearbn.bn.running_mean = bn.running_mean # type: ignore[assignment]
|
||||
qat_linearbn.bn.running_var = bn.running_var # type: ignore[assignment]
|
||||
qat_linearbn.bn.num_batches_tracked = bn.num_batches_tracked # type: ignore[assignment]
|
||||
return qat_linearbn
|
||||
|
||||
def to_float(self):
|
||||
linear = torch.nn.Linear(self.in_features, self.out_features)
|
||||
if self.bn.running_var is None or self.bn.running_mean is None:
|
||||
raise AssertionError(
|
||||
"self.bn.running_var and self.bn.running_mean must not be None"
|
||||
)
|
||||
linear.weight, linear.bias = fuse_linear_bn_weights(
|
||||
self.weight,
|
||||
self.bias,
|
||||
self.bn.running_mean,
|
||||
self.bn.running_var,
|
||||
self.bn.eps,
|
||||
self.bn.weight,
|
||||
self.bn.bias,
|
||||
)
|
||||
return linear
|
||||
+74
@@ -0,0 +1,74 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
import torch.ao.nn.intrinsic as nni
|
||||
import torch.ao.nn.qat as nnqat
|
||||
import torch.nn.functional as F
|
||||
from torch.ao.nn.intrinsic.modules.fused import _FusedModule
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch.ao.quantization.qconfig import QConfigAny
|
||||
|
||||
|
||||
__all__ = ["LinearReLU"]
|
||||
|
||||
|
||||
class LinearReLU(nnqat.Linear, _FusedModule):
|
||||
r"""
|
||||
A LinearReLU module fused from Linear and ReLU modules, attached with
|
||||
FakeQuantize modules for weight, used in
|
||||
quantization aware training.
|
||||
|
||||
We adopt the same interface as :class:`torch.nn.Linear`.
|
||||
|
||||
Similar to `torch.ao.nn.intrinsic.LinearReLU`, with FakeQuantize modules initialized to
|
||||
default.
|
||||
|
||||
Attributes:
|
||||
weight: fake quant module for weight
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> m = nn.qat.LinearReLU(20, 30)
|
||||
>>> input = torch.randn(128, 20)
|
||||
>>> output = m(input)
|
||||
>>> print(output.size())
|
||||
torch.Size([128, 30])
|
||||
"""
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
_FLOAT_MODULE = nni.LinearReLU
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_features: int,
|
||||
out_features: int,
|
||||
bias: bool = True,
|
||||
qconfig: QConfigAny = None,
|
||||
) -> None:
|
||||
super().__init__(in_features, out_features, bias, qconfig)
|
||||
|
||||
def forward(self, input: torch.Tensor) -> torch.Tensor:
|
||||
return F.relu(F.linear(input, self.weight_fake_quant(self.weight), self.bias))
|
||||
|
||||
@classmethod
|
||||
def from_float(
|
||||
cls,
|
||||
mod: torch.nn.Module,
|
||||
use_precomputed_fake_quant: bool = False,
|
||||
) -> LinearReLU:
|
||||
return super().from_float(mod, use_precomputed_fake_quant) # type: ignore[no-untyped-call,no-any-return]
|
||||
|
||||
def to_float(self) -> nni.LinearReLU:
|
||||
linear = torch.nn.Linear(
|
||||
self.in_features, self.out_features, self.bias is not None
|
||||
)
|
||||
linear.weight = torch.nn.Parameter(self.weight.detach())
|
||||
if self.bias is not None:
|
||||
linear.bias = torch.nn.Parameter(self.bias.detach())
|
||||
relu = torch.nn.ReLU()
|
||||
return torch.ao.nn.intrinsic.LinearReLU(linear, relu) # type: ignore[no-untyped-call]
|
||||
@@ -0,0 +1,15 @@
|
||||
from .modules import * # noqa: F403
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BNReLU2d",
|
||||
"BNReLU3d",
|
||||
"ConvReLU1d",
|
||||
"ConvReLU2d",
|
||||
"ConvReLU3d",
|
||||
"LinearReLU",
|
||||
"LinearLeakyReLU",
|
||||
"LinearTanh",
|
||||
"ConvAdd2d",
|
||||
"ConvAddReLU2d",
|
||||
]
|
||||
+1
@@ -0,0 +1 @@
|
||||
from .modules import * # noqa: F403
|
||||
+6
@@ -0,0 +1,6 @@
|
||||
from .linear_relu import LinearReLU
|
||||
|
||||
|
||||
__all__ = [
|
||||
"LinearReLU",
|
||||
]
|
||||
+72
@@ -0,0 +1,72 @@
|
||||
from typing import Any
|
||||
from typing_extensions import Self
|
||||
|
||||
import torch
|
||||
import torch.ao.nn.intrinsic as nni
|
||||
import torch.ao.nn.quantized.dynamic as nnqd
|
||||
|
||||
|
||||
__all__ = ["LinearReLU"]
|
||||
|
||||
|
||||
class LinearReLU(nnqd.Linear):
|
||||
r"""
|
||||
A LinearReLU module fused from Linear and ReLU modules that can be used
|
||||
for dynamic quantization.
|
||||
Supports both, FP16 and INT8 quantization.
|
||||
|
||||
We adopt the same interface as :class:`torch.ao.nn.quantized.dynamic.Linear`.
|
||||
|
||||
Attributes:
|
||||
Same as torch.ao.nn.quantized.dynamic.Linear
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> m = nn.intrinsic.quantized.dynamic.LinearReLU(20, 30)
|
||||
>>> input = torch.randn(128, 20)
|
||||
>>> output = m(input)
|
||||
>>> print(output.size())
|
||||
torch.Size([128, 30])
|
||||
"""
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
_FLOAT_MODULE = nni.LinearReLU
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_features: int,
|
||||
out_features: int,
|
||||
bias: bool = True,
|
||||
dtype: torch.dtype = torch.qint8,
|
||||
) -> None:
|
||||
super().__init__(in_features, out_features, bias, dtype)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
if self._packed_params.dtype == torch.qint8:
|
||||
# TODO check if we should set reduce_rage = True by default here
|
||||
Y = torch.ops.quantized.linear_relu_dynamic(
|
||||
x, self._packed_params._packed_params, reduce_range=True
|
||||
)
|
||||
elif self._packed_params.dtype == torch.float16:
|
||||
Y = torch.ops.quantized.linear_relu_dynamic_fp16(
|
||||
x, self._packed_params._packed_params
|
||||
)
|
||||
else:
|
||||
raise RuntimeError("Unsupported dtype on dynamic quantized linear relu!")
|
||||
return Y.to(x.dtype)
|
||||
|
||||
def _get_name(self) -> str:
|
||||
return "DynamicQuantizedLinearReLU"
|
||||
|
||||
@classmethod
|
||||
def from_float(
|
||||
cls, mod: torch.nn.Module, use_precomputed_fake_quant: bool = False
|
||||
) -> Self:
|
||||
return super().from_float(
|
||||
mod, use_precomputed_fake_quant=use_precomputed_fake_quant
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_reference(cls, ref_qlinear_relu: Any) -> Self: # type: ignore[override]
|
||||
return super().from_reference(ref_qlinear_relu[0])
|
||||
+18
@@ -0,0 +1,18 @@
|
||||
from .bn_relu import BNReLU2d, BNReLU3d
|
||||
from .conv_add import ConvAdd2d, ConvAddReLU2d
|
||||
from .conv_relu import ConvReLU1d, ConvReLU2d, ConvReLU3d
|
||||
from .linear_relu import LinearLeakyReLU, LinearReLU, LinearTanh
|
||||
|
||||
|
||||
__all__ = [
|
||||
"LinearReLU",
|
||||
"ConvReLU1d",
|
||||
"ConvReLU2d",
|
||||
"ConvReLU3d",
|
||||
"BNReLU2d",
|
||||
"BNReLU3d",
|
||||
"LinearLeakyReLU",
|
||||
"LinearTanh",
|
||||
"ConvAdd2d",
|
||||
"ConvAddReLU2d",
|
||||
]
|
||||
+113
@@ -0,0 +1,113 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import torch
|
||||
import torch.ao.nn.intrinsic
|
||||
import torch.ao.nn.intrinsic.qat
|
||||
import torch.ao.nn.quantized as nnq
|
||||
|
||||
|
||||
__all__ = ["BNReLU2d", "BNReLU3d"]
|
||||
|
||||
|
||||
class BNReLU2d(nnq.BatchNorm2d):
|
||||
r"""
|
||||
A BNReLU2d module is a fused module of BatchNorm2d and ReLU
|
||||
|
||||
We adopt the same interface as :class:`torch.ao.nn.quantized.BatchNorm2d`.
|
||||
|
||||
Attributes:
|
||||
Same as torch.ao.nn.quantized.BatchNorm2d
|
||||
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE = torch.ao.nn.intrinsic.BNReLU2d
|
||||
|
||||
def __init__(self, num_features, eps=1e-5, momentum=0.1, device=None, dtype=None):
|
||||
super().__init__(
|
||||
num_features, eps=eps, momentum=momentum, device=device, dtype=dtype
|
||||
)
|
||||
|
||||
def forward(self, input):
|
||||
r"""Applies fused BatchNorm2d and ReLU."""
|
||||
# Temporarily using len(shape) instead of ndim due to JIT issue
|
||||
# https://github.com/pytorch/pytorch/issues/23890
|
||||
if len(input.shape) != 4:
|
||||
raise ValueError("Input shape must be `(N, C, H, W)`!")
|
||||
return torch.ops.quantized.batch_norm2d_relu(
|
||||
input,
|
||||
self.weight,
|
||||
self.bias,
|
||||
self.running_mean,
|
||||
self.running_var,
|
||||
self.eps,
|
||||
self.scale,
|
||||
self.zero_point,
|
||||
)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedBNReLU2d"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False): # type: ignore[override]
|
||||
r"""Creates a quantized module from a float module."""
|
||||
# TODO: Add qat support for BNReLU2d
|
||||
return super().from_float(
|
||||
mod, use_precomputed_fake_quant=use_precomputed_fake_quant
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_reference(cls, bn_relu, output_scale, output_zero_point):
|
||||
r"""Creates a quantized module from a reference module."""
|
||||
return super().from_reference(bn_relu[0], output_scale, output_zero_point)
|
||||
|
||||
|
||||
class BNReLU3d(nnq.BatchNorm3d):
|
||||
r"""
|
||||
A BNReLU3d module is a fused module of BatchNorm3d and ReLU
|
||||
|
||||
We adopt the same interface as :class:`torch.ao.nn.quantized.BatchNorm3d`.
|
||||
|
||||
Attributes:
|
||||
Same as torch.ao.nn.quantized.BatchNorm3d
|
||||
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE = torch.ao.nn.intrinsic.BNReLU3d
|
||||
|
||||
def __init__(self, num_features, eps=1e-5, momentum=0.1, device=None, dtype=None):
|
||||
super().__init__(
|
||||
num_features, eps=eps, momentum=momentum, device=device, dtype=dtype
|
||||
)
|
||||
|
||||
def forward(self, input):
|
||||
r"""Applies fused BatchNorm3d and ReLU."""
|
||||
# Temporarily using len(shape) instead of ndim due to JIT issue
|
||||
# https://github.com/pytorch/pytorch/issues/23890
|
||||
if len(input.shape) != 5:
|
||||
raise ValueError("Input shape must be `(N, C, D, H, W)`!")
|
||||
return torch.ops.quantized.batch_norm3d_relu(
|
||||
input,
|
||||
self.weight,
|
||||
self.bias,
|
||||
self.running_mean,
|
||||
self.running_var,
|
||||
self.eps,
|
||||
self.scale,
|
||||
self.zero_point,
|
||||
)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedBNReLU3d"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False): # type: ignore[override]
|
||||
r"""Creates a quantized module from a float module."""
|
||||
# TODO: Add qat support for BNReLU3d
|
||||
return super().from_float(
|
||||
mod, use_precomputed_fake_quant=use_precomputed_fake_quant
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_reference(cls, bn_relu, output_scale, output_zero_point):
|
||||
r"""Creates a quantized module from a reference module."""
|
||||
return super().from_reference(bn_relu[0], output_scale, output_zero_point)
|
||||
+153
@@ -0,0 +1,153 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import torch
|
||||
import torch.ao.nn.intrinsic
|
||||
import torch.ao.nn.intrinsic.qat
|
||||
import torch.ao.nn.quantized as nnq
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
_reverse_repeat_padding = nnq.modules.conv._reverse_repeat_padding
|
||||
|
||||
|
||||
class ConvAdd2d(nnq.Conv2d):
|
||||
r"""
|
||||
A ConvAdd2d module is a fused module of Conv2d and Add
|
||||
|
||||
We adopt the same interface as :class:`torch.ao.nn.quantized.Conv2d`.
|
||||
|
||||
Attributes:
|
||||
Same as torch.ao.nn.quantized.Conv2d
|
||||
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE = torch.ao.nn.intrinsic.ConvAdd2d # type: ignore[assignment]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
bias=True,
|
||||
padding_mode="zeros",
|
||||
device=None,
|
||||
dtype=None,
|
||||
):
|
||||
super().__init__(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
padding_mode=padding_mode,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
def forward(self, input, extra_input): # type: ignore[override]
|
||||
r"""Applies fused quantized Conv2d and addition."""
|
||||
# Temporarily using len(shape) instead of ndim due to JIT issue
|
||||
# https://github.com/pytorch/pytorch/issues/23890
|
||||
if len(input.shape) != 4:
|
||||
raise ValueError("Input shape must be `(N, C, H, W)`!")
|
||||
if self.padding_mode != "zeros":
|
||||
_reversed_padding_repeated_twice = _reverse_repeat_padding(self.padding)
|
||||
input = F.pad(
|
||||
input, _reversed_padding_repeated_twice, mode=self.padding_mode
|
||||
)
|
||||
return torch.ops.quantized.conv2d_add(
|
||||
input, extra_input, self._packed_params, self.scale, self.zero_point
|
||||
)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedConvAdd2d"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False): # type: ignore[override]
|
||||
r"""Creates a quantized module from a float module."""
|
||||
return super().from_float(
|
||||
mod, use_precomputed_fake_quant=use_precomputed_fake_quant
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_reference(cls, ref_qconv, output_scale, output_zero_point):
|
||||
r"""Creates a quantized module from a reference module."""
|
||||
return super().from_reference(ref_qconv[0], output_scale, output_zero_point)
|
||||
|
||||
|
||||
class ConvAddReLU2d(nnq.Conv2d):
|
||||
r"""
|
||||
A ConvAddReLU2d module is a fused module of Conv2d, Add and Relu
|
||||
|
||||
We adopt the same interface as :class:`torch.ao.nn.quantized.Conv2d`.
|
||||
|
||||
Attributes:
|
||||
Same as torch.ao.nn.quantized.Conv2d
|
||||
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE = torch.ao.nn.intrinsic.ConvAddReLU2d # type: ignore[assignment]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
bias=True,
|
||||
padding_mode="zeros",
|
||||
device=None,
|
||||
dtype=None,
|
||||
):
|
||||
super().__init__(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
padding_mode=padding_mode,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
def forward(self, input, extra_input): # type: ignore[override]
|
||||
r"""Applies fused quantized Conv2d, addition, and ReLU."""
|
||||
# Temporarily using len(shape) instead of ndim due to JIT issue
|
||||
# https://github.com/pytorch/pytorch/issues/23890
|
||||
if len(input.shape) != 4:
|
||||
raise ValueError("Input shape must be `(N, C, H, W)`!")
|
||||
if self.padding_mode != "zeros":
|
||||
_reversed_padding_repeated_twice = _reverse_repeat_padding(self.padding)
|
||||
input = F.pad(
|
||||
input, _reversed_padding_repeated_twice, mode=self.padding_mode
|
||||
)
|
||||
return torch.ops.quantized.conv2d_add_relu(
|
||||
input, extra_input, self._packed_params, self.scale, self.zero_point
|
||||
)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedConvAddReLU2d"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False): # type: ignore[override]
|
||||
r"""Creates a quantized module from a float module."""
|
||||
return super().from_float(
|
||||
mod, use_precomputed_fake_quant=use_precomputed_fake_quant
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_reference(cls, ref_qconv, output_scale, output_zero_point):
|
||||
r"""Creates a quantized module from a reference module."""
|
||||
return super().from_reference(ref_qconv[0], output_scale, output_zero_point)
|
||||
+289
@@ -0,0 +1,289 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import torch
|
||||
import torch.ao.nn.intrinsic
|
||||
import torch.ao.nn.intrinsic.qat
|
||||
import torch.ao.nn.quantized as nnq
|
||||
import torch.nn.functional as F
|
||||
from torch.nn.utils import fuse_conv_bn_weights
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ConvReLU1d",
|
||||
"ConvReLU2d",
|
||||
"ConvReLU3d",
|
||||
]
|
||||
|
||||
_reverse_repeat_padding = nnq.modules.conv._reverse_repeat_padding
|
||||
|
||||
|
||||
# TODO: factor out the common parts to ConvNd
|
||||
class ConvReLU1d(nnq.Conv1d):
|
||||
r"""
|
||||
A ConvReLU1d module is a fused module of Conv1d and ReLU
|
||||
|
||||
We adopt the same interface as :class:`torch.ao.nn.quantized.Conv1d`.
|
||||
|
||||
Attributes:
|
||||
Same as torch.ao.nn.quantized.Conv1d
|
||||
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE = torch.ao.nn.intrinsic.ConvReLU1d # type: ignore[assignment]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
bias=True,
|
||||
padding_mode="zeros",
|
||||
device=None,
|
||||
dtype=None,
|
||||
):
|
||||
super().__init__(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
padding_mode=padding_mode,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
def forward(self, input):
|
||||
r"""Applies fused quantized Conv1d and ReLU."""
|
||||
# Temporarily using len(shape) instead of ndim due to JIT issue
|
||||
# https://github.com/pytorch/pytorch/issues/23890
|
||||
if len(input.shape) != 3:
|
||||
raise ValueError("Input shape must be `(N, C, L)`!")
|
||||
if self.padding_mode != "zeros":
|
||||
# Padding in Conv1d is stored as (p, p), need to get (p,)
|
||||
_reversed_padding_repeated_twice = _reverse_repeat_padding(self.padding[:1])
|
||||
input = F.pad(
|
||||
input, _reversed_padding_repeated_twice, mode=self.padding_mode
|
||||
)
|
||||
return torch.ops.quantized.conv1d_relu(
|
||||
input, self._packed_params, self.scale, self.zero_point
|
||||
)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedConvReLU1d"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False): # type: ignore[override]
|
||||
r"""Creates a quantized module from a float module."""
|
||||
if type(mod) is torch.ao.nn.intrinsic.qat.ConvBnReLU1d:
|
||||
if mod.bn.running_var is None or mod.bn.running_mean is None:
|
||||
raise AssertionError(
|
||||
"mod.bn.running_var and mod.bn.running_mean must not be None"
|
||||
)
|
||||
mod.weight, mod.bias = fuse_conv_bn_weights(
|
||||
mod.weight,
|
||||
mod.bias,
|
||||
mod.bn.running_mean,
|
||||
mod.bn.running_var,
|
||||
mod.bn.eps,
|
||||
mod.bn.weight,
|
||||
mod.bn.bias,
|
||||
)
|
||||
return super().from_float(mod, use_precomputed_fake_quant)
|
||||
|
||||
@classmethod
|
||||
def from_reference(cls, ref_qconv, output_scale, output_zero_point):
|
||||
r"""Creates a quantized module from a reference module."""
|
||||
if type(ref_qconv) is torch.ao.nn.intrinsic.ConvBnReLU1d:
|
||||
raise AssertionError(
|
||||
"BatchNorm1d should be fused into Conv1d before converting to reference module"
|
||||
)
|
||||
return super().from_reference(ref_qconv[0], output_scale, output_zero_point)
|
||||
|
||||
|
||||
class ConvReLU2d(nnq.Conv2d):
|
||||
r"""
|
||||
A ConvReLU2d module is a fused module of Conv2d and ReLU
|
||||
|
||||
We adopt the same interface as :class:`torch.ao.nn.quantized.Conv2d`.
|
||||
|
||||
Attributes:
|
||||
Same as torch.ao.nn.quantized.Conv2d
|
||||
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE = torch.ao.nn.intrinsic.ConvReLU2d # type: ignore[assignment]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
bias=True,
|
||||
padding_mode="zeros",
|
||||
device=None,
|
||||
dtype=None,
|
||||
):
|
||||
super().__init__(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
padding_mode=padding_mode,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
def forward(self, input):
|
||||
r"""Applies fused quantized Conv2d and ReLU."""
|
||||
# Temporarily using len(shape) instead of ndim due to JIT issue
|
||||
# https://github.com/pytorch/pytorch/issues/23890
|
||||
if len(input.shape) != 4:
|
||||
raise ValueError("Input shape must be `(N, C, H, W)`!")
|
||||
if self.padding_mode != "zeros":
|
||||
_reversed_padding_repeated_twice = _reverse_repeat_padding(self.padding)
|
||||
input = F.pad(
|
||||
input, _reversed_padding_repeated_twice, mode=self.padding_mode
|
||||
)
|
||||
return torch.ops.quantized.conv2d_relu(
|
||||
input, self._packed_params, self.scale, self.zero_point
|
||||
)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedConvReLU2d"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False): # type: ignore[override]
|
||||
r"""Creates a quantized module from a float module."""
|
||||
if type(mod) is torch.ao.nn.intrinsic.qat.ConvBnReLU2d:
|
||||
if mod.bn.running_var is None or mod.bn.running_mean is None:
|
||||
raise AssertionError(
|
||||
"mod.bn.running_var and mod.bn.running_mean must not be None"
|
||||
)
|
||||
mod.weight, mod.bias = fuse_conv_bn_weights(
|
||||
mod.weight,
|
||||
mod.bias,
|
||||
mod.bn.running_mean,
|
||||
mod.bn.running_var,
|
||||
mod.bn.eps,
|
||||
mod.bn.weight,
|
||||
mod.bn.bias,
|
||||
)
|
||||
return super().from_float(
|
||||
mod, use_precomputed_fake_quant=use_precomputed_fake_quant
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_reference(cls, ref_qconv, output_scale, output_zero_point):
|
||||
r"""Creates a quantized module from a reference module."""
|
||||
if type(ref_qconv) is torch.ao.nn.intrinsic.ConvBnReLU2d:
|
||||
raise AssertionError(
|
||||
"BatchNorm2d should be fused into Conv2d before converting to reference module"
|
||||
)
|
||||
return super().from_reference(ref_qconv[0], output_scale, output_zero_point)
|
||||
|
||||
|
||||
class ConvReLU3d(nnq.Conv3d):
|
||||
r"""
|
||||
A ConvReLU3d module is a fused module of Conv3d and ReLU
|
||||
|
||||
We adopt the same interface as :class:`torch.ao.nn.quantized.Conv3d`.
|
||||
|
||||
Attributes: Same as torch.ao.nn.quantized.Conv3d
|
||||
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE = torch.ao.nn.intrinsic.ConvReLU3d # type: ignore[assignment]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
bias=True,
|
||||
padding_mode="zeros",
|
||||
device=None,
|
||||
dtype=None,
|
||||
):
|
||||
if padding_mode == "reflect":
|
||||
raise AssertionError("Conv3d does not support reflection padding")
|
||||
super().__init__(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
padding_mode=padding_mode,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
def forward(self, input):
|
||||
r"""Applies fused quantized Conv3d and ReLU."""
|
||||
# Temporarily using len(shape) instead of ndim due to JIT issue
|
||||
# https://github.com/pytorch/pytorch/issues/23890
|
||||
if len(input.shape) != 5:
|
||||
raise ValueError("Input shape must be `(N, C, D, H, W)`!")
|
||||
if self.padding_mode != "zeros":
|
||||
_reversed_padding_repeated_twice = _reverse_repeat_padding(self.padding)
|
||||
input = F.pad(
|
||||
input, _reversed_padding_repeated_twice, mode=self.padding_mode
|
||||
)
|
||||
return torch.ops.quantized.conv3d_relu(
|
||||
input, self._packed_params, self.scale, self.zero_point
|
||||
)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedConvReLU3d"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False): # type: ignore[override]
|
||||
r"""Creates a quantized module from a float module."""
|
||||
if type(mod) is torch.ao.nn.intrinsic.qat.ConvBnReLU3d:
|
||||
if mod.bn.running_var is None or mod.bn.running_mean is None:
|
||||
raise AssertionError(
|
||||
"mod.bn.running_var and mod.bn.running_mean must not be None"
|
||||
)
|
||||
mod.weight, mod.bias = fuse_conv_bn_weights(
|
||||
mod.weight,
|
||||
mod.bias,
|
||||
mod.bn.running_mean,
|
||||
mod.bn.running_var,
|
||||
mod.bn.eps,
|
||||
mod.bn.weight,
|
||||
mod.bn.bias,
|
||||
)
|
||||
return super().from_float(
|
||||
mod, use_precomputed_fake_quant=use_precomputed_fake_quant
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_reference(cls, ref_qconv, output_scale, output_zero_point):
|
||||
r"""Creates a quantized module from a reference module."""
|
||||
if type(ref_qconv) is torch.ao.nn.intrinsic.ConvBnReLU3d:
|
||||
raise AssertionError(
|
||||
"BatchNorm3d should be fused into Conv3d before converting to reference module"
|
||||
)
|
||||
return super().from_reference(ref_qconv[0], output_scale, output_zero_point)
|
||||
+198
@@ -0,0 +1,198 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import torch
|
||||
import torch.ao.nn.intrinsic as nni
|
||||
import torch.ao.nn.quantized as nnq
|
||||
from torch.ao.nn.quantized.modules.utils import _quantize_weight
|
||||
|
||||
|
||||
__all__ = [
|
||||
"LinearReLU",
|
||||
"LinearLeakyReLU",
|
||||
"LinearTanh",
|
||||
]
|
||||
|
||||
|
||||
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):
|
||||
if type(mod) is not nni.LinearLeakyReLU:
|
||||
raise AssertionError("Input float module should be LinearLeakyReLU")
|
||||
if not hasattr(mod, "qconfig"):
|
||||
raise AssertionError("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() # type: ignore[union-attr, operator]
|
||||
weight_post_process(mod.weight)
|
||||
dtype = weight_post_process.dtype
|
||||
act_scale, act_zp = activation_post_process.calculate_qparams() # type: ignore[union-attr,operator]
|
||||
if dtype != torch.qint8:
|
||||
raise AssertionError(
|
||||
f"Weight observer must have dtype torch.qint8, got {dtype}"
|
||||
)
|
||||
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) # type: ignore[arg-type]
|
||||
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):
|
||||
if type(mod) is not nni.LinearTanh:
|
||||
raise AssertionError("Input float module should be LinearTanh")
|
||||
if not hasattr(mod, "qconfig"):
|
||||
raise AssertionError("Input float module must have qconfig defined")
|
||||
activation_post_process = mod.activation_post_process
|
||||
mod = mod[0]
|
||||
weight_post_process = mod.qconfig.weight() # type: ignore[union-attr,operator]
|
||||
weight_post_process(mod.weight)
|
||||
dtype = weight_post_process.dtype
|
||||
act_scale, act_zp = activation_post_process.calculate_qparams() # type: ignore[union-attr,operator]
|
||||
if dtype != torch.qint8:
|
||||
raise AssertionError(
|
||||
f"Weight observer must have dtype torch.qint8, got {dtype}"
|
||||
)
|
||||
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) # type: ignore[arg-type]
|
||||
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
|
||||
Reference in New Issue
Block a user