Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,31 @@
|
||||
# torch.ao is a package with a lot of interdependencies.
|
||||
# We will use lazy import to avoid cyclic dependencies here.
|
||||
|
||||
from typing import TYPE_CHECKING as _TYPE_CHECKING
|
||||
|
||||
|
||||
if _TYPE_CHECKING:
|
||||
from types import ModuleType
|
||||
|
||||
from torch.ao import ( # noqa: TC004
|
||||
nn as nn,
|
||||
ns as ns,
|
||||
pruning as pruning,
|
||||
quantization as quantization,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"nn",
|
||||
"ns",
|
||||
"pruning",
|
||||
"quantization",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str) -> "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,35 @@
|
||||
# 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/
|
||||
|
||||
from typing import TYPE_CHECKING as _TYPE_CHECKING
|
||||
|
||||
|
||||
if _TYPE_CHECKING:
|
||||
from types import ModuleType
|
||||
|
||||
from torch.ao.nn import ( # noqa: TC004
|
||||
intrinsic as intrinsic,
|
||||
qat as qat,
|
||||
quantizable as quantizable,
|
||||
quantized as quantized,
|
||||
sparse as sparse,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"intrinsic",
|
||||
"qat",
|
||||
"quantizable",
|
||||
"quantized",
|
||||
"sparse",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str) -> "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 @@
|
||||
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
|
||||
@@ -0,0 +1 @@
|
||||
from .modules import * # noqa: F403
|
||||
@@ -0,0 +1 @@
|
||||
from .modules import * # noqa: F403
|
||||
@@ -0,0 +1,4 @@
|
||||
from .linear import Linear
|
||||
|
||||
|
||||
__all__ = ["Linear"]
|
||||
@@ -0,0 +1,40 @@
|
||||
from typing import Optional, TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch.ao.quantization.qconfig import QConfig
|
||||
|
||||
|
||||
__all__ = ["Linear"]
|
||||
|
||||
|
||||
class Linear(torch.ao.nn.qat.Linear):
|
||||
r"""
|
||||
A linear module attached with FakeQuantize modules for weight,
|
||||
used for dynamic quantization aware training.
|
||||
|
||||
We adopt the same interface as `torch.nn.Linear`, please see
|
||||
https://pytorch.org/docs/stable/nn.html#torch.nn.Linear
|
||||
for documentation.
|
||||
|
||||
Similar to `torch.nn.Linear`, with FakeQuantize modules initialized to
|
||||
default.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_features: int,
|
||||
out_features: int,
|
||||
bias: bool = True,
|
||||
qconfig: Optional["QConfig"] = None,
|
||||
device: int | str | torch.device | None = None,
|
||||
dtype: str | None = None,
|
||||
) -> None:
|
||||
super().__init__(in_features, out_features, bias, qconfig, device, dtype)
|
||||
if not torch.ao.quantization.qconfig._activation_is_memoryless(qconfig): # type: ignore[arg-type]
|
||||
raise ValueError(
|
||||
"Dynamic QAT requires a memoryless observer."
|
||||
+ "This means a MovingAverage observer with averaging constant equal to 1"
|
||||
)
|
||||
@@ -0,0 +1,13 @@
|
||||
from .conv import Conv1d, Conv2d, Conv3d
|
||||
from .embedding_ops import Embedding, EmbeddingBag
|
||||
from .linear import Linear
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Linear",
|
||||
"Conv1d",
|
||||
"Conv2d",
|
||||
"Conv3d",
|
||||
"Embedding",
|
||||
"EmbeddingBag",
|
||||
]
|
||||
@@ -0,0 +1,318 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from typing import ClassVar, Literal
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.ao.nn.intrinsic import _FusedModule
|
||||
from torch.nn.common_types import _size_1_t, _size_2_t, _size_3_t
|
||||
from torch.nn.modules.utils import _pair, _single, _triple
|
||||
|
||||
|
||||
__all__ = ["Conv1d", "Conv2d", "Conv3d"]
|
||||
|
||||
|
||||
class _ConvNd(nn.modules.conv._ConvNd):
|
||||
_FLOAT_MODULE: ClassVar[type[nn.modules.conv._ConvNd]]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size: tuple[int, ...],
|
||||
stride: tuple[int, ...],
|
||||
padding: str | tuple[int, ...],
|
||||
dilation: tuple[int, ...],
|
||||
transposed: bool,
|
||||
output_padding: tuple[int, ...],
|
||||
groups: int,
|
||||
bias: bool,
|
||||
padding_mode: Literal["zeros", "reflect", "replicate", "circular"],
|
||||
qconfig=None,
|
||||
device=None,
|
||||
dtype=None,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
nn.modules.conv._ConvNd.__init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
dilation,
|
||||
transposed,
|
||||
output_padding,
|
||||
groups,
|
||||
bias,
|
||||
padding_mode,
|
||||
**factory_kwargs,
|
||||
)
|
||||
if not qconfig:
|
||||
raise AssertionError("qconfig must be provided for QAT module")
|
||||
self.qconfig = qconfig
|
||||
self.weight_fake_quant = qconfig.weight(factory_kwargs=factory_kwargs)
|
||||
|
||||
def forward(self, input):
|
||||
return self._conv_forward(input, self.weight_fake_quant(self.weight), self.bias)
|
||||
|
||||
@staticmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False):
|
||||
r"""Create a qat module from a float module
|
||||
|
||||
Args:
|
||||
`mod`: a float module, either produced by torch.ao.quantization utilities
|
||||
or directly from user
|
||||
"""
|
||||
if type(mod) is not cls._FLOAT_MODULE:
|
||||
raise AssertionError(
|
||||
f"qat.{cls.__name__}.from_float only works for "
|
||||
f"{cls._FLOAT_MODULE.__name__}, got {type(mod).__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")
|
||||
if issubclass(type(mod), _FusedModule):
|
||||
mod = mod[0]
|
||||
qconfig = mod.qconfig
|
||||
qat_conv = cls(
|
||||
mod.in_channels,
|
||||
mod.out_channels,
|
||||
mod.kernel_size,
|
||||
stride=mod.stride,
|
||||
padding=mod.padding,
|
||||
dilation=mod.dilation,
|
||||
groups=mod.groups,
|
||||
bias=mod.bias is not None,
|
||||
padding_mode=mod.padding_mode,
|
||||
qconfig=qconfig,
|
||||
)
|
||||
qat_conv.weight = mod.weight
|
||||
qat_conv.bias = mod.bias
|
||||
return qat_conv
|
||||
|
||||
def to_float(self):
|
||||
"""This works for both single qat conv, and the qat conv - relu modules
|
||||
to convert the qat module to a floating point module
|
||||
"""
|
||||
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())
|
||||
# conv relu
|
||||
if issubclass(cls, _FusedModule):
|
||||
modules = [conv]
|
||||
if not hasattr(cls, "_FLOAT_RELU_MODULE"):
|
||||
raise AssertionError(
|
||||
f"{cls.__name__} must have _FLOAT_RELU_MODULE attribute"
|
||||
)
|
||||
relu = cls._FLOAT_RELU_MODULE()
|
||||
modules.append(relu)
|
||||
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
fused = cls._FLOAT_MODULE(*modules)
|
||||
fused.train(self.training)
|
||||
return fused
|
||||
else:
|
||||
return conv
|
||||
|
||||
|
||||
class Conv1d(_ConvNd, nn.Conv1d):
|
||||
r"""
|
||||
A Conv1d module attached with FakeQuantize modules for weight,
|
||||
used for quantization aware training.
|
||||
|
||||
We adopt the same interface as :class:`~torch.nn.Conv1d`
|
||||
|
||||
Similar to :class:`~torch.nn.Conv2d`, with FakeQuantize modules initialized to
|
||||
default.
|
||||
|
||||
Attributes:
|
||||
weight_fake_quant: fake quant module for weight
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE: ClassVar[type[nn.Conv1d]] = nn.Conv1d
|
||||
_FLOAT_CONV_MODULE: ClassVar[type[nn.Conv1d]] = nn.Conv1d
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size: _size_1_t,
|
||||
stride: _size_1_t = 1,
|
||||
padding: str | _size_1_t = 0,
|
||||
dilation: _size_1_t = 1,
|
||||
groups: int = 1,
|
||||
bias: bool = True,
|
||||
padding_mode: Literal["zeros", "reflect", "replicate", "circular"] = "zeros",
|
||||
qconfig=None,
|
||||
device=None,
|
||||
dtype=None,
|
||||
) -> None:
|
||||
kernel_size_ = _single(kernel_size)
|
||||
stride_ = _single(stride)
|
||||
padding_ = padding if isinstance(padding, str) else _single(padding)
|
||||
dilation_ = _single(dilation)
|
||||
super().__init__(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size_,
|
||||
stride=stride_,
|
||||
padding=padding_,
|
||||
dilation=dilation_,
|
||||
transposed=False,
|
||||
output_padding=_single(0),
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
padding_mode=padding_mode,
|
||||
qconfig=qconfig,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False): # type: ignore[override]
|
||||
return super().from_float(
|
||||
cls, mod, use_precomputed_fake_quant=use_precomputed_fake_quant
|
||||
)
|
||||
|
||||
|
||||
class Conv2d(_ConvNd, nn.Conv2d):
|
||||
r"""
|
||||
A Conv2d module attached with FakeQuantize modules for weight,
|
||||
used for quantization aware training.
|
||||
|
||||
We adopt the same interface as `torch.nn.Conv2d`, please see
|
||||
https://pytorch.org/docs/stable/nn.html?highlight=conv2d#torch.nn.Conv2d
|
||||
for documentation.
|
||||
|
||||
Similar to `torch.nn.Conv2d`, with FakeQuantize modules initialized to
|
||||
default.
|
||||
|
||||
Attributes:
|
||||
weight_fake_quant: fake quant module for weight
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE: ClassVar[type[nn.Conv2d]] = nn.Conv2d
|
||||
_FLOAT_CONV_MODULE: ClassVar[type[nn.Conv2d]] = nn.Conv2d
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size: _size_2_t,
|
||||
stride: _size_2_t = 1,
|
||||
padding: str | _size_2_t = 0,
|
||||
dilation: _size_2_t = 1,
|
||||
groups: int = 1,
|
||||
bias: bool = True,
|
||||
padding_mode: Literal["zeros", "reflect", "replicate", "circular"] = "zeros",
|
||||
qconfig=None,
|
||||
device=None,
|
||||
dtype=None,
|
||||
) -> None:
|
||||
kernel_size_ = _pair(kernel_size)
|
||||
stride_ = _pair(stride)
|
||||
padding_ = padding if isinstance(padding, str) else _pair(padding)
|
||||
dilation_ = _pair(dilation)
|
||||
super().__init__(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size_,
|
||||
stride=stride_,
|
||||
padding=padding_,
|
||||
dilation=dilation_,
|
||||
transposed=False,
|
||||
output_padding=_pair(0),
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
padding_mode=padding_mode,
|
||||
qconfig=qconfig,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
def forward(self, input):
|
||||
return 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]
|
||||
return super().from_float(
|
||||
cls, mod, use_precomputed_fake_quant=use_precomputed_fake_quant
|
||||
)
|
||||
|
||||
|
||||
class Conv3d(_ConvNd, nn.Conv3d):
|
||||
r"""
|
||||
A Conv3d module attached with FakeQuantize modules for weight,
|
||||
used for quantization aware training.
|
||||
|
||||
We adopt the same interface as `torch.nn.Conv3d`, please see
|
||||
https://pytorch.org/docs/stable/nn.html?highlight=conv3d#torch.nn.Conv3d
|
||||
for documentation.
|
||||
|
||||
Similar to `torch.nn.Conv3d`, with FakeQuantize modules initialized to
|
||||
default.
|
||||
|
||||
Attributes:
|
||||
weight_fake_quant: fake quant module for weight
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE: ClassVar[type[nn.Conv3d]] = nn.Conv3d
|
||||
_FLOAT_CONV_MODULE: ClassVar[type[nn.Conv3d]] = nn.Conv3d
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size: _size_3_t,
|
||||
stride: _size_3_t = 1,
|
||||
padding: str | _size_3_t = 0,
|
||||
dilation: _size_3_t = 1,
|
||||
groups: int = 1,
|
||||
bias: bool = True,
|
||||
padding_mode: Literal["zeros", "reflect", "replicate", "circular"] = "zeros",
|
||||
qconfig=None,
|
||||
device=None,
|
||||
dtype=None,
|
||||
) -> None:
|
||||
kernel_size_ = _triple(kernel_size)
|
||||
stride_ = _triple(stride)
|
||||
padding_ = padding if isinstance(padding, str) else _triple(padding)
|
||||
dilation_ = _triple(dilation)
|
||||
super().__init__(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size_,
|
||||
stride=stride_,
|
||||
padding=padding_,
|
||||
dilation=dilation_,
|
||||
transposed=False,
|
||||
output_padding=_triple(0),
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
padding_mode=padding_mode,
|
||||
qconfig=qconfig,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
def forward(self, input):
|
||||
return 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]
|
||||
return super().from_float(
|
||||
cls, mod, use_precomputed_fake_quant=use_precomputed_fake_quant
|
||||
)
|
||||
@@ -0,0 +1,264 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
__all__ = ["Embedding", "EmbeddingBag"]
|
||||
|
||||
|
||||
class Embedding(nn.Embedding):
|
||||
r"""
|
||||
An embedding bag module attached with FakeQuantize modules for weight,
|
||||
used for quantization aware training.
|
||||
|
||||
We adopt the same interface as `torch.nn.Embedding`, please see
|
||||
https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html#torch.nn.Embedding
|
||||
for documentation.
|
||||
|
||||
Similar to `torch.nn.Embedding`, with FakeQuantize modules initialized to
|
||||
default.
|
||||
|
||||
Attributes:
|
||||
weight: fake quant module for weight
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE = nn.Embedding
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_embeddings,
|
||||
embedding_dim,
|
||||
padding_idx=None,
|
||||
max_norm=None,
|
||||
norm_type=2.0,
|
||||
scale_grad_by_freq=False,
|
||||
sparse=False,
|
||||
_weight=None,
|
||||
device=None,
|
||||
dtype=None,
|
||||
qconfig=None,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__(
|
||||
num_embeddings,
|
||||
embedding_dim,
|
||||
padding_idx,
|
||||
max_norm,
|
||||
norm_type,
|
||||
scale_grad_by_freq,
|
||||
sparse,
|
||||
_weight,
|
||||
**factory_kwargs,
|
||||
)
|
||||
if not qconfig:
|
||||
raise AssertionError("qconfig must be provided for QAT module")
|
||||
weight_qscheme = qconfig.weight().qscheme
|
||||
if weight_qscheme != torch.per_channel_affine_float_qparams:
|
||||
raise AssertionError(
|
||||
"Embedding weights requires a qscheme of torch.per_channel_affine_float_qparams Got "
|
||||
+ str(weight_qscheme)
|
||||
)
|
||||
self.qconfig = qconfig
|
||||
self.weight_fake_quant = qconfig.weight(factory_kwargs=factory_kwargs)
|
||||
|
||||
def forward(self, input) -> Tensor:
|
||||
return F.embedding(
|
||||
input,
|
||||
self.weight_fake_quant(self.weight),
|
||||
self.padding_idx,
|
||||
self.max_norm,
|
||||
self.norm_type,
|
||||
self.scale_grad_by_freq,
|
||||
self.sparse,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False):
|
||||
r"""Create a qat module from a float module
|
||||
|
||||
Args: `mod` a float module, either produced by torch.ao.quantization utilities
|
||||
or directly from user
|
||||
"""
|
||||
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")
|
||||
weight_qscheme = mod.qconfig.weight().qscheme # type: ignore[union-attr, operator]
|
||||
if weight_qscheme != torch.per_channel_affine_float_qparams:
|
||||
raise AssertionError(
|
||||
"Embedding weights requires a qscheme of torch.per_channel_affine_float_qparams Got "
|
||||
+ str(weight_qscheme)
|
||||
)
|
||||
|
||||
qconfig = mod.qconfig
|
||||
qat_embedding_bag = cls(
|
||||
mod.num_embeddings,
|
||||
mod.embedding_dim,
|
||||
mod.padding_idx,
|
||||
mod.max_norm,
|
||||
mod.norm_type,
|
||||
mod.scale_grad_by_freq,
|
||||
mod.sparse,
|
||||
mod.weight,
|
||||
qconfig=qconfig,
|
||||
)
|
||||
|
||||
return qat_embedding_bag
|
||||
|
||||
def to_float(self):
|
||||
embedding_bag = torch.nn.Embedding(
|
||||
self.num_embeddings,
|
||||
self.embedding_dim,
|
||||
self.padding_idx,
|
||||
self.max_norm,
|
||||
self.norm_type,
|
||||
self.scale_grad_by_freq,
|
||||
self.sparse,
|
||||
None,
|
||||
)
|
||||
embedding_bag.weight = torch.nn.Parameter(self.weight.detach())
|
||||
embedding_bag.train(self.training)
|
||||
return embedding_bag
|
||||
|
||||
|
||||
class EmbeddingBag(nn.EmbeddingBag):
|
||||
r"""
|
||||
An embedding bag module attached with FakeQuantize modules for weight,
|
||||
used for quantization aware training.
|
||||
|
||||
We adopt the same interface as `torch.nn.EmbeddingBag`, please see
|
||||
https://pytorch.org/docs/stable/generated/torch.nn.EmbeddingBag.html#torch.nn.EmbeddingBag
|
||||
for documentation.
|
||||
|
||||
Similar to `torch.nn.EmbeddingBag`, with FakeQuantize modules initialized to
|
||||
default.
|
||||
|
||||
Attributes:
|
||||
weight: fake quant module for weight
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE = nn.EmbeddingBag
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_embeddings,
|
||||
embedding_dim,
|
||||
max_norm=None,
|
||||
norm_type=2.0,
|
||||
scale_grad_by_freq=False,
|
||||
mode="mean",
|
||||
sparse=False,
|
||||
_weight=None,
|
||||
include_last_offset=False,
|
||||
padding_idx=None,
|
||||
qconfig=None,
|
||||
device=None,
|
||||
dtype=None,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__(
|
||||
num_embeddings,
|
||||
embedding_dim,
|
||||
max_norm,
|
||||
norm_type,
|
||||
scale_grad_by_freq,
|
||||
mode,
|
||||
sparse,
|
||||
_weight,
|
||||
include_last_offset,
|
||||
padding_idx,
|
||||
**factory_kwargs,
|
||||
)
|
||||
if not qconfig:
|
||||
raise AssertionError("qconfig must be provided for QAT module")
|
||||
weight_qscheme = qconfig.weight().qscheme
|
||||
if weight_qscheme != torch.per_channel_affine_float_qparams:
|
||||
raise AssertionError(
|
||||
"Embedding Bag weights requires a qscheme of torch.per_channel_affine_float_qparams Got "
|
||||
+ str(weight_qscheme)
|
||||
)
|
||||
self.qconfig = qconfig
|
||||
self.weight_fake_quant = qconfig.weight(factory_kwargs=factory_kwargs)
|
||||
|
||||
def forward(self, input, offsets=None, per_sample_weights=None) -> Tensor:
|
||||
return F.embedding_bag(
|
||||
input,
|
||||
self.weight_fake_quant(self.weight),
|
||||
offsets,
|
||||
self.max_norm,
|
||||
self.norm_type,
|
||||
self.scale_grad_by_freq,
|
||||
self.mode,
|
||||
self.sparse,
|
||||
per_sample_weights,
|
||||
self.include_last_offset,
|
||||
self.padding_idx,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False):
|
||||
r"""Create a qat module from a float module
|
||||
|
||||
Args: `mod` a float module, either produced by torch.ao.quantization utilities
|
||||
or directly from user
|
||||
"""
|
||||
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")
|
||||
weight_qscheme = mod.qconfig.weight().qscheme # type: ignore[union-attr, operator]
|
||||
if weight_qscheme != torch.per_channel_affine_float_qparams:
|
||||
raise AssertionError(
|
||||
"Embedding Bag weights requires a qscheme of torch.per_channel_affine_float_qparams Got "
|
||||
+ str(weight_qscheme)
|
||||
)
|
||||
|
||||
qconfig = mod.qconfig
|
||||
qat_embedding_bag = cls(
|
||||
mod.num_embeddings,
|
||||
mod.embedding_dim,
|
||||
mod.max_norm,
|
||||
mod.norm_type,
|
||||
mod.scale_grad_by_freq,
|
||||
mod.mode,
|
||||
mod.sparse,
|
||||
mod.weight,
|
||||
mod.include_last_offset,
|
||||
mod.padding_idx,
|
||||
qconfig=qconfig,
|
||||
)
|
||||
|
||||
return qat_embedding_bag
|
||||
|
||||
def to_float(self):
|
||||
embedding_bag = torch.nn.EmbeddingBag(
|
||||
self.num_embeddings,
|
||||
self.embedding_dim,
|
||||
self.max_norm,
|
||||
self.norm_type,
|
||||
self.scale_grad_by_freq,
|
||||
self.mode,
|
||||
self.sparse,
|
||||
None,
|
||||
self.include_last_offset,
|
||||
self.padding_idx,
|
||||
)
|
||||
embedding_bag.weight = torch.nn.Parameter(self.weight.detach())
|
||||
embedding_bag.train(self.training)
|
||||
return embedding_bag
|
||||
@@ -0,0 +1,99 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.ao.nn.intrinsic import LinearReLU
|
||||
from torch.nn.utils.parametrize import (
|
||||
is_parametrized,
|
||||
transfer_parametrizations_and_params,
|
||||
type_before_parametrizations,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["Linear"]
|
||||
|
||||
|
||||
class Linear(nn.Linear):
|
||||
r"""
|
||||
A linear module attached with FakeQuantize modules for weight,
|
||||
used for quantization aware training.
|
||||
|
||||
We adopt the same interface as `torch.nn.Linear`, please see
|
||||
https://pytorch.org/docs/stable/nn.html#torch.nn.Linear
|
||||
for documentation.
|
||||
|
||||
Similar to `torch.nn.Linear`, with FakeQuantize modules initialized to
|
||||
default.
|
||||
|
||||
Attributes:
|
||||
weight: fake quant module for weight
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE = nn.Linear
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_features,
|
||||
out_features,
|
||||
bias=True,
|
||||
qconfig=None,
|
||||
device=None,
|
||||
dtype=None,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__(in_features, out_features, bias, **factory_kwargs)
|
||||
if not qconfig:
|
||||
raise AssertionError("qconfig must be provided for QAT module")
|
||||
self.qconfig = qconfig
|
||||
self.weight_fake_quant = qconfig.weight(factory_kwargs=factory_kwargs)
|
||||
|
||||
def forward(self, input):
|
||||
return F.linear(input, self.weight_fake_quant(self.weight), self.bias)
|
||||
|
||||
@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
|
||||
"""
|
||||
if type_before_parametrizations(mod) != cls._FLOAT_MODULE:
|
||||
raise AssertionError(
|
||||
f"qat.{cls.__name__}.from_float only works for "
|
||||
f"{cls._FLOAT_MODULE.__name__}, got {type_before_parametrizations(mod).__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")
|
||||
if type_before_parametrizations(mod) == LinearReLU:
|
||||
mod = mod[0]
|
||||
|
||||
qconfig = mod.qconfig
|
||||
qat_linear = cls(
|
||||
mod.in_features,
|
||||
mod.out_features,
|
||||
bias=mod.bias is not None,
|
||||
qconfig=qconfig,
|
||||
)
|
||||
|
||||
if is_parametrized(mod, "weight"):
|
||||
transfer_parametrizations_and_params(mod, qat_linear, "weight")
|
||||
else:
|
||||
qat_linear.weight = mod.weight
|
||||
|
||||
if is_parametrized(mod, "bias"):
|
||||
transfer_parametrizations_and_params(mod, qat_linear, "bias")
|
||||
else:
|
||||
qat_linear.bias = mod.bias
|
||||
|
||||
return qat_linear
|
||||
|
||||
def to_float(self):
|
||||
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())
|
||||
linear.train(self.training)
|
||||
return linear
|
||||
@@ -0,0 +1 @@
|
||||
from .modules import * # noqa: F403
|
||||
@@ -0,0 +1,9 @@
|
||||
from .activation import MultiheadAttention
|
||||
from .rnn import LSTM, LSTMCell
|
||||
|
||||
|
||||
__all__ = [
|
||||
"LSTM",
|
||||
"LSTMCell",
|
||||
"MultiheadAttention",
|
||||
]
|
||||
+623
@@ -0,0 +1,623 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
import torch.jit # this is needed to avoid a circular import
|
||||
import torch.nn.functional as F
|
||||
from torch import nn, Tensor
|
||||
|
||||
|
||||
__all__ = ["MultiheadAttention"]
|
||||
|
||||
|
||||
class MultiheadAttention(nn.MultiheadAttention):
|
||||
_FLOAT_MODULE = nn.MultiheadAttention
|
||||
|
||||
r"""Quantizable implementation of the MultiheadAttention.
|
||||
|
||||
Note::
|
||||
Please, refer to :class:`~torch.nn.MultiheadAttention` for more
|
||||
information
|
||||
|
||||
Allows the model to jointly attend to information from different
|
||||
representation subspaces.
|
||||
See reference: Attention Is All You Need
|
||||
|
||||
The original MHA module is not quantizable.
|
||||
This reimplements it by explicitly instantiating the linear layers.
|
||||
|
||||
.. math::
|
||||
\text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O
|
||||
\text{where} head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)
|
||||
|
||||
Args:
|
||||
embed_dim: total dimension of the model.
|
||||
num_heads: parallel attention heads.
|
||||
dropout: a Dropout layer on attn_output_weights. Default: 0.0.
|
||||
bias: add bias as module parameter. Default: True.
|
||||
add_bias_kv: add bias to the key and value sequences at dim=0.
|
||||
add_zero_attn: add a new batch of zeros to the key and
|
||||
value sequences at dim=1.
|
||||
kdim: total number of features in key. Default: None.
|
||||
vdim: total number of features in value. Default: None.
|
||||
batch_first: If ``True``, then the input and output tensors are provided
|
||||
as (batch, seq, feature). Default: ``False`` (seq, batch, feature).
|
||||
|
||||
Note that if :attr:`kdim` and :attr:`vdim` are None, they will be set
|
||||
to :attr:`embed_dim` such that query, key, and value have the same
|
||||
number of features.
|
||||
|
||||
Examples::
|
||||
|
||||
>>> import torch.ao.nn.quantizable as nnqa
|
||||
>>> multihead_attn = nnqa.MultiheadAttention(embed_dim, num_heads)
|
||||
>>> attn_output, attn_output_weights = multihead_attn(query, key, value)
|
||||
|
||||
Note::
|
||||
Please, follow the quantization flow to convert the quantizable MHA.
|
||||
"""
|
||||
__constants__ = ["batch_first"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
embed_dim: int,
|
||||
num_heads: int,
|
||||
dropout: float = 0.0,
|
||||
bias: bool = True,
|
||||
add_bias_kv: bool = False,
|
||||
add_zero_attn: bool = False,
|
||||
kdim: int | None = None,
|
||||
vdim: int | None = None,
|
||||
batch_first: bool = False,
|
||||
device=None,
|
||||
dtype=None,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__(
|
||||
embed_dim,
|
||||
num_heads,
|
||||
dropout,
|
||||
bias,
|
||||
add_bias_kv,
|
||||
add_zero_attn,
|
||||
kdim,
|
||||
vdim,
|
||||
batch_first,
|
||||
**factory_kwargs,
|
||||
)
|
||||
self.linear_Q = nn.Linear(
|
||||
self.embed_dim, self.embed_dim, bias=bias, **factory_kwargs
|
||||
)
|
||||
self.linear_K = nn.Linear(
|
||||
self.kdim, self.embed_dim, bias=bias, **factory_kwargs
|
||||
)
|
||||
self.linear_V = nn.Linear(
|
||||
self.vdim, self.embed_dim, bias=bias, **factory_kwargs
|
||||
)
|
||||
# for the type: ignore, see https://github.com/pytorch/pytorch/issues/58969
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
self.out_proj = nn.Linear(
|
||||
self.embed_dim, self.embed_dim, bias=bias, **factory_kwargs
|
||||
) # type: ignore[assignment]
|
||||
|
||||
# Functionals
|
||||
self.q_scaling_product = torch.ao.nn.quantized.FloatFunctional()
|
||||
# note: importing torch.ao.nn.quantized at top creates a circular import
|
||||
|
||||
# Quant/Dequant
|
||||
self.quant_attn_output = torch.ao.quantization.QuantStub()
|
||||
self.quant_attn_output_weights = torch.ao.quantization.QuantStub()
|
||||
self.dequant_q = torch.ao.quantization.DeQuantStub()
|
||||
self.dequant_k = torch.ao.quantization.DeQuantStub()
|
||||
self.dequant_v = torch.ao.quantization.DeQuantStub()
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizableMultiheadAttention"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, other):
|
||||
if type(other) is not cls._FLOAT_MODULE:
|
||||
raise AssertionError(
|
||||
f"Expected type {cls._FLOAT_MODULE}, got {type(other)}"
|
||||
)
|
||||
if not hasattr(other, "qconfig"):
|
||||
raise AssertionError("The float module must have 'qconfig'")
|
||||
# Setting the dropout to 0.0!
|
||||
observed = cls(
|
||||
other.embed_dim,
|
||||
other.num_heads,
|
||||
other.dropout,
|
||||
(other.in_proj_bias is not None),
|
||||
(other.bias_k is not None),
|
||||
other.add_zero_attn,
|
||||
other.kdim,
|
||||
other.vdim,
|
||||
other.batch_first,
|
||||
)
|
||||
observed.bias_k = other.bias_k
|
||||
observed.bias_v = other.bias_v
|
||||
observed.qconfig = other.qconfig
|
||||
|
||||
# Set the linear weights
|
||||
# for the type: ignores, see https://github.com/pytorch/pytorch/issues/58969
|
||||
observed.out_proj.weight = other.out_proj.weight
|
||||
observed.out_proj.bias = other.out_proj.bias
|
||||
if other._qkv_same_embed_dim:
|
||||
# Use separate params
|
||||
bias = other.in_proj_bias
|
||||
_start = 0
|
||||
_end = _start + other.embed_dim
|
||||
weight = other.in_proj_weight[_start:_end, :]
|
||||
if bias is not None:
|
||||
bias = torch.nn.Parameter(bias[_start:_end], bias.requires_grad)
|
||||
observed.linear_Q.weight = torch.nn.Parameter(weight, weight.requires_grad)
|
||||
observed.linear_Q.bias = bias
|
||||
|
||||
bias = other.in_proj_bias
|
||||
_start = _end
|
||||
_end = _start + other.embed_dim
|
||||
weight = other.in_proj_weight[_start:_end, :]
|
||||
if bias is not None:
|
||||
bias = torch.nn.Parameter(bias[_start:_end], bias.requires_grad)
|
||||
observed.linear_K.weight = torch.nn.Parameter(weight, weight.requires_grad)
|
||||
observed.linear_K.bias = bias
|
||||
|
||||
bias = other.in_proj_bias
|
||||
_start = _end
|
||||
weight = other.in_proj_weight[_start:, :]
|
||||
if bias is not None:
|
||||
bias = torch.nn.Parameter(bias[_start:], bias.requires_grad)
|
||||
observed.linear_V.weight = torch.nn.Parameter(weight, weight.requires_grad)
|
||||
observed.linear_V.bias = bias
|
||||
else:
|
||||
observed.linear_Q.weight = nn.Parameter(other.q_proj_weight)
|
||||
observed.linear_K.weight = nn.Parameter(other.k_proj_weight)
|
||||
observed.linear_V.weight = nn.Parameter(other.v_proj_weight)
|
||||
if other.in_proj_bias is None:
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
observed.linear_Q.bias = None
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
observed.linear_K.bias = None
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
observed.linear_V.bias = None
|
||||
else:
|
||||
observed.linear_Q.bias = nn.Parameter(
|
||||
other.in_proj_bias[0 : other.embed_dim]
|
||||
)
|
||||
observed.linear_K.bias = nn.Parameter(
|
||||
other.in_proj_bias[other.embed_dim : (other.embed_dim * 2)]
|
||||
)
|
||||
observed.linear_V.bias = nn.Parameter(
|
||||
other.in_proj_bias[(other.embed_dim * 2) :]
|
||||
)
|
||||
observed.eval()
|
||||
# Explicit prepare
|
||||
observed = torch.ao.quantization.prepare(observed, inplace=True)
|
||||
return observed
|
||||
|
||||
@torch.jit.unused
|
||||
def dequantize(self):
|
||||
r"""Utility to convert the quantized MHA back to float.
|
||||
|
||||
The motivation for this is that it is not trivial to convert the weights
|
||||
from the format that is used in the quantized version back to the
|
||||
float.
|
||||
"""
|
||||
fp = self._FLOAT_MODULE(
|
||||
self.embed_dim,
|
||||
self.num_heads,
|
||||
self.dropout,
|
||||
(self.linear_Q._weight_bias()[1] is not None), # type: ignore[operator]
|
||||
(self.bias_k is not None),
|
||||
self.add_zero_attn,
|
||||
self.kdim,
|
||||
self.vdim,
|
||||
self.batch_first,
|
||||
)
|
||||
if fp._qkv_same_embed_dim != self._qkv_same_embed_dim:
|
||||
raise AssertionError(
|
||||
f"_qkv_same_embed_dim mismatch: {fp._qkv_same_embed_dim} != {self._qkv_same_embed_dim}"
|
||||
)
|
||||
if self.bias_k is not None:
|
||||
fp.bias_k = nn.Parameter(self.bias_k.dequantize())
|
||||
if self.bias_v is not None:
|
||||
fp.bias_v = nn.Parameter(self.bias_v.dequantize())
|
||||
|
||||
# Set the linear weights
|
||||
# Note: Because the linear layers are quantized, mypy does not know how
|
||||
# to deal with them -- might need to ignore the typing checks.
|
||||
# for the type: ignore[has-type], see https://github.com/pytorch/pytorch/issues/58969
|
||||
w, b = self.out_proj._weight_bias() # type: ignore[operator, has-type]
|
||||
fp.out_proj.weight = nn.Parameter(w.dequantize())
|
||||
if b is not None:
|
||||
fp.out_proj.bias = nn.Parameter(b)
|
||||
|
||||
wQ, bQ = self.linear_Q._weight_bias() # type: ignore[operator]
|
||||
wQ = wQ.dequantize()
|
||||
wK, bK = self.linear_K._weight_bias() # type: ignore[operator]
|
||||
wK = wK.dequantize()
|
||||
wV, bV = self.linear_V._weight_bias() # type: ignore[operator]
|
||||
wV = wV.dequantize()
|
||||
if fp._qkv_same_embed_dim:
|
||||
# Use separate params
|
||||
_start = 0
|
||||
_end = _start + fp.embed_dim
|
||||
fp.in_proj_weight[_start:_end, :] = wQ
|
||||
if fp.in_proj_bias is not None:
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
if not all(bQ == 0):
|
||||
raise AssertionError("Expected all bQ elements to be 0")
|
||||
fp.in_proj_bias[_start:_end] = bQ
|
||||
|
||||
_start = _end
|
||||
_end = _start + fp.embed_dim
|
||||
fp.in_proj_weight[_start:_end, :] = wK
|
||||
if fp.in_proj_bias is not None:
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
if not all(bK == 0):
|
||||
raise AssertionError("Expected all bK elements to be 0")
|
||||
fp.in_proj_bias[_start:_end] = bK
|
||||
|
||||
_start = _end
|
||||
fp.in_proj_weight[_start:, :] = wV
|
||||
if fp.in_proj_bias is not None:
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
if not all(bV == 0):
|
||||
raise AssertionError("Expected all bV elements to be 0")
|
||||
fp.in_proj_bias[_start:] = bV
|
||||
else:
|
||||
fp.q_proj_weight = nn.Parameter(wQ)
|
||||
fp.k_proj_weight = nn.Parameter(wK)
|
||||
fp.v_proj_weight = nn.Parameter(wV)
|
||||
if fp.in_proj_bias is None:
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
self.linear_Q.bias = None
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
self.linear_K.bias = None
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
self.linear_V.bias = None
|
||||
else:
|
||||
fp.in_proj_bias[0 : fp.embed_dim] = bQ
|
||||
fp.in_proj_bias[fp.embed_dim : (fp.embed_dim * 2)] = bK
|
||||
fp.in_proj_bias[(fp.embed_dim * 2) :] = bV
|
||||
|
||||
return fp
|
||||
|
||||
@classmethod
|
||||
def from_observed(cls, other):
|
||||
# The whole flow is float -> observed -> quantized
|
||||
# This class does float -> observed only
|
||||
# See nn.quantized.MultiheadAttention
|
||||
raise NotImplementedError(
|
||||
"It looks like you are trying to prepare an "
|
||||
"MHA module. Please, see "
|
||||
"the examples on quantizable MHAs."
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
query: Tensor,
|
||||
key: Tensor,
|
||||
value: Tensor,
|
||||
key_padding_mask: Tensor | None = None,
|
||||
need_weights: bool = True,
|
||||
attn_mask: Tensor | None = None,
|
||||
average_attn_weights: bool = True,
|
||||
is_causal: bool = False,
|
||||
) -> tuple[Tensor, Tensor | None]:
|
||||
r"""
|
||||
Note::
|
||||
Please, refer to :func:`~torch.nn.MultiheadAttention.forward` for more
|
||||
information
|
||||
|
||||
Args:
|
||||
query, key, value: map a query and a set of key-value pairs to an output.
|
||||
See "Attention Is All You Need" for more details.
|
||||
key_padding_mask: if provided, specified padding elements in the key will
|
||||
be ignored by the attention. When given a binary mask and a value is True,
|
||||
the corresponding value on the attention layer will be ignored.
|
||||
need_weights: output attn_output_weights.
|
||||
attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
|
||||
the batches while a 3D mask allows to specify a different mask for the entries of each batch.
|
||||
|
||||
Shape:
|
||||
- Inputs:
|
||||
- query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
|
||||
the embedding dimension. :math:`(N, L, E)` if ``batch_first`` is ``True``.
|
||||
- key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
|
||||
the embedding dimension. :math:`(N, S, E)` if ``batch_first`` is ``True``.
|
||||
- value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
|
||||
the embedding dimension. :math:`(N, S, E)` if ``batch_first`` is ``True``.
|
||||
- key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
|
||||
If a BoolTensor is provided, the positions with the
|
||||
value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
|
||||
- attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
|
||||
3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
|
||||
S is the source sequence length. attn_mask ensure that position i is allowed to attend the unmasked
|
||||
positions. If a BoolTensor is provided, positions with ``True``
|
||||
is not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
|
||||
is provided, it will be added to the attention weight.
|
||||
- is_causal: If specified, applies a causal mask as attention mask. Mutually exclusive with providing attn_mask.
|
||||
Default: ``False``.
|
||||
- average_attn_weights: If true, indicates that the returned ``attn_weights`` should be averaged across
|
||||
heads. Otherwise, ``attn_weights`` are provided separately per head. Note that this flag only has an
|
||||
effect when ``need_weights=True.``. Default: True (i.e. average weights across heads)
|
||||
|
||||
- Outputs:
|
||||
- attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
|
||||
E is the embedding dimension. :math:`(N, L, E)` if ``batch_first`` is ``True``.
|
||||
- attn_output_weights: If ``average_attn_weights=True``, returns attention weights averaged
|
||||
across heads of shape :math:`(N, L, S)`, where N is the batch size, L is the target sequence length,
|
||||
S is the source sequence length. If ``average_attn_weights=False``, returns attention weights per
|
||||
head of shape :math:`(N, num_heads, L, S)`.
|
||||
"""
|
||||
return self._forward_impl(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
key_padding_mask,
|
||||
need_weights,
|
||||
attn_mask,
|
||||
average_attn_weights,
|
||||
is_causal,
|
||||
)
|
||||
|
||||
def _forward_impl(
|
||||
self,
|
||||
query: Tensor,
|
||||
key: Tensor,
|
||||
value: Tensor,
|
||||
key_padding_mask: Tensor | None = None,
|
||||
need_weights: bool = True,
|
||||
attn_mask: Tensor | None = None,
|
||||
average_attn_weights: bool = True,
|
||||
is_causal: bool = False,
|
||||
) -> tuple[Tensor, Tensor | None]:
|
||||
# This version will not deal with the static key/value pairs.
|
||||
# Keeping it here for future changes.
|
||||
#
|
||||
# TODO: This method has some duplicate lines with the
|
||||
# `torch.nn.functional.multi_head_attention`. Will need to refactor.
|
||||
static_k = None
|
||||
static_v = None
|
||||
|
||||
if attn_mask is not None and is_causal:
|
||||
raise AssertionError("Only allow causal mask or attn_mask")
|
||||
|
||||
if is_causal:
|
||||
raise AssertionError("causal mask not supported by AO MHA module")
|
||||
|
||||
if self.batch_first:
|
||||
query, key, value = (x.transpose(0, 1) for x in (query, key, value))
|
||||
|
||||
tgt_len, bsz, embed_dim_to_check = query.size()
|
||||
if self.embed_dim != embed_dim_to_check:
|
||||
raise AssertionError(
|
||||
f"embed_dim mismatch: {self.embed_dim} != {embed_dim_to_check}"
|
||||
)
|
||||
# allow MHA to have different sizes for the feature dimension
|
||||
if key.size(0) != value.size(0) or key.size(1) != value.size(1):
|
||||
raise AssertionError(
|
||||
f"key and value size mismatch: key.size()={key.size()}, value.size()={value.size()}"
|
||||
)
|
||||
|
||||
head_dim = self.embed_dim // self.num_heads
|
||||
if head_dim * self.num_heads != self.embed_dim:
|
||||
raise AssertionError("embed_dim must be divisible by num_heads")
|
||||
scaling = float(head_dim) ** -0.5
|
||||
|
||||
q = self.linear_Q(query)
|
||||
k = self.linear_K(key)
|
||||
v = self.linear_V(value)
|
||||
|
||||
q = self.q_scaling_product.mul_scalar(q, scaling)
|
||||
|
||||
if attn_mask is not None:
|
||||
if attn_mask.dtype == torch.uint8:
|
||||
warnings.warn(
|
||||
"Byte tensor for `attn_mask` in `nn.MultiheadAttention` is deprecated. "
|
||||
"Use bool tensor instead.",
|
||||
stacklevel=3,
|
||||
)
|
||||
attn_mask = attn_mask.to(torch.bool)
|
||||
if not attn_mask.is_floating_point() and attn_mask.dtype != torch.bool:
|
||||
raise AssertionError(
|
||||
f"Only float and bool types are supported for attn_mask, not {attn_mask.dtype}"
|
||||
)
|
||||
|
||||
if attn_mask.dim() == 2:
|
||||
attn_mask = attn_mask.unsqueeze(0)
|
||||
if list(attn_mask.size()) != [1, query.size(0), key.size(0)]:
|
||||
raise RuntimeError("The size of the 2D attn_mask is not correct.")
|
||||
elif attn_mask.dim() == 3:
|
||||
if list(attn_mask.size()) != [
|
||||
bsz * self.num_heads,
|
||||
query.size(0),
|
||||
key.size(0),
|
||||
]:
|
||||
raise RuntimeError("The size of the 3D attn_mask is not correct.")
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"attn_mask's dimension {attn_mask.dim()} is not supported"
|
||||
)
|
||||
# attn_mask's dim is 3 now.
|
||||
|
||||
# convert ByteTensor key_padding_mask to bool
|
||||
if key_padding_mask is not None and key_padding_mask.dtype == torch.uint8:
|
||||
warnings.warn(
|
||||
"Byte tensor for `key_padding_mask` in `nn.MultiheadAttention` is deprecated. "
|
||||
"Use bool tensor instead.",
|
||||
stacklevel=3,
|
||||
)
|
||||
key_padding_mask = key_padding_mask.to(torch.bool)
|
||||
if self.bias_k is not None and self.bias_v is not None:
|
||||
if static_k is None and static_v is None:
|
||||
# Explicitly check that bias_k and bias_v are not None
|
||||
# in a way that TorchScript can understand.
|
||||
bias_k = self.bias_k
|
||||
if bias_k is None:
|
||||
raise AssertionError("bias_k must not be None")
|
||||
bias_v = self.bias_v
|
||||
if bias_v is None:
|
||||
raise AssertionError("bias_v must not be None")
|
||||
|
||||
k = torch.cat([k, bias_k.repeat(1, bsz, 1)])
|
||||
v = torch.cat([v, bias_v.repeat(1, bsz, 1)])
|
||||
if attn_mask is not None:
|
||||
attn_mask = F.pad(attn_mask, (0, 1))
|
||||
if key_padding_mask is not None:
|
||||
key_padding_mask = F.pad(key_padding_mask, (0, 1))
|
||||
else:
|
||||
if static_k is not None:
|
||||
raise AssertionError("bias cannot be added to static key.")
|
||||
if static_v is not None:
|
||||
raise AssertionError("bias cannot be added to static value.")
|
||||
else:
|
||||
if self.bias_k is not None:
|
||||
raise AssertionError(
|
||||
"self.bias_k must be None when self.bias_v is None"
|
||||
)
|
||||
if self.bias_v is not None:
|
||||
raise AssertionError(
|
||||
"self.bias_v must be None when self.bias_k is None"
|
||||
)
|
||||
|
||||
q = q.contiguous().view(tgt_len, bsz * self.num_heads, head_dim).transpose(0, 1)
|
||||
if k is not None:
|
||||
k = k.contiguous().view(-1, bsz * self.num_heads, head_dim).transpose(0, 1)
|
||||
if v is not None:
|
||||
v = v.contiguous().view(-1, bsz * self.num_heads, head_dim).transpose(0, 1)
|
||||
|
||||
if static_k is not None:
|
||||
if static_k.size(0) != bsz * self.num_heads:
|
||||
raise AssertionError(
|
||||
f"static_k.size(0) must be {bsz * self.num_heads}, got {static_k.size(0)}"
|
||||
)
|
||||
if static_k.size(2) != head_dim:
|
||||
raise AssertionError(
|
||||
f"static_k.size(2) must be {head_dim}, got {static_k.size(2)}"
|
||||
)
|
||||
k = static_k
|
||||
|
||||
if static_v is not None:
|
||||
if static_v.size(0) != bsz * self.num_heads:
|
||||
raise AssertionError(
|
||||
f"static_v.size(0) must be {bsz * self.num_heads}, got {static_v.size(0)}"
|
||||
)
|
||||
if static_v.size(2) != head_dim:
|
||||
raise AssertionError(
|
||||
f"static_v.size(2) must be {head_dim}, got {static_v.size(2)}"
|
||||
)
|
||||
v = static_v
|
||||
|
||||
src_len = k.size(1)
|
||||
|
||||
if key_padding_mask is not None:
|
||||
if key_padding_mask.size(0) != bsz:
|
||||
raise AssertionError(
|
||||
f"key_padding_mask.size(0) must be {bsz}, got {key_padding_mask.size(0)}"
|
||||
)
|
||||
if key_padding_mask.size(1) != src_len:
|
||||
raise AssertionError(
|
||||
f"key_padding_mask.size(1) must be {src_len}, got {key_padding_mask.size(1)}"
|
||||
)
|
||||
|
||||
if self.add_zero_attn:
|
||||
src_len += 1
|
||||
|
||||
k_zeros = torch.zeros((k.size(0), 1) + k.size()[2:])
|
||||
|
||||
if k.is_quantized:
|
||||
k_zeros = torch.quantize_per_tensor(
|
||||
k_zeros,
|
||||
k.q_scale(),
|
||||
k.q_zero_point(),
|
||||
k.dtype,
|
||||
)
|
||||
|
||||
k = torch.cat([k, k_zeros], dim=1)
|
||||
|
||||
v_zeros = torch.zeros((v.size(0), 1) + k.size()[2:])
|
||||
|
||||
if v.is_quantized:
|
||||
v_zeros = torch.quantize_per_tensor(
|
||||
v_zeros,
|
||||
v.q_scale(),
|
||||
v.q_zero_point(),
|
||||
v.dtype,
|
||||
)
|
||||
|
||||
v = torch.cat([v, v_zeros], dim=1)
|
||||
|
||||
if attn_mask is not None:
|
||||
attn_mask = F.pad(attn_mask, (0, 1))
|
||||
if key_padding_mask is not None:
|
||||
key_padding_mask = F.pad(key_padding_mask, (0, 1))
|
||||
|
||||
# Leaving the quantized zone here
|
||||
q = self.dequant_q(q)
|
||||
k = self.dequant_k(k)
|
||||
v = self.dequant_v(v)
|
||||
attn_output_weights = torch.bmm(q, k.transpose(1, 2))
|
||||
expected_size = [bsz * self.num_heads, tgt_len, src_len]
|
||||
if list(attn_output_weights.size()) != expected_size:
|
||||
raise AssertionError(
|
||||
f"attn_output_weights size mismatch: expected {expected_size}, "
|
||||
f"got {list(attn_output_weights.size())}"
|
||||
)
|
||||
|
||||
if attn_mask is not None:
|
||||
if attn_mask.dtype == torch.bool:
|
||||
attn_output_weights.masked_fill_(attn_mask, float("-inf"))
|
||||
else:
|
||||
attn_output_weights += attn_mask
|
||||
|
||||
if key_padding_mask is not None:
|
||||
attn_output_weights = attn_output_weights.view(
|
||||
bsz, self.num_heads, tgt_len, src_len
|
||||
)
|
||||
attn_output_weights = attn_output_weights.masked_fill(
|
||||
key_padding_mask.unsqueeze(1).unsqueeze(2),
|
||||
float("-inf"),
|
||||
)
|
||||
attn_output_weights = attn_output_weights.view(
|
||||
bsz * self.num_heads, tgt_len, src_len
|
||||
)
|
||||
|
||||
attn_output_weights = F.softmax(attn_output_weights, dim=-1)
|
||||
attn_output_weights = F.dropout(
|
||||
attn_output_weights, p=self.dropout, training=self.training
|
||||
)
|
||||
|
||||
attn_output = torch.bmm(attn_output_weights, v)
|
||||
expected_output_size = [bsz * self.num_heads, tgt_len, head_dim]
|
||||
if list(attn_output.size()) != expected_output_size:
|
||||
raise AssertionError(
|
||||
f"attn_output size mismatch: expected {expected_output_size}, "
|
||||
f"got {list(attn_output.size())}"
|
||||
)
|
||||
if self.batch_first:
|
||||
attn_output = attn_output.view(bsz, tgt_len, self.embed_dim)
|
||||
else:
|
||||
attn_output = (
|
||||
attn_output.transpose(0, 1)
|
||||
.contiguous()
|
||||
.view(tgt_len, bsz, self.embed_dim)
|
||||
)
|
||||
|
||||
# Reentering the quantized zone
|
||||
attn_output = self.quant_attn_output(attn_output)
|
||||
# for the type: ignore[has-type], see https://github.com/pytorch/pytorch/issues/58969
|
||||
attn_output = self.out_proj(attn_output) # type: ignore[has-type]
|
||||
attn_output_weights = self.quant_attn_output_weights(attn_output_weights)
|
||||
|
||||
if need_weights:
|
||||
# average attention weights over heads
|
||||
attn_output_weights = attn_output_weights.view(
|
||||
bsz, self.num_heads, tgt_len, src_len
|
||||
)
|
||||
if average_attn_weights:
|
||||
attn_output_weights = attn_output_weights.mean(dim=1)
|
||||
return attn_output, attn_output_weights
|
||||
else:
|
||||
return attn_output, None
|
||||
@@ -0,0 +1,618 @@
|
||||
"""
|
||||
We will recreate all the RNN modules as we require the modules to be decomposed
|
||||
into its building blocks to be able to observe.
|
||||
"""
|
||||
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import numbers
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
__all__ = ["LSTMCell", "LSTM"]
|
||||
|
||||
|
||||
class LSTMCell(torch.nn.Module):
|
||||
r"""A quantizable long short-term memory (LSTM) cell.
|
||||
|
||||
For the description and the argument types, please, refer to :class:`~torch.nn.LSTMCell`
|
||||
|
||||
`split_gates`: specify True to compute the input/forget/cell/output gates separately
|
||||
to avoid an intermediate tensor which is subsequently chunk'd. This optimization can
|
||||
be beneficial for on-device inference latency. This flag is cascaded down from the
|
||||
parent classes.
|
||||
|
||||
Examples::
|
||||
|
||||
>>> import torch.ao.nn.quantizable as nnqa
|
||||
>>> rnn = nnqa.LSTMCell(10, 20)
|
||||
>>> input = torch.randn(6, 10)
|
||||
>>> hx = torch.randn(3, 20)
|
||||
>>> cx = torch.randn(3, 20)
|
||||
>>> output = []
|
||||
>>> for i in range(6):
|
||||
... hx, cx = rnn(input[i], (hx, cx))
|
||||
... output.append(hx)
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE = torch.nn.LSTMCell
|
||||
__constants__ = ["split_gates"] # for jit.script
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_dim: int,
|
||||
hidden_dim: int,
|
||||
bias: bool = True,
|
||||
device=None,
|
||||
dtype=None,
|
||||
*,
|
||||
split_gates=False,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
self.input_size = input_dim
|
||||
self.hidden_size = hidden_dim
|
||||
self.bias = bias
|
||||
self.split_gates = split_gates
|
||||
|
||||
if not split_gates:
|
||||
self.igates: torch.nn.Module = torch.nn.Linear(
|
||||
input_dim, 4 * hidden_dim, bias=bias, **factory_kwargs
|
||||
)
|
||||
self.hgates: torch.nn.Module = torch.nn.Linear(
|
||||
hidden_dim, 4 * hidden_dim, bias=bias, **factory_kwargs
|
||||
)
|
||||
self.gates: torch.nn.Module = torch.ao.nn.quantized.FloatFunctional()
|
||||
else:
|
||||
# keep separate Linear layers for each gate
|
||||
self.igates = torch.nn.ModuleDict()
|
||||
self.hgates = torch.nn.ModuleDict()
|
||||
self.gates = torch.nn.ModuleDict()
|
||||
for g in ["input", "forget", "cell", "output"]:
|
||||
# pyre-fixme[29]: `Union[torch._tensor.Tensor, torch.nn.modules.module.Module]`
|
||||
self.igates[g] = torch.nn.Linear(
|
||||
input_dim, hidden_dim, bias=bias, **factory_kwargs
|
||||
)
|
||||
# pyre-fixme[29]: `Union[torch._tensor.Tensor, torch.nn.modules.module.Module]`
|
||||
self.hgates[g] = torch.nn.Linear(
|
||||
hidden_dim, hidden_dim, bias=bias, **factory_kwargs
|
||||
)
|
||||
# pyre-fixme[29]: `Union[torch._tensor.Tensor, torch.nn.modules.module.Module]`
|
||||
self.gates[g] = torch.ao.nn.quantized.FloatFunctional()
|
||||
|
||||
self.input_gate = torch.nn.Sigmoid()
|
||||
self.forget_gate = torch.nn.Sigmoid()
|
||||
self.cell_gate = torch.nn.Tanh()
|
||||
self.output_gate = torch.nn.Sigmoid()
|
||||
|
||||
self.fgate_cx = torch.ao.nn.quantized.FloatFunctional()
|
||||
self.igate_cgate = torch.ao.nn.quantized.FloatFunctional()
|
||||
self.fgate_cx_igate_cgate = torch.ao.nn.quantized.FloatFunctional()
|
||||
|
||||
self.ogate_cy = torch.ao.nn.quantized.FloatFunctional()
|
||||
|
||||
self.initial_hidden_state_qparams: tuple[float, int] = (1.0, 0)
|
||||
self.initial_cell_state_qparams: tuple[float, int] = (1.0, 0)
|
||||
self.hidden_state_dtype: torch.dtype = torch.quint8
|
||||
self.cell_state_dtype: torch.dtype = torch.quint8
|
||||
|
||||
def forward(
|
||||
self, x: Tensor, hidden: tuple[Tensor, Tensor] | None = None
|
||||
) -> tuple[Tensor, Tensor]:
|
||||
if hidden is None or hidden[0] is None or hidden[1] is None:
|
||||
hidden = self.initialize_hidden(x.shape[0], x.is_quantized)
|
||||
hx, cx = hidden
|
||||
|
||||
if not self.split_gates:
|
||||
igates = self.igates(x)
|
||||
hgates = self.hgates(hx)
|
||||
gates = self.gates.add(igates, hgates) # type: ignore[operator]
|
||||
|
||||
input_gate, forget_gate, cell_gate, out_gate = gates.chunk(4, 1)
|
||||
|
||||
input_gate = self.input_gate(input_gate)
|
||||
forget_gate = self.forget_gate(forget_gate)
|
||||
cell_gate = self.cell_gate(cell_gate)
|
||||
out_gate = self.output_gate(out_gate)
|
||||
else:
|
||||
# apply each input + hidden projection and add together
|
||||
gate = {}
|
||||
for (key, gates), igates, hgates in zip(
|
||||
self.gates.items(), # type: ignore[operator]
|
||||
self.igates.values(), # type: ignore[operator]
|
||||
self.hgates.values(), # type: ignore[operator]
|
||||
):
|
||||
gate[key] = gates.add(igates(x), hgates(hx))
|
||||
|
||||
input_gate = self.input_gate(gate["input"])
|
||||
forget_gate = self.forget_gate(gate["forget"])
|
||||
cell_gate = self.cell_gate(gate["cell"])
|
||||
out_gate = self.output_gate(gate["output"])
|
||||
|
||||
fgate_cx = self.fgate_cx.mul(forget_gate, cx)
|
||||
igate_cgate = self.igate_cgate.mul(input_gate, cell_gate)
|
||||
fgate_cx_igate_cgate = self.fgate_cx_igate_cgate.add(fgate_cx, igate_cgate)
|
||||
cy = fgate_cx_igate_cgate
|
||||
|
||||
# TODO: make this tanh a member of the module so its qparams can be configured
|
||||
tanh_cy = torch.tanh(cy)
|
||||
hy = self.ogate_cy.mul(out_gate, tanh_cy)
|
||||
return hy, cy
|
||||
|
||||
def initialize_hidden(
|
||||
self, batch_size: int, is_quantized: bool = False
|
||||
) -> tuple[Tensor, Tensor]:
|
||||
h, c = (
|
||||
torch.zeros((batch_size, self.hidden_size)),
|
||||
torch.zeros((batch_size, self.hidden_size)),
|
||||
)
|
||||
if is_quantized:
|
||||
(h_scale, h_zp) = self.initial_hidden_state_qparams
|
||||
(c_scale, c_zp) = self.initial_cell_state_qparams
|
||||
h = torch.quantize_per_tensor(
|
||||
h, scale=h_scale, zero_point=h_zp, dtype=self.hidden_state_dtype
|
||||
)
|
||||
c = torch.quantize_per_tensor(
|
||||
c, scale=c_scale, zero_point=c_zp, dtype=self.cell_state_dtype
|
||||
)
|
||||
return h, c
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizableLSTMCell"
|
||||
|
||||
@classmethod
|
||||
def from_params(cls, wi, wh, bi=None, bh=None, split_gates=False):
|
||||
"""Uses the weights and biases to create a new LSTM cell.
|
||||
|
||||
Args:
|
||||
wi, wh: Weights for the input and hidden layers
|
||||
bi, bh: Biases for the input and hidden layers
|
||||
"""
|
||||
if (bi is None) != (bh is None):
|
||||
raise AssertionError("bi and bh must both be None or both have values")
|
||||
input_size = wi.shape[1]
|
||||
hidden_size = wh.shape[1]
|
||||
cell = cls(
|
||||
input_dim=input_size,
|
||||
hidden_dim=hidden_size,
|
||||
bias=(bi is not None),
|
||||
split_gates=split_gates,
|
||||
)
|
||||
|
||||
if not split_gates:
|
||||
cell.igates.weight = torch.nn.Parameter(wi)
|
||||
if bi is not None:
|
||||
cell.igates.bias = torch.nn.Parameter(bi)
|
||||
cell.hgates.weight = torch.nn.Parameter(wh)
|
||||
if bh is not None:
|
||||
cell.hgates.bias = torch.nn.Parameter(bh)
|
||||
else:
|
||||
# split weight/bias
|
||||
for w, b, gates in zip([wi, wh], [bi, bh], [cell.igates, cell.hgates]):
|
||||
for w_chunk, gate in zip(w.chunk(4, dim=0), gates.values()): # type: ignore[operator]
|
||||
gate.weight = torch.nn.Parameter(w_chunk)
|
||||
|
||||
if b is not None:
|
||||
for b_chunk, gate in zip(b.chunk(4, dim=0), gates.values()): # type: ignore[operator]
|
||||
gate.bias = torch.nn.Parameter(b_chunk)
|
||||
|
||||
return cell
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, other, use_precomputed_fake_quant=False, split_gates=False):
|
||||
if type(other) is not cls._FLOAT_MODULE:
|
||||
raise AssertionError(
|
||||
f"Expected module type {cls._FLOAT_MODULE}, got {type(other)}"
|
||||
)
|
||||
if not hasattr(other, "qconfig"):
|
||||
raise AssertionError("The float module must have 'qconfig'")
|
||||
observed = cls.from_params(
|
||||
other.weight_ih,
|
||||
other.weight_hh,
|
||||
other.bias_ih,
|
||||
other.bias_hh,
|
||||
split_gates=split_gates,
|
||||
)
|
||||
observed.qconfig = other.qconfig
|
||||
observed.igates.qconfig = other.qconfig
|
||||
observed.hgates.qconfig = other.qconfig
|
||||
if split_gates:
|
||||
# also apply qconfig directly to Linear modules
|
||||
for g in observed.igates.values():
|
||||
g.qconfig = other.qconfig
|
||||
for g in observed.hgates.values():
|
||||
g.qconfig = other.qconfig
|
||||
return observed
|
||||
|
||||
|
||||
class _LSTMSingleLayer(torch.nn.Module):
|
||||
r"""A single one-directional LSTM layer.
|
||||
|
||||
The difference between a layer and a cell is that the layer can process a
|
||||
sequence, while the cell only expects an instantaneous value.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_dim: int,
|
||||
hidden_dim: int,
|
||||
bias: bool = True,
|
||||
device=None,
|
||||
dtype=None,
|
||||
*,
|
||||
split_gates=False,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
self.cell = LSTMCell(
|
||||
input_dim, hidden_dim, bias=bias, split_gates=split_gates, **factory_kwargs
|
||||
)
|
||||
|
||||
def forward(self, x: Tensor, hidden: tuple[Tensor, Tensor] | None = None):
|
||||
result = []
|
||||
seq_len = x.shape[0]
|
||||
for i in range(seq_len):
|
||||
hidden = self.cell(x[i], hidden)
|
||||
result.append(hidden[0]) # type: ignore[index]
|
||||
result_tensor = torch.stack(result, 0)
|
||||
return result_tensor, hidden
|
||||
|
||||
@classmethod
|
||||
def from_params(cls, *args, **kwargs):
|
||||
cell = LSTMCell.from_params(*args, **kwargs)
|
||||
layer = cls(
|
||||
cell.input_size, cell.hidden_size, cell.bias, split_gates=cell.split_gates
|
||||
)
|
||||
layer.cell = cell
|
||||
return layer
|
||||
|
||||
|
||||
class _LSTMLayer(torch.nn.Module):
|
||||
r"""A single bi-directional LSTM layer."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_dim: int,
|
||||
hidden_dim: int,
|
||||
bias: bool = True,
|
||||
batch_first: bool = False,
|
||||
bidirectional: bool = False,
|
||||
device=None,
|
||||
dtype=None,
|
||||
*,
|
||||
split_gates=False,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
self.batch_first = batch_first
|
||||
self.bidirectional = bidirectional
|
||||
self.layer_fw = _LSTMSingleLayer(
|
||||
input_dim, hidden_dim, bias=bias, split_gates=split_gates, **factory_kwargs
|
||||
)
|
||||
if self.bidirectional:
|
||||
self.layer_bw = _LSTMSingleLayer(
|
||||
input_dim,
|
||||
hidden_dim,
|
||||
bias=bias,
|
||||
split_gates=split_gates,
|
||||
**factory_kwargs,
|
||||
)
|
||||
|
||||
def forward(self, x: Tensor, hidden: tuple[Tensor, Tensor] | None = None):
|
||||
if self.batch_first:
|
||||
x = x.transpose(0, 1)
|
||||
if hidden is None:
|
||||
hx_fw, cx_fw = (None, None)
|
||||
else:
|
||||
hx_fw, cx_fw = hidden
|
||||
hidden_bw: tuple[Tensor, Tensor] | None = None
|
||||
if self.bidirectional:
|
||||
if hx_fw is None:
|
||||
hx_bw = None
|
||||
else:
|
||||
hx_bw = hx_fw[1]
|
||||
hx_fw = hx_fw[0]
|
||||
if cx_fw is None:
|
||||
cx_bw = None
|
||||
else:
|
||||
cx_bw = cx_fw[1]
|
||||
cx_fw = cx_fw[0]
|
||||
if hx_bw is not None and cx_bw is not None:
|
||||
hidden_bw = hx_bw, cx_bw
|
||||
if hx_fw is None and cx_fw is None:
|
||||
hidden_fw = None
|
||||
else:
|
||||
hidden_fw = (
|
||||
torch.jit._unwrap_optional(hx_fw),
|
||||
torch.jit._unwrap_optional(cx_fw),
|
||||
)
|
||||
result_fw, hidden_fw = self.layer_fw(x, hidden_fw)
|
||||
|
||||
if hasattr(self, "layer_bw") and self.bidirectional:
|
||||
x_reversed = x.flip(0)
|
||||
result_bw, hidden_bw = self.layer_bw(x_reversed, hidden_bw)
|
||||
result_bw = result_bw.flip(0)
|
||||
|
||||
result = torch.cat([result_fw, result_bw], result_fw.dim() - 1)
|
||||
if hidden_fw is None and hidden_bw is None:
|
||||
h = None
|
||||
c = None
|
||||
elif hidden_fw is None:
|
||||
(h, c) = torch.jit._unwrap_optional(hidden_bw)
|
||||
elif hidden_bw is None:
|
||||
(h, c) = torch.jit._unwrap_optional(hidden_fw)
|
||||
else:
|
||||
h = torch.stack([hidden_fw[0], hidden_bw[0]], 0) # type: ignore[list-item]
|
||||
c = torch.stack([hidden_fw[1], hidden_bw[1]], 0) # type: ignore[list-item]
|
||||
else:
|
||||
result = result_fw
|
||||
h, c = torch.jit._unwrap_optional(hidden_fw) # type: ignore[assignment]
|
||||
|
||||
if self.batch_first:
|
||||
result.transpose_(0, 1)
|
||||
|
||||
return result, (h, c)
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, other, layer_idx=0, qconfig=None, **kwargs):
|
||||
r"""
|
||||
There is no FP equivalent of this class. This function is here just to
|
||||
mimic the behavior of the `prepare` within the `torch.ao.quantization`
|
||||
flow.
|
||||
"""
|
||||
if not hasattr(other, "qconfig") and qconfig is None:
|
||||
raise AssertionError("other must have qconfig or qconfig must be provided")
|
||||
|
||||
input_size = kwargs.get("input_size", other.input_size)
|
||||
hidden_size = kwargs.get("hidden_size", other.hidden_size)
|
||||
bias = kwargs.get("bias", other.bias)
|
||||
batch_first = kwargs.get("batch_first", other.batch_first)
|
||||
bidirectional = kwargs.get("bidirectional", other.bidirectional)
|
||||
split_gates = kwargs.get("split_gates", False)
|
||||
|
||||
layer = cls(
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
input_size,
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
hidden_size,
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
bias,
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
batch_first,
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
bidirectional,
|
||||
split_gates=split_gates,
|
||||
)
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
layer.qconfig = getattr(other, "qconfig", qconfig)
|
||||
wi = getattr(other, f"weight_ih_l{layer_idx}")
|
||||
wh = getattr(other, f"weight_hh_l{layer_idx}")
|
||||
bi = getattr(other, f"bias_ih_l{layer_idx}", None)
|
||||
bh = getattr(other, f"bias_hh_l{layer_idx}", None)
|
||||
|
||||
layer.layer_fw = _LSTMSingleLayer.from_params(
|
||||
wi, wh, bi, bh, split_gates=split_gates
|
||||
)
|
||||
|
||||
if other.bidirectional:
|
||||
wi = getattr(other, f"weight_ih_l{layer_idx}_reverse")
|
||||
wh = getattr(other, f"weight_hh_l{layer_idx}_reverse")
|
||||
bi = getattr(other, f"bias_ih_l{layer_idx}_reverse", None)
|
||||
bh = getattr(other, f"bias_hh_l{layer_idx}_reverse", None)
|
||||
layer.layer_bw = _LSTMSingleLayer.from_params(
|
||||
wi, wh, bi, bh, split_gates=split_gates
|
||||
)
|
||||
return layer
|
||||
|
||||
|
||||
class LSTM(torch.nn.Module):
|
||||
r"""A quantizable long short-term memory (LSTM).
|
||||
|
||||
For the description and the argument types, please, refer to :class:`~torch.nn.LSTM`
|
||||
|
||||
Attributes:
|
||||
layers : instances of the `_LSTMLayer`
|
||||
|
||||
.. note::
|
||||
To access the weights and biases, you need to access them per layer.
|
||||
See examples below.
|
||||
|
||||
Examples::
|
||||
|
||||
>>> import torch.ao.nn.quantizable as nnqa
|
||||
>>> rnn = nnqa.LSTM(10, 20, 2)
|
||||
>>> input = torch.randn(5, 3, 10)
|
||||
>>> h0 = torch.randn(2, 3, 20)
|
||||
>>> c0 = torch.randn(2, 3, 20)
|
||||
>>> output, (hn, cn) = rnn(input, (h0, c0))
|
||||
>>> # To get the weights:
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> print(rnn.layers[0].weight_ih)
|
||||
tensor([[...]])
|
||||
>>> print(rnn.layers[0].weight_hh)
|
||||
AssertionError: There is no reverse path in the non-bidirectional layer
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE = torch.nn.LSTM
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_size: int,
|
||||
hidden_size: int,
|
||||
num_layers: int = 1,
|
||||
bias: bool = True,
|
||||
batch_first: bool = False,
|
||||
dropout: float = 0.0,
|
||||
bidirectional: bool = False,
|
||||
device=None,
|
||||
dtype=None,
|
||||
*,
|
||||
split_gates: bool = False,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
self.input_size = input_size
|
||||
self.hidden_size = hidden_size
|
||||
self.num_layers = num_layers
|
||||
self.bias = bias
|
||||
self.batch_first = batch_first
|
||||
self.dropout = float(dropout)
|
||||
self.bidirectional = bidirectional
|
||||
self.training = False # Default to eval mode. If we want to train, we will explicitly set to training.
|
||||
|
||||
if (
|
||||
not isinstance(dropout, numbers.Number)
|
||||
or not 0 <= dropout <= 1
|
||||
or isinstance(dropout, bool)
|
||||
):
|
||||
raise ValueError(
|
||||
"dropout should be a number in range [0, 1] "
|
||||
"representing the probability of an element being "
|
||||
"zeroed"
|
||||
)
|
||||
|
||||
if dropout > 0:
|
||||
warnings.warn(
|
||||
"dropout option for quantizable LSTM is ignored. "
|
||||
"If you are training, please, use nn.LSTM version "
|
||||
"followed by `prepare` step.",
|
||||
stacklevel=2,
|
||||
)
|
||||
if num_layers == 1:
|
||||
warnings.warn(
|
||||
"dropout option adds dropout after all but last "
|
||||
"recurrent layer, so non-zero dropout expects "
|
||||
f"num_layers greater than 1, but got dropout={dropout} "
|
||||
f"and num_layers={num_layers}",
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
layers = [
|
||||
_LSTMLayer(
|
||||
self.input_size,
|
||||
self.hidden_size,
|
||||
self.bias,
|
||||
batch_first=False,
|
||||
bidirectional=self.bidirectional,
|
||||
split_gates=split_gates,
|
||||
**factory_kwargs,
|
||||
)
|
||||
]
|
||||
layers.extend(
|
||||
_LSTMLayer(
|
||||
self.hidden_size,
|
||||
self.hidden_size,
|
||||
self.bias,
|
||||
batch_first=False,
|
||||
bidirectional=self.bidirectional,
|
||||
split_gates=split_gates,
|
||||
**factory_kwargs,
|
||||
)
|
||||
for _ in range(1, num_layers)
|
||||
)
|
||||
self.layers = torch.nn.ModuleList(layers)
|
||||
|
||||
def forward(self, x: Tensor, hidden: tuple[Tensor, Tensor] | None = None):
|
||||
if self.batch_first:
|
||||
x = x.transpose(0, 1)
|
||||
|
||||
max_batch_size = x.size(1)
|
||||
num_directions = 2 if self.bidirectional else 1
|
||||
if hidden is None:
|
||||
zeros = torch.zeros(
|
||||
num_directions,
|
||||
max_batch_size,
|
||||
self.hidden_size,
|
||||
dtype=torch.float,
|
||||
device=x.device,
|
||||
)
|
||||
zeros.squeeze_(0)
|
||||
if x.is_quantized:
|
||||
zeros = torch.quantize_per_tensor(
|
||||
zeros, scale=1.0, zero_point=0, dtype=x.dtype
|
||||
)
|
||||
hxcx = [(zeros, zeros) for _ in range(self.num_layers)]
|
||||
else:
|
||||
hidden_non_opt = torch.jit._unwrap_optional(hidden)
|
||||
if isinstance(hidden_non_opt[0], Tensor):
|
||||
hx = hidden_non_opt[0].reshape(
|
||||
self.num_layers, num_directions, max_batch_size, self.hidden_size
|
||||
)
|
||||
cx = hidden_non_opt[1].reshape(
|
||||
self.num_layers, num_directions, max_batch_size, self.hidden_size
|
||||
)
|
||||
hxcx = [
|
||||
(hx[idx].squeeze(0), cx[idx].squeeze(0))
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
else:
|
||||
hxcx = hidden_non_opt
|
||||
|
||||
hx_list = []
|
||||
cx_list = []
|
||||
for idx, layer in enumerate(self.layers):
|
||||
x, (h, c) = layer(x, hxcx[idx])
|
||||
hx_list.append(torch.jit._unwrap_optional(h))
|
||||
cx_list.append(torch.jit._unwrap_optional(c))
|
||||
hx_tensor = torch.stack(hx_list)
|
||||
cx_tensor = torch.stack(cx_list)
|
||||
|
||||
# We are creating another dimension for bidirectional case
|
||||
# need to collapse it
|
||||
hx_tensor = hx_tensor.reshape(-1, hx_tensor.shape[-2], hx_tensor.shape[-1])
|
||||
cx_tensor = cx_tensor.reshape(-1, cx_tensor.shape[-2], cx_tensor.shape[-1])
|
||||
|
||||
if self.batch_first:
|
||||
x = x.transpose(0, 1)
|
||||
|
||||
return x, (hx_tensor, cx_tensor)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizableLSTM"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, other, qconfig=None, split_gates=False):
|
||||
if not isinstance(other, cls._FLOAT_MODULE):
|
||||
raise AssertionError(
|
||||
f"Expected module type {cls._FLOAT_MODULE}, got {type(other)}"
|
||||
)
|
||||
if not hasattr(other, "qconfig") and not qconfig:
|
||||
raise AssertionError("other must have qconfig or qconfig must be provided")
|
||||
observed = cls(
|
||||
other.input_size,
|
||||
other.hidden_size,
|
||||
other.num_layers,
|
||||
other.bias,
|
||||
other.batch_first,
|
||||
other.dropout,
|
||||
other.bidirectional,
|
||||
split_gates=split_gates,
|
||||
)
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
observed.qconfig = getattr(other, "qconfig", qconfig)
|
||||
for idx in range(other.num_layers):
|
||||
observed.layers[idx] = _LSTMLayer.from_float(
|
||||
other, idx, qconfig, batch_first=False, split_gates=split_gates
|
||||
)
|
||||
|
||||
# Prepare the model
|
||||
if other.training:
|
||||
observed.train()
|
||||
observed = torch.ao.quantization.prepare_qat(observed, inplace=True)
|
||||
else:
|
||||
observed.eval()
|
||||
observed = torch.ao.quantization.prepare(observed, inplace=True)
|
||||
return observed
|
||||
|
||||
@classmethod
|
||||
def from_observed(cls, other):
|
||||
# The whole flow is float -> observed -> quantized
|
||||
# This class does float -> observed only
|
||||
raise NotImplementedError(
|
||||
"It looks like you are trying to convert a "
|
||||
"non-quantizable LSTM module. Please, see "
|
||||
"the examples on quantizable LSTMs."
|
||||
)
|
||||
@@ -0,0 +1,39 @@
|
||||
from . import functional
|
||||
from .modules import * # noqa: F403
|
||||
from .modules import MaxPool2d
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BatchNorm2d",
|
||||
"BatchNorm3d",
|
||||
"Conv1d",
|
||||
"Conv2d",
|
||||
"Conv3d",
|
||||
"ConvTranspose1d",
|
||||
"ConvTranspose2d",
|
||||
"ConvTranspose3d",
|
||||
"DeQuantize",
|
||||
"ELU",
|
||||
"Embedding",
|
||||
"EmbeddingBag",
|
||||
"GroupNorm",
|
||||
"Hardswish",
|
||||
"InstanceNorm1d",
|
||||
"InstanceNorm2d",
|
||||
"InstanceNorm3d",
|
||||
"LayerNorm",
|
||||
"LeakyReLU",
|
||||
"Linear",
|
||||
"LSTM",
|
||||
"MultiheadAttention",
|
||||
"Quantize",
|
||||
"ReLU6",
|
||||
"Sigmoid",
|
||||
"Softmax",
|
||||
"Dropout",
|
||||
"PReLU",
|
||||
# Wrapper modules
|
||||
"FloatFunctional",
|
||||
"FXFloatFunctional",
|
||||
"QFunctional",
|
||||
]
|
||||
@@ -0,0 +1 @@
|
||||
from .modules import * # noqa: F403
|
||||
+26
@@ -0,0 +1,26 @@
|
||||
from .conv import (
|
||||
Conv1d,
|
||||
Conv2d,
|
||||
Conv3d,
|
||||
ConvTranspose1d,
|
||||
ConvTranspose2d,
|
||||
ConvTranspose3d,
|
||||
)
|
||||
from .linear import Linear
|
||||
from .rnn import GRU, GRUCell, LSTM, LSTMCell, RNNCell
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Linear",
|
||||
"LSTM",
|
||||
"GRU",
|
||||
"LSTMCell",
|
||||
"RNNCell",
|
||||
"GRUCell",
|
||||
"Conv1d",
|
||||
"Conv2d",
|
||||
"Conv3d",
|
||||
"ConvTranspose1d",
|
||||
"ConvTranspose2d",
|
||||
"ConvTranspose3d",
|
||||
]
|
||||
+531
@@ -0,0 +1,531 @@
|
||||
# mypy: allow-untyped-defs
|
||||
r"""Dynamically quantized convolution modules."""
|
||||
|
||||
import warnings
|
||||
from typing import ClassVar, Literal
|
||||
|
||||
import torch
|
||||
import torch.ao.nn.quantized as nnq
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
from torch._ops import ops
|
||||
from torch.ao.nn.quantized.modules.conv import _reverse_repeat_padding
|
||||
from torch.nn.common_types import _size_1_t
|
||||
from torch.nn.modules.utils import _pair, _single, _triple
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Conv1d",
|
||||
"Conv2d",
|
||||
"Conv3d",
|
||||
"ConvTranspose1d",
|
||||
"ConvTranspose2d",
|
||||
"ConvTranspose3d",
|
||||
]
|
||||
|
||||
|
||||
class Conv1d(nnq.Conv1d):
|
||||
r"""A dynamically quantized conv module with floating point tensors as inputs and outputs.
|
||||
|
||||
For details on input arguments, parameters, and implementation see
|
||||
:class:`~torch.nn.Conv1d` and :class:`~torch.ao.nn.quantized.dynamic.Conv1d` and
|
||||
|
||||
Attributes:
|
||||
weight (Tensor): packed tensor derived from the learnable weight
|
||||
parameter.
|
||||
scale (Tensor): scalar for the output scale
|
||||
zero_point (Tensor): scalar for the output zero point
|
||||
|
||||
See :class:`~torch.nn.Conv1d` for other attributes.
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> m = nn.quantized.dynamic.Conv1d(16, 33, 3, stride=2)
|
||||
>>> input = torch.randn(20, 16, 100)
|
||||
>>> output = m(input)
|
||||
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE: ClassVar[type[nn.Conv1d]] = nn.Conv1d
|
||||
_NNIQAT_CONV_BN_MODULE: ClassVar[type[nn.Module] | None] = None
|
||||
_NNI_CONV_RELU_MODULE: ClassVar[type[nn.Module] | None] = None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size: _size_1_t,
|
||||
stride: _size_1_t = 1,
|
||||
padding: _size_1_t = 0,
|
||||
dilation: _size_1_t = 1,
|
||||
groups: int = 1,
|
||||
bias: bool = True,
|
||||
padding_mode: Literal["zeros", "reflect", "replicate", "circular"] = "zeros",
|
||||
device=None,
|
||||
dtype=None,
|
||||
reduce_range=True,
|
||||
):
|
||||
warnings.warn(
|
||||
f"The current implementation of the {self._get_name()} module has poor numerical accuracy and its use is not recommended", # noqa: B950
|
||||
stacklevel=2,
|
||||
)
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
kernel_size = _single(kernel_size)
|
||||
stride = _single(stride)
|
||||
|
||||
padding = padding if isinstance(padding, str) else _single(padding)
|
||||
dilation = _single(dilation)
|
||||
|
||||
super().__init__(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
dilation,
|
||||
groups,
|
||||
bias,
|
||||
padding_mode,
|
||||
**factory_kwargs,
|
||||
)
|
||||
|
||||
def _get_name(self):
|
||||
return "DynamicQuantizedConv1d"
|
||||
|
||||
def forward(self, input: Tensor, reduce_range: bool = True) -> Tensor:
|
||||
# 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 ops.quantized.conv1d_dynamic(input, self._packed_params, reduce_range)
|
||||
|
||||
|
||||
class Conv2d(nnq.Conv2d):
|
||||
r"""A dynamically quantized conv module with floating point tensors as inputs and outputs.
|
||||
|
||||
For details on input arguments, parameters, and implementation see
|
||||
:class:`~torch.nn.Conv2d` and :class:`~torch.ao.nn.quantized.dynamic.Conv2d` and
|
||||
|
||||
Attributes:
|
||||
weight (Tensor): packed tensor derived from the learnable weight
|
||||
parameter.
|
||||
scale (Tensor): scalar for the output scale
|
||||
zero_point (Tensor): scalar for the output zero point
|
||||
|
||||
See :class:`~torch.nn.Conv2d` for other attributes.
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> # With square kernels and equal stride
|
||||
>>> m = nn.quantized.dynamic.Conv2d(16, 33, 3, stride=2)
|
||||
>>> # non-square kernels and unequal stride and with padding
|
||||
>>> m = nn.quantized.dynamic.Conv2d(16, 33, (3, 5), stride=(2, 1), padding=(4, 2))
|
||||
>>> # non-square kernels and unequal stride and with padding and dilation
|
||||
>>> m = nn.quantized.dynamic.Conv2d(16, 33, (3, 5), stride=(2, 1), padding=(4, 2), dilation=(3, 1))
|
||||
>>> input = torch.randn(20, 16, 50, 100)
|
||||
>>> output = m(input)
|
||||
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE: ClassVar[type[nn.Conv2d]] = nn.Conv2d
|
||||
_NNIQAT_CONV_BN_MODULE: ClassVar[type[nn.Module] | None] = None
|
||||
_NNI_CONV_RELU_MODULE: ClassVar[type[nn.Module] | None] = None
|
||||
|
||||
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,
|
||||
):
|
||||
warnings.warn(
|
||||
f"The current implementation of the {self._get_name()} module "
|
||||
"has poor numerical accuracy and its use is not recommended",
|
||||
stacklevel=2,
|
||||
)
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
kernel_size = _pair(kernel_size)
|
||||
stride = _pair(stride)
|
||||
padding = _pair(padding)
|
||||
dilation = _pair(dilation)
|
||||
|
||||
super().__init__(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
dilation,
|
||||
groups,
|
||||
bias,
|
||||
padding_mode,
|
||||
**factory_kwargs,
|
||||
)
|
||||
|
||||
def _get_name(self):
|
||||
return "DynamicQuantizedConv2d"
|
||||
|
||||
def forward(self, input: Tensor, reduce_range: bool = True) -> Tensor:
|
||||
# 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 ops.quantized.conv2d_dynamic(input, self._packed_params, reduce_range)
|
||||
|
||||
|
||||
class Conv3d(nnq.Conv3d):
|
||||
r"""A dynamically quantized conv module with floating point tensors as inputs and outputs.
|
||||
|
||||
For details on input arguments, parameters, and implementation see
|
||||
:class:`~torch.nn.Conv3d` and :class:`~torch.ao.nn.quantized.dynamic.Conv3d` and
|
||||
|
||||
Attributes:
|
||||
weight (Tensor): packed tensor derived from the learnable weight
|
||||
parameter.
|
||||
scale (Tensor): scalar for the output scale
|
||||
zero_point (Tensor): scalar for the output zero point
|
||||
|
||||
See :class:`~torch.nn.Conv3d` for other attributes.
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> # With square kernels and equal stride
|
||||
>>> m = nn.quantized.dynamic.Conv3d(16, 33, 3, stride=2)
|
||||
>>> # non-square kernels and unequal stride and with padding
|
||||
>>> m = nn.quantized.dynamic.Conv3d(16, 33, (3, 5, 5), stride=(1, 2, 2), padding=(1, 2, 2))
|
||||
>>> # non-square kernels and unequal stride and with padding and dilation
|
||||
>>> m = nn.quantized.dynamic.Conv3d(16, 33, (3, 5, 5), stride=(1, 2, 2), padding=(1, 2, 2), dilation=(1, 2, 2))
|
||||
>>> input = torch.randn(20, 16, 56, 56, 56)
|
||||
>>> output = m(input)
|
||||
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE: ClassVar[type[nn.Conv3d]] = nn.Conv3d
|
||||
_NNIQAT_CONV_BN_MODULE: ClassVar[type[nn.Module] | None] = None
|
||||
_NNI_CONV_RELU_MODULE: ClassVar[type[nn.Module] | None] = None
|
||||
|
||||
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,
|
||||
):
|
||||
warnings.warn(
|
||||
f"The current implementation of the {self._get_name()} module has poor numerical accuracy and its use is not recommended", # noqa: B950
|
||||
stacklevel=2,
|
||||
)
|
||||
if padding_mode == "reflect":
|
||||
raise AssertionError("Conv3d does not support reflection padding")
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
kernel_size = _triple(kernel_size)
|
||||
stride = _triple(stride)
|
||||
padding = _triple(padding)
|
||||
dilation = _triple(dilation)
|
||||
super()._init(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
dilation,
|
||||
False,
|
||||
_triple(0),
|
||||
groups,
|
||||
bias,
|
||||
padding_mode,
|
||||
**factory_kwargs,
|
||||
)
|
||||
|
||||
def _get_name(self):
|
||||
return "DynamicQuantizedConv3d"
|
||||
|
||||
def forward(self, input: Tensor, reduce_range: bool = True) -> Tensor:
|
||||
# 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 ops.quantized.conv3d_dynamic(input, self._packed_params, reduce_range)
|
||||
|
||||
|
||||
class ConvTranspose1d(nnq.ConvTranspose1d):
|
||||
r"""A dynamically quantized transposed convolution module with floating point tensors as inputs and outputs.
|
||||
|
||||
For details on input arguments, parameters, and implementation see
|
||||
:class:`~torch.nn.ConvTranspose1d`.
|
||||
|
||||
For special notes, please, see :class:`~torch.ao.nn.quantized.dynamic.Conv1d`
|
||||
|
||||
Attributes:
|
||||
weight (Tensor): packed tensor derived from the learnable weight
|
||||
parameter.
|
||||
scale (Tensor): scalar for the output scale
|
||||
zero_point (Tensor): scalar for the output zero point
|
||||
See :class:`~torch.nn.ConvTranspose1d` for other attributes.
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> # With square kernels and equal stride
|
||||
>>> m = nndq.ConvTranspose1d(16, 33, 3, stride=2)
|
||||
>>> # non-square kernels and unequal stride and with padding
|
||||
>>> m = nndq.ConvTranspose1d(16, 33, (3, 5), stride=(2, 1), padding=(4, 2))
|
||||
>>> output = m(input)
|
||||
>>> # exact output size can be also specified as an argument
|
||||
>>> downsample = nndq.Conv1d(16, 16, 3, stride=2, padding=1)
|
||||
>>> upsample = nndq.ConvTranspose1d(16, 16, 3, stride=2, padding=1)
|
||||
>>> h = downsample(input)
|
||||
>>> h.size()
|
||||
torch.Size([1, 16, 6])
|
||||
>>> output = upsample(h, output_size=input.size())
|
||||
>>> output.size()
|
||||
torch.Size([1, 16, 12])
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE: ClassVar[type[nn.ConvTranspose1d]] = nn.ConvTranspose1d
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
output_padding=0,
|
||||
groups=1,
|
||||
bias=True,
|
||||
dilation=1,
|
||||
padding_mode="zeros",
|
||||
device=None,
|
||||
dtype=None,
|
||||
):
|
||||
warnings.warn(
|
||||
f"The current implementation of the {self._get_name()} module has poor numerical accuracy and its use is not recommended", # noqa: B950
|
||||
stacklevel=2,
|
||||
)
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
output_padding,
|
||||
groups,
|
||||
bias,
|
||||
dilation,
|
||||
padding_mode,
|
||||
**factory_kwargs,
|
||||
)
|
||||
|
||||
def _get_name(self):
|
||||
return "DynamicQuantizedConvTranspose1d"
|
||||
|
||||
def forward(self, input: Tensor, reduce_range: bool = True) -> Tensor:
|
||||
# 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)`!")
|
||||
return torch.ops.quantized.conv_transpose1d_dynamic(
|
||||
input, self._packed_params, reduce_range
|
||||
)
|
||||
|
||||
|
||||
class ConvTranspose2d(nnq.ConvTranspose2d):
|
||||
r"""A dynamically quantized transposed convolution module with floating point tensors as inputs and outputs.
|
||||
|
||||
For details on input arguments, parameters, and implementation see
|
||||
:class:`~torch.nn.ConvTranspose2d`.
|
||||
|
||||
For special notes, please, see :class:`~torch.ao.nn.quantized.dynamic.Conv2d`
|
||||
|
||||
Attributes:
|
||||
weight (Tensor): packed tensor derived from the learnable weight
|
||||
parameter.
|
||||
scale (Tensor): scalar for the output scale
|
||||
zero_point (Tensor): scalar for the output zero point
|
||||
See :class:`~torch.nn.ConvTranspose2d` for other attributes.
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> # With square kernels and equal stride
|
||||
>>> m = nnq.ConvTranspose2d(16, 33, 3, stride=2)
|
||||
>>> # non-square kernels and unequal stride and with padding
|
||||
>>> m = nnq.ConvTranspose2d(16, 33, (3, 5), stride=(2, 1), padding=(4, 2))
|
||||
>>> output = m(input)
|
||||
>>> # exact output size can be also specified as an argument
|
||||
>>> downsample = nnq.Conv2d(16, 16, 3, stride=2, padding=1)
|
||||
>>> upsample = nnq.ConvTranspose2d(16, 16, 3, stride=2, padding=1)
|
||||
>>> h = downsample(input)
|
||||
>>> h.size()
|
||||
torch.Size([1, 16, 6, 6])
|
||||
>>> output = upsample(h, output_size=input.size())
|
||||
>>> output.size()
|
||||
torch.Size([1, 16, 12, 12])
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE: ClassVar[type[nn.ConvTranspose2d]] = nn.ConvTranspose2d
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
output_padding=0,
|
||||
groups=1,
|
||||
bias=True,
|
||||
dilation=1,
|
||||
padding_mode="zeros",
|
||||
device=None,
|
||||
dtype=None,
|
||||
):
|
||||
warnings.warn(
|
||||
f"The current implementation of the {self._get_name()} module has poor numerical accuracy and its use is not recommended", # noqa: B950
|
||||
stacklevel=2,
|
||||
)
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
output_padding,
|
||||
groups,
|
||||
bias,
|
||||
dilation,
|
||||
padding_mode,
|
||||
**factory_kwargs,
|
||||
)
|
||||
|
||||
def _get_name(self):
|
||||
return "DynamicQuantizedConvTranspose2d"
|
||||
|
||||
def forward(self, input: Tensor, reduce_range: bool = True) -> Tensor:
|
||||
# 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 ops.quantized.conv_transpose2d_dynamic(
|
||||
input, self._packed_params, reduce_range
|
||||
)
|
||||
|
||||
|
||||
class ConvTranspose3d(nnq.ConvTranspose3d):
|
||||
r"""A dynamically quantized transposed convolution module with floating point tensors as inputs and outputs.
|
||||
|
||||
For details on input arguments, parameters, and implementation see
|
||||
:class:`~torch.nn.ConvTranspose3d`.
|
||||
|
||||
For special notes, please, see :class:`~torch.ao.nn.quantized.dynamic.Conv3d`
|
||||
|
||||
Attributes:
|
||||
weight (Tensor): packed tensor derived from the learnable weight
|
||||
parameter.
|
||||
scale (Tensor): scalar for the output scale
|
||||
zero_point (Tensor): scalar for the output zero point
|
||||
See :class:`~torch.nn.ConvTranspose3d` for other attributes.
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> # With cubic kernels and equal stride
|
||||
>>> m = nnq.ConvTranspose3d(16, 33, 3, stride=2)
|
||||
>>> # non-cubic kernels and unequal stride and with padding
|
||||
>>> m = nnq.ConvTranspose3d(16, 33, (3, 3, 5), stride=(2, 1, 1), padding=(4, 2, 2))
|
||||
>>> output = m(input)
|
||||
>>> # exact output size can be also specified as an argument
|
||||
>>> downsample = nnq.Conv3d(16, 16, 3, stride=2, padding=1)
|
||||
>>> upsample = nnq.ConvTranspose3d(16, 16, 3, stride=2, padding=1)
|
||||
>>> h = downsample(input)
|
||||
>>> h.size()
|
||||
torch.Size([1, 16, 6, 6, 6])
|
||||
>>> output = upsample(h, output_size=input.size())
|
||||
>>> output.size()
|
||||
torch.Size([1, 16, 12, 12, 12])
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE: ClassVar[type[nn.ConvTranspose3d]] = nn.ConvTranspose3d
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
output_padding=0,
|
||||
groups=1,
|
||||
bias=True,
|
||||
dilation=1,
|
||||
padding_mode="zeros",
|
||||
device=None,
|
||||
dtype=None,
|
||||
):
|
||||
warnings.warn(
|
||||
f"The current implementation of the {self._get_name()} module has poor numerical accuracy and its use is not recommended", # noqa: B950
|
||||
stacklevel=2,
|
||||
)
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
output_padding,
|
||||
groups,
|
||||
bias,
|
||||
dilation,
|
||||
padding_mode,
|
||||
**factory_kwargs,
|
||||
)
|
||||
|
||||
def _get_name(self):
|
||||
return "DynamicQuantizedConvTranspose3d"
|
||||
|
||||
def forward(self, input: Tensor, reduce_range: bool = True) -> Tensor:
|
||||
# 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, T, H, W)`!")
|
||||
return ops.quantized.conv_transpose3d_dynamic(
|
||||
input, self._packed_params, reduce_range
|
||||
)
|
||||
+170
@@ -0,0 +1,170 @@
|
||||
# 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__ = [
|
||||
"Linear",
|
||||
]
|
||||
|
||||
|
||||
class Linear(nnq.Linear):
|
||||
r"""
|
||||
A dynamic quantized linear module with floating point tensor as inputs and outputs.
|
||||
We adopt the same interface as `torch.nn.Linear`, please see
|
||||
https://pytorch.org/docs/stable/nn.html#torch.nn.Linear for documentation.
|
||||
|
||||
Similar to :class:`torch.nn.Linear`, attributes will be randomly
|
||||
initialized at module creation time and will be overwritten later
|
||||
|
||||
Attributes:
|
||||
weight (Tensor): the non-learnable quantized weights of the module which are of
|
||||
shape :math:`(\text{out\_features}, \text{in\_features})`.
|
||||
bias (Tensor): the non-learnable floating point bias of the module of shape
|
||||
:math:`(\text{out\_features})`. If :attr:`bias` is ``True``,
|
||||
the values are initialized to zero.
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> m = nn.quantized.dynamic.Linear(20, 30)
|
||||
>>> input = torch.randn(128, 20)
|
||||
>>> output = m(input)
|
||||
>>> print(output.size())
|
||||
torch.Size([128, 30])
|
||||
"""
|
||||
|
||||
# version used in this class is different from the parent class nnq.Linear
|
||||
_version = 4
|
||||
|
||||
def __init__(self, in_features, out_features, bias_=True, dtype=torch.qint8):
|
||||
super().__init__(in_features, out_features, bias_, dtype=dtype)
|
||||
# We don't muck around with buffers or attributes or anything here
|
||||
# to keep the module simple. *everything* is simply a Python attribute.
|
||||
# Serialization logic is explicitly handled in the below serialization and
|
||||
# deserialization modules
|
||||
self.version = 4
|
||||
|
||||
def forward(self, x):
|
||||
# Note that we can handle self.bias == None case.
|
||||
if self._packed_params.dtype == torch.qint8:
|
||||
if self.version is None or self.version < 4:
|
||||
Y = torch.ops.quantized.linear_dynamic(
|
||||
x, self._packed_params._packed_params
|
||||
)
|
||||
else:
|
||||
Y = torch.ops.quantized.linear_dynamic(
|
||||
x, self._packed_params._packed_params, reduce_range=True
|
||||
)
|
||||
elif self._packed_params.dtype == torch.float16:
|
||||
Y = torch.ops.quantized.linear_dynamic_fp16(
|
||||
x, self._packed_params._packed_params
|
||||
)
|
||||
else:
|
||||
raise RuntimeError("Unsupported dtype on dynamic quantized linear!")
|
||||
return Y.to(x.dtype)
|
||||
|
||||
def _get_name(self):
|
||||
return "DynamicQuantizedLinear"
|
||||
|
||||
def extra_repr(self):
|
||||
extra_repr_str = f"in_features={self.in_features}, out_features={self.out_features}, dtype={self._packed_params.dtype}"
|
||||
if self._packed_params.dtype == torch.qint8:
|
||||
extra_repr_str += f", qscheme={self.weight().qscheme()}"
|
||||
return extra_repr_str
|
||||
|
||||
def _load_from_state_dict(
|
||||
self,
|
||||
state_dict,
|
||||
prefix,
|
||||
local_metadata,
|
||||
strict,
|
||||
missing_keys,
|
||||
unexpected_keys,
|
||||
error_msgs,
|
||||
):
|
||||
version = local_metadata.get("version", None)
|
||||
self.version = version
|
||||
super()._load_from_state_dict(
|
||||
state_dict,
|
||||
prefix,
|
||||
local_metadata,
|
||||
False,
|
||||
missing_keys,
|
||||
unexpected_keys,
|
||||
error_msgs,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False):
|
||||
r"""Create a dynamic quantized module from a float module or qparams_dict
|
||||
|
||||
Args:
|
||||
mod (Module): a float module, either produced by torch.ao.quantization
|
||||
utilities or provided by the user
|
||||
"""
|
||||
float_modules = [
|
||||
torch.nn.Linear,
|
||||
torch.nn.modules.linear.NonDynamicallyQuantizableLinear,
|
||||
torch.ao.nn.intrinsic.modules.fused.LinearReLU,
|
||||
torch.ao.nn.qat.dynamic.Linear,
|
||||
]
|
||||
|
||||
if type(mod) not in float_modules:
|
||||
raise AssertionError(
|
||||
"nn.quantized.dynamic.Linear.from_float only works for one of"
|
||||
+ str([float_mod.__name__ for float_mod in float_modules])
|
||||
+ f", got {type(mod)}"
|
||||
)
|
||||
if not hasattr(mod, "qconfig"):
|
||||
raise AssertionError("Input float module must have qconfig defined")
|
||||
if type(mod) is nni.LinearReLU:
|
||||
mod = mod[0]
|
||||
|
||||
if mod.qconfig is not None and mod.qconfig.weight is not None:
|
||||
weight_observer = mod.qconfig.weight()
|
||||
else:
|
||||
# We have the circular import issues if we import the qconfig in the beginning of this file:
|
||||
# https://github.com/pytorch/pytorch/pull/24231. The current workaround is to postpone the
|
||||
# import until we need it.
|
||||
from torch.ao.quantization.qconfig import default_dynamic_qconfig
|
||||
|
||||
weight_observer = default_dynamic_qconfig.weight()
|
||||
dtype = weight_observer.dtype
|
||||
if dtype not in [torch.qint8, torch.float16]:
|
||||
raise AssertionError(
|
||||
f"The only supported dtypes for dynamic quantized linear are qint8 and float16, got: {dtype}"
|
||||
)
|
||||
weight_observer(mod.weight)
|
||||
if dtype == torch.qint8:
|
||||
qweight = _quantize_weight(mod.weight.float(), weight_observer)
|
||||
elif dtype == torch.float16:
|
||||
qweight = mod.weight.float()
|
||||
else:
|
||||
raise RuntimeError(
|
||||
"Unsupported dtype specified for dynamic quantized Linear!"
|
||||
)
|
||||
qlinear = cls(mod.in_features, mod.out_features, dtype=dtype)
|
||||
|
||||
qlinear.set_weight_bias(qweight, mod.bias)
|
||||
return qlinear
|
||||
|
||||
@classmethod
|
||||
def from_reference(cls, ref_qlinear): # type: ignore[override]
|
||||
"""Create a (fbgemm/qnnpack) dynamic quantized module from a reference quantized
|
||||
module
|
||||
Args:
|
||||
ref_qlinear (Module): a reference quantized module, either produced by
|
||||
torch.ao.quantization functions or provided by the user
|
||||
"""
|
||||
qlinear = cls(
|
||||
ref_qlinear.in_features,
|
||||
ref_qlinear.out_features,
|
||||
dtype=ref_qlinear.weight_dtype,
|
||||
)
|
||||
qweight = ref_qlinear.get_quantized_weight()
|
||||
bias = ref_qlinear.bias
|
||||
qlinear.set_weight_bias(qweight, bias)
|
||||
return qlinear
|
||||
+1375
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,782 @@
|
||||
# mypy: allow-untyped-defs
|
||||
r"""Functional interface (quantized)."""
|
||||
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.jit.annotations import BroadcastingList2
|
||||
from torch.nn.modules.utils import _pair, _triple
|
||||
|
||||
from .modules.utils import _pair_from_first
|
||||
|
||||
|
||||
# Although some of the functions and docstrings are mirrored from the torch.nn,
|
||||
# we want to have them here for future changes.
|
||||
|
||||
__all__ = [
|
||||
"avg_pool2d",
|
||||
"avg_pool3d",
|
||||
"adaptive_avg_pool2d",
|
||||
"adaptive_avg_pool3d",
|
||||
"conv1d",
|
||||
"conv2d",
|
||||
"conv3d",
|
||||
"interpolate",
|
||||
"linear",
|
||||
"max_pool1d",
|
||||
"max_pool2d",
|
||||
"celu",
|
||||
"leaky_relu",
|
||||
"hardtanh",
|
||||
"hardswish",
|
||||
"threshold",
|
||||
"elu",
|
||||
"hardsigmoid",
|
||||
"clamp",
|
||||
"upsample",
|
||||
"upsample_bilinear",
|
||||
"upsample_nearest",
|
||||
]
|
||||
|
||||
|
||||
def avg_pool2d(
|
||||
input,
|
||||
kernel_size,
|
||||
stride=None,
|
||||
padding=0,
|
||||
ceil_mode=False,
|
||||
count_include_pad=True,
|
||||
divisor_override=None,
|
||||
):
|
||||
r"""
|
||||
Applies 2D average-pooling operation in :math:`kH \times kW` regions by step size
|
||||
:math:`sH \times sW` steps. The number of output features is equal to the number of
|
||||
input planes.
|
||||
|
||||
.. note:: The input quantization parameters propagate to the output.
|
||||
|
||||
See :class:`~torch.ao.nn.quantized.AvgPool2d` for details and output shape.
|
||||
|
||||
Args:
|
||||
input: quantized input tensor :math:`(\text{minibatch} , \text{in\_channels} , iH , iW)`
|
||||
kernel_size: size of the pooling region. Can be a single number or a
|
||||
tuple `(kH, kW)`
|
||||
stride: stride of the pooling operation. Can be a single number or a
|
||||
tuple `(sH, sW)`. Default: :attr:`kernel_size`
|
||||
padding: implicit zero paddings on both sides of the input. Can be a
|
||||
single number or a tuple `(padH, padW)`. Default: 0
|
||||
ceil_mode: when True, will use `ceil` instead of `floor` in the formula
|
||||
to compute the output shape. Default: ``False``
|
||||
count_include_pad: when True, will include the zero-padding in the
|
||||
averaging calculation. Default: ``True``
|
||||
divisor_override: if specified, it will be used as divisor, otherwise
|
||||
size of the pooling region will be used. Default: None
|
||||
"""
|
||||
if not input.is_quantized:
|
||||
raise ValueError("Input to 'quantized.avg_pool2d' must be quantized!")
|
||||
return torch.nn.functional.avg_pool2d(
|
||||
input,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
ceil_mode,
|
||||
count_include_pad,
|
||||
divisor_override,
|
||||
)
|
||||
|
||||
|
||||
def avg_pool3d(
|
||||
input,
|
||||
kernel_size,
|
||||
stride=None,
|
||||
padding=0,
|
||||
ceil_mode=False,
|
||||
count_include_pad=True,
|
||||
divisor_override=None,
|
||||
):
|
||||
r"""
|
||||
Applies 3D average-pooling operation in :math:`kD \ times kH \times kW` regions by step size
|
||||
:math:`sD \times sH \times sW` steps. The number of output features is equal to the number of
|
||||
input planes.
|
||||
|
||||
.. note:: The input quantization parameters propagate to the output.
|
||||
|
||||
Args:
|
||||
input: quantized input tensor :math:`(\text{minibatch} , \text{in\_channels} , iH , iW)`
|
||||
kernel_size: size of the pooling region. Can be a single number or a
|
||||
tuple `(kD, kH, kW)`
|
||||
stride: stride of the pooling operation. Can be a single number or a
|
||||
tuple `(sD, sH, sW)`. Default: :attr:`kernel_size`
|
||||
padding: implicit zero paddings on both sides of the input. Can be a
|
||||
single number or a tuple `(padD, padH, padW)`. Default: 0
|
||||
ceil_mode: when True, will use `ceil` instead of `floor` in the formula
|
||||
to compute the output shape. Default: ``False``
|
||||
count_include_pad: when True, will include the zero-padding in the
|
||||
averaging calculation. Default: ``True``
|
||||
divisor_override: if specified, it will be used as divisor, otherwise
|
||||
size of the pooling region will be used. Default: None
|
||||
"""
|
||||
if not input.is_quantized:
|
||||
raise ValueError("Input to 'quantized.avg_pool3d' must be quantized!")
|
||||
return torch.nn.functional.avg_pool3d(
|
||||
input,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
ceil_mode,
|
||||
count_include_pad,
|
||||
divisor_override,
|
||||
)
|
||||
|
||||
|
||||
def adaptive_avg_pool2d(input: Tensor, output_size: BroadcastingList2[int]) -> Tensor:
|
||||
r"""
|
||||
Applies a 2D adaptive average pooling over a quantized input signal composed
|
||||
of several quantized input planes.
|
||||
|
||||
.. note:: The input quantization parameters propagate to the output.
|
||||
|
||||
See :class:`~torch.ao.nn.quantized.AdaptiveAvgPool2d` for details and output shape.
|
||||
|
||||
Args:
|
||||
output_size: the target output size (single integer or
|
||||
double-integer tuple)
|
||||
"""
|
||||
if not input.is_quantized:
|
||||
raise ValueError(
|
||||
"Input to 'quantized.functional.adaptive_avg_pool2d' must be quantized!"
|
||||
)
|
||||
return torch.nn.functional.adaptive_avg_pool2d(input, output_size)
|
||||
|
||||
|
||||
def adaptive_avg_pool3d(input: Tensor, output_size: BroadcastingList2[int]) -> Tensor:
|
||||
r"""
|
||||
Applies a 3D adaptive average pooling over a quantized input signal composed
|
||||
of several quantized input planes.
|
||||
|
||||
.. note:: The input quantization parameters propagate to the output.
|
||||
|
||||
See :class:`~torch.ao.nn.quantized.AdaptiveAvgPool3d` for details and output shape.
|
||||
|
||||
Args:
|
||||
output_size: the target output size (single integer or
|
||||
double-integer tuple)
|
||||
"""
|
||||
if not input.is_quantized:
|
||||
raise ValueError(
|
||||
"Input to 'quantized.functional.adaptive_avg_pool3d' must be quantized!"
|
||||
)
|
||||
return torch.nn.functional.adaptive_avg_pool3d(input, output_size)
|
||||
|
||||
|
||||
def conv1d(
|
||||
input,
|
||||
weight,
|
||||
bias,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
padding_mode="zeros",
|
||||
scale=1.0,
|
||||
zero_point=0,
|
||||
dtype=torch.quint8,
|
||||
):
|
||||
r"""
|
||||
Applies a 1D convolution over a quantized 1D input composed of several input
|
||||
planes.
|
||||
|
||||
See :class:`~torch.ao.nn.quantized.Conv1d` for details and output shape.
|
||||
|
||||
Args:
|
||||
input: quantized input tensor of shape :math:`(\text{minibatch} , \text{in\_channels} , iW)`
|
||||
weight: quantized filters of shape :math:`(\text{out\_channels} , \frac{\text{in\_channels}}{\text{groups}} , iW)`
|
||||
bias: **non-quantized** bias tensor of shape :math:`(\text{out\_channels})`. The tensor type must be `torch.float`.
|
||||
stride: the stride of the convolving kernel. Can be a single number or a
|
||||
tuple `(sW,)`. Default: 1
|
||||
padding: implicit paddings on both sides of the input. Can be a
|
||||
single number or a tuple `(padW,)`. Default: 0
|
||||
dilation: the spacing between kernel elements. Can be a single number or
|
||||
a tuple `(dW,)`. Default: 1
|
||||
groups: split input into groups, :math:`\text{in\_channels}` should be divisible by the
|
||||
number of groups. Default: 1
|
||||
padding_mode: the padding mode to use. Only "zeros" is supported for quantized convolution at the moment. Default: "zeros"
|
||||
scale: quantization scale for the output. Default: 1.0
|
||||
zero_point: quantization zero_point for the output. Default: 0
|
||||
dtype: quantization data type to use. Default: ``torch.quint8``
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_QENGINE)
|
||||
>>> from torch.ao.nn.quantized import functional as qF
|
||||
>>> filters = torch.randn(33, 16, 3, dtype=torch.float)
|
||||
>>> inputs = torch.randn(20, 16, 50, dtype=torch.float)
|
||||
>>> bias = torch.randn(33, dtype=torch.float)
|
||||
>>>
|
||||
>>> scale, zero_point = 1.0, 0
|
||||
>>> dtype_inputs = torch.quint8
|
||||
>>> dtype_filters = torch.qint8
|
||||
>>>
|
||||
>>> q_filters = torch.quantize_per_tensor(filters, scale, zero_point, dtype_filters)
|
||||
>>> q_inputs = torch.quantize_per_tensor(inputs, scale, zero_point, dtype_inputs)
|
||||
>>> qF.conv1d(q_inputs, q_filters, bias, padding=1, scale=scale, zero_point=zero_point)
|
||||
""" # noqa: E501
|
||||
if padding_mode != "zeros":
|
||||
raise NotImplementedError("Only zero-padding is supported!")
|
||||
if input.dtype != torch.quint8:
|
||||
raise NotImplementedError(
|
||||
"Only torch.quint8 is supported for activation tensor!"
|
||||
)
|
||||
if weight.dtype != torch.qint8:
|
||||
raise NotImplementedError("Only torch.qint8 is supported for weight tensor!")
|
||||
if input.ndim != 3:
|
||||
raise ValueError("Input shape must be `(N, C, L)`!")
|
||||
stride = _pair_from_first(stride)
|
||||
padding = _pair_from_first(padding)
|
||||
dilation = _pair_from_first(dilation)
|
||||
|
||||
packed_params = torch.ops.quantized.conv1d_prepack(
|
||||
weight, bias, stride, padding, dilation, groups
|
||||
)
|
||||
return torch.ops.quantized.conv1d(input, packed_params, scale, zero_point)
|
||||
|
||||
|
||||
def conv2d(
|
||||
input,
|
||||
weight,
|
||||
bias,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
padding_mode="zeros",
|
||||
scale=1.0,
|
||||
zero_point=0,
|
||||
dtype=torch.quint8,
|
||||
):
|
||||
r"""
|
||||
Applies a 2D convolution over a quantized 2D input composed of several input
|
||||
planes.
|
||||
|
||||
See :class:`~torch.ao.nn.quantized.Conv2d` for details and output shape.
|
||||
|
||||
Args:
|
||||
input: quantized input tensor of shape :math:`(\text{minibatch} , \text{in\_channels} , iH , iW)`
|
||||
weight: quantized filters of shape :math:`(\text{out\_channels} , \frac{\text{in\_channels}}{\text{groups}} , kH , kW)`
|
||||
bias: **non-quantized** bias tensor of shape :math:`(\text{out\_channels})`. The tensor type must be `torch.float`.
|
||||
stride: the stride of the convolving kernel. Can be a single number or a
|
||||
tuple `(sH, sW)`. Default: 1
|
||||
padding: implicit paddings on both sides of the input. Can be a
|
||||
single number or a tuple `(padH, padW)`. Default: 0
|
||||
dilation: the spacing between kernel elements. Can be a single number or
|
||||
a tuple `(dH, dW)`. Default: 1
|
||||
groups: split input into groups, :math:`\text{in\_channels}` should be divisible by the
|
||||
number of groups. Default: 1
|
||||
padding_mode: the padding mode to use. Only "zeros" is supported for quantized convolution at the moment. Default: "zeros"
|
||||
scale: quantization scale for the output. Default: 1.0
|
||||
zero_point: quantization zero_point for the output. Default: 0
|
||||
dtype: quantization data type to use. Default: ``torch.quint8``
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_QENGINE)
|
||||
>>> from torch.ao.nn.quantized import functional as qF
|
||||
>>> filters = torch.randn(8, 4, 3, 3, dtype=torch.float)
|
||||
>>> inputs = torch.randn(1, 4, 5, 5, dtype=torch.float)
|
||||
>>> bias = torch.randn(8, dtype=torch.float)
|
||||
>>>
|
||||
>>> scale, zero_point = 1.0, 0
|
||||
>>> dtype_inputs = torch.quint8
|
||||
>>> dtype_filters = torch.qint8
|
||||
>>>
|
||||
>>> q_filters = torch.quantize_per_tensor(filters, scale, zero_point, dtype_filters)
|
||||
>>> q_inputs = torch.quantize_per_tensor(inputs, scale, zero_point, dtype_inputs)
|
||||
>>> qF.conv2d(q_inputs, q_filters, bias, padding=1, scale=scale, zero_point=zero_point)
|
||||
""" # noqa: E501
|
||||
if padding_mode != "zeros":
|
||||
raise NotImplementedError("Only zero-padding is supported!")
|
||||
if input.dtype != torch.quint8:
|
||||
raise NotImplementedError(
|
||||
"Only torch.quint8 is supported for activation tensor!"
|
||||
)
|
||||
if weight.dtype != torch.qint8:
|
||||
raise NotImplementedError("Only torch.qint8 is supported for weight tensor!")
|
||||
if input.ndim != 4:
|
||||
raise ValueError("Input shape must be `(N, C, H, W)`!")
|
||||
stride = _pair(stride)
|
||||
padding = _pair(padding)
|
||||
dilation = _pair(dilation)
|
||||
|
||||
packed_params = torch.ops.quantized.conv2d_prepack(
|
||||
weight, bias, stride, padding, dilation, groups
|
||||
)
|
||||
return torch.ops.quantized.conv2d(input, packed_params, scale, zero_point)
|
||||
|
||||
|
||||
def conv3d(
|
||||
input,
|
||||
weight,
|
||||
bias,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
padding_mode="zeros",
|
||||
scale=1.0,
|
||||
zero_point=0,
|
||||
dtype=torch.quint8,
|
||||
):
|
||||
r"""
|
||||
Applies a 3D convolution over a quantized 3D input composed of several input
|
||||
planes.
|
||||
|
||||
See :class:`~torch.ao.nn.quantized.Conv3d` for details and output shape.
|
||||
|
||||
Args:
|
||||
input: quantized input tensor of shape
|
||||
:math:`(\text{minibatch} , \text{in\_channels} , iD , iH , iW)`
|
||||
weight: quantized filters of shape
|
||||
:math:`(\text{out\_channels} , \frac{\text{in\_channels}}{\text{groups}} , kD , kH , kW)`
|
||||
bias: **non-quantized** bias tensor of shape
|
||||
:math:`(\text{out\_channels})`. The tensor type must be `torch.float`.
|
||||
stride: the stride of the convolving kernel. Can be a single number or a
|
||||
tuple `(sD, sH, sW)`. Default: 1
|
||||
padding: implicit paddings on both sides of the input. Can be a
|
||||
single number or a tuple `(padD, padH, padW)`. Default: 0
|
||||
dilation: the spacing between kernel elements. Can be a single number or
|
||||
a tuple `(dD, dH, dW)`. Default: 1
|
||||
groups: split input into groups, :math:`\text{in\_channels}` should be
|
||||
divisible by the number of groups. Default: 1
|
||||
padding_mode: the padding mode to use. Only "zeros" is supported for
|
||||
quantized convolution at the moment. Default: "zeros"
|
||||
scale: quantization scale for the output. Default: 1.0
|
||||
zero_point: quantization zero_point for the output. Default: 0
|
||||
dtype: quantization data type to use. Default: ``torch.quint8``
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_QENGINE)
|
||||
>>> from torch.ao.nn.quantized import functional as qF
|
||||
>>> filters = torch.randn(8, 4, 3, 3, 3, dtype=torch.float)
|
||||
>>> inputs = torch.randn(1, 4, 5, 5, 5, dtype=torch.float)
|
||||
>>> bias = torch.randn(8, dtype=torch.float)
|
||||
>>>
|
||||
>>> scale, zero_point = 1.0, 0
|
||||
>>> dtype_inputs = torch.quint8
|
||||
>>> dtype_filters = torch.qint8
|
||||
>>>
|
||||
>>> q_filters = torch.quantize_per_tensor(filters, scale, zero_point, dtype_filters)
|
||||
>>> q_inputs = torch.quantize_per_tensor(inputs, scale, zero_point, dtype_inputs)
|
||||
>>> qF.conv3d(q_inputs, q_filters, bias, padding=1, scale=scale, zero_point=zero_point)
|
||||
""" # noqa: E501
|
||||
if padding_mode != "zeros":
|
||||
raise NotImplementedError("Only zero-padding is supported!")
|
||||
if input.dtype != torch.quint8:
|
||||
raise NotImplementedError(
|
||||
"Only torch.quint8 is supported for activation tensor!"
|
||||
)
|
||||
if weight.dtype != torch.qint8:
|
||||
raise NotImplementedError("Only torch.qint8 is supported for weight tensor!")
|
||||
if input.ndim != 5:
|
||||
raise ValueError("Input shape must be `(N, C, D, H, W)`!")
|
||||
stride = _triple(stride)
|
||||
padding = _triple(padding)
|
||||
dilation = _triple(dilation)
|
||||
|
||||
packed_params = torch.ops.quantized.conv3d_prepack(
|
||||
weight, bias, stride, padding, dilation, groups
|
||||
)
|
||||
return torch.ops.quantized.conv3d(input, packed_params, scale, zero_point)
|
||||
|
||||
|
||||
def interpolate(
|
||||
input, size=None, scale_factor=None, mode="nearest", align_corners=None
|
||||
):
|
||||
r"""Down/up samples the input to either the given :attr:`size` or the given
|
||||
:attr:`scale_factor`
|
||||
|
||||
See :func:`torch.nn.functional.interpolate` for implementation details.
|
||||
|
||||
The input dimensions are interpreted in the form:
|
||||
`mini-batch x channels x [optional depth] x [optional height] x width`.
|
||||
|
||||
.. note:: The input quantization parameters propagate to the output.
|
||||
|
||||
.. note:: Only 2D/3D input is supported for quantized inputs
|
||||
|
||||
.. note:: Only the following modes are supported for the quantized inputs:
|
||||
|
||||
- `bilinear`
|
||||
- `nearest`
|
||||
|
||||
Args:
|
||||
input (Tensor): the input tensor
|
||||
size (int or Tuple[int] or Tuple[int, int] or Tuple[int, int, int]):
|
||||
output spatial size.
|
||||
scale_factor (float or Tuple[float]): multiplier for spatial size. Has to match input size if it is a tuple.
|
||||
mode (str): algorithm used for upsampling:
|
||||
``'nearest'`` | ``'bilinear'``
|
||||
align_corners (bool, optional): Geometrically, we consider the pixels of the
|
||||
input and output as squares rather than points.
|
||||
If set to ``True``, the input and output tensors are aligned by the
|
||||
center points of their corner pixels, preserving the values at the corner pixels.
|
||||
If set to ``False``, the input and output tensors are aligned by the corner
|
||||
points of their corner pixels, and the interpolation uses edge value padding
|
||||
for out-of-boundary values, making this operation *independent* of input size
|
||||
when :attr:`scale_factor` is kept the same. This only has an effect when :attr:`mode`
|
||||
is ``'bilinear'``.
|
||||
Default: ``False``
|
||||
"""
|
||||
if not input.is_quantized:
|
||||
raise ValueError("Input to 'quantized.interpolate' must be quantized!")
|
||||
return torch.nn.functional.interpolate(
|
||||
input, size, scale_factor, mode, align_corners
|
||||
)
|
||||
|
||||
|
||||
def linear(
|
||||
input: Tensor,
|
||||
weight: Tensor,
|
||||
bias: Tensor | None = None,
|
||||
scale: float | None = None,
|
||||
zero_point: int | None = None,
|
||||
) -> Tensor:
|
||||
r"""
|
||||
Applies a linear transformation to the incoming quantized data:
|
||||
:math:`y = xA^T + b`.
|
||||
See :class:`~torch.ao.nn.quantized.Linear`
|
||||
|
||||
.. note::
|
||||
|
||||
Current implementation packs weights on every call, which has penalty on performance.
|
||||
If you want to avoid the overhead, use :class:`~torch.ao.nn.quantized.Linear`.
|
||||
|
||||
Args:
|
||||
input (Tensor): Quantized input of type `torch.quint8`
|
||||
weight (Tensor): Quantized weight of type `torch.qint8`
|
||||
bias (Tensor): None or fp32 bias of type `torch.float`
|
||||
scale (double): output scale. If None, derived from the input scale
|
||||
zero_point (long): output zero point. If None, derived from the input zero_point
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, *, in\_features)` where `*` means any number of
|
||||
additional dimensions
|
||||
- Weight: :math:`(out\_features, in\_features)`
|
||||
- Bias: :math:`(out\_features)`
|
||||
- Output: :math:`(N, *, out\_features)`
|
||||
"""
|
||||
if scale is None:
|
||||
scale = input.q_scale()
|
||||
if zero_point is None:
|
||||
zero_point = input.q_zero_point()
|
||||
_packed_params = torch.ops.quantized.linear_prepack(weight, bias)
|
||||
return torch.ops.quantized.linear(input, _packed_params, scale, zero_point)
|
||||
|
||||
|
||||
def max_pool1d(
|
||||
input,
|
||||
kernel_size,
|
||||
stride=None,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
ceil_mode=False,
|
||||
return_indices=False,
|
||||
):
|
||||
r"""Applies a 1D max pooling over a quantized input signal composed of
|
||||
several quantized input planes.
|
||||
|
||||
.. note:: The input quantization parameters are propagated to the output.
|
||||
|
||||
See :class:`~torch.ao.nn.quantized.MaxPool1d` for details.
|
||||
"""
|
||||
if return_indices:
|
||||
raise NotImplementedError("return_indices is not yet implemented!")
|
||||
if stride is None:
|
||||
stride = torch.jit.annotate(list[int], [])
|
||||
return torch.nn.functional.max_pool1d(
|
||||
input,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
dilation,
|
||||
ceil_mode=ceil_mode,
|
||||
return_indices=return_indices,
|
||||
)
|
||||
|
||||
|
||||
def max_pool2d(
|
||||
input,
|
||||
kernel_size,
|
||||
stride=None,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
ceil_mode=False,
|
||||
return_indices=False,
|
||||
):
|
||||
r"""Applies a 2D max pooling over a quantized input signal composed of
|
||||
several quantized input planes.
|
||||
|
||||
.. note:: The input quantization parameters are propagated to the output.
|
||||
|
||||
See :class:`~torch.ao.nn.quantized.MaxPool2d` for details.
|
||||
"""
|
||||
if return_indices:
|
||||
raise NotImplementedError("return_indices is not yet implemented!")
|
||||
if stride is None:
|
||||
stride = torch.jit.annotate(list[int], [])
|
||||
return torch.nn.functional.max_pool2d(
|
||||
input,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
dilation,
|
||||
ceil_mode=ceil_mode,
|
||||
return_indices=return_indices,
|
||||
)
|
||||
|
||||
|
||||
def celu(input: Tensor, scale: float, zero_point: int, alpha: float = 1.0) -> Tensor:
|
||||
r"""celu(input, scale, zero_point, alpha=1.) -> Tensor
|
||||
|
||||
Applies the quantized CELU function element-wise.
|
||||
|
||||
.. math::
|
||||
\text{CELU}(x) = \max(0,x) + \min(0, \alpha * (\exp(x / \alpha) - 1))
|
||||
|
||||
Args:
|
||||
input: quantized input
|
||||
alpha: the :math:`\alpha` value for the CELU formulation. Default: 1.0
|
||||
"""
|
||||
if not input.is_quantized:
|
||||
raise ValueError("Input to 'quantized.celu' must be quantized!")
|
||||
return torch.ops.quantized.celu(input, scale, zero_point, alpha)
|
||||
|
||||
|
||||
def leaky_relu(
|
||||
input: Tensor,
|
||||
negative_slope: float = 0.01,
|
||||
inplace: bool = False,
|
||||
scale: float | None = None,
|
||||
zero_point: int | None = None,
|
||||
):
|
||||
r"""
|
||||
Quantized version of the.
|
||||
leaky_relu(input, negative_slope=0.01, inplace=False, scale, zero_point) -> Tensor
|
||||
|
||||
Applies element-wise,
|
||||
:math:`\text{LeakyReLU}(x) = \max(0, x) + \text{negative\_slope} * \min(0, x)`
|
||||
|
||||
Args:
|
||||
input: Quantized input
|
||||
negative_slope: The slope of the negative input
|
||||
inplace: Inplace modification of the input tensor
|
||||
scale, zero_point: Scale and zero point of the output tensor.
|
||||
|
||||
See :class:`~torch.nn.LeakyReLU` for more details.
|
||||
"""
|
||||
if scale is not None and zero_point is not None:
|
||||
if inplace:
|
||||
raise AssertionError("Cannot rescale with `inplace`")
|
||||
output = torch._empty_affine_quantized(
|
||||
input.shape, scale=scale, zero_point=int(zero_point), dtype=input.dtype
|
||||
)
|
||||
torch._C._nn.leaky_relu(input, negative_slope, out=output)
|
||||
return output
|
||||
if inplace:
|
||||
result = torch._C._nn.leaky_relu_(input, negative_slope)
|
||||
else:
|
||||
result = torch._C._nn.leaky_relu(input, negative_slope)
|
||||
return result
|
||||
|
||||
|
||||
def hardtanh(
|
||||
input: Tensor, min_val: float = -1.0, max_val: float = 1.0, inplace: bool = False
|
||||
) -> Tensor:
|
||||
r"""This is the quantized version of :func:`~torch.nn.functional.hardtanh`."""
|
||||
if not input.is_quantized:
|
||||
raise ValueError("Input to 'quantized.hardtanh' must be quantized!")
|
||||
if inplace:
|
||||
return torch._C._nn.hardtanh_(input, min_val, max_val)
|
||||
return torch._C._nn.hardtanh(input, min_val, max_val)
|
||||
|
||||
|
||||
def hardswish(input: Tensor, scale: float, zero_point: int) -> Tensor:
|
||||
r"""This is the quantized version of :func:`~torch.nn.functional.hardswish`.
|
||||
|
||||
Args:
|
||||
input: quantized input
|
||||
scale: quantization scale of the output tensor
|
||||
zero_point: quantization zero point of the output tensor
|
||||
"""
|
||||
if not input.is_quantized:
|
||||
raise ValueError("Input to 'quantized.hardswish' must be quantized!")
|
||||
return torch._ops.ops.quantized.hardswish(input, scale, zero_point)
|
||||
|
||||
|
||||
def threshold(input: Tensor, threshold: float, value: float) -> Tensor:
|
||||
r"""Applies the quantized version of the threshold function element-wise:
|
||||
|
||||
.. math::
|
||||
x = \begin{cases}
|
||||
x & \text{if~} x > \text{threshold} \\
|
||||
\text{value} & \text{otherwise}
|
||||
\end{cases}
|
||||
|
||||
See :class:`~torch.nn.Threshold` for more details.
|
||||
"""
|
||||
if not input.is_quantized:
|
||||
raise ValueError("Input to 'quantized.threshold' must be quantized!")
|
||||
if threshold is None:
|
||||
raise ValueError("Input to 'threshold' must be specified!")
|
||||
if value is None:
|
||||
raise ValueError("Input to 'value' must be specified!")
|
||||
return torch._ops.ops.quantized.threshold(input, threshold, value)
|
||||
|
||||
|
||||
def elu(input: Tensor, scale: float, zero_point: int, alpha: float = 1.0) -> Tensor:
|
||||
r"""This is the quantized version of :func:`~torch.nn.functional.elu`.
|
||||
|
||||
Args:
|
||||
input: quantized input
|
||||
scale: quantization scale of the output tensor
|
||||
zero_point: quantization zero point of the output tensor
|
||||
alpha: the alpha constant
|
||||
"""
|
||||
if not input.is_quantized:
|
||||
raise ValueError("Input to 'quantized.elu' must be quantized!")
|
||||
return torch.ops.quantized.elu(input, scale, zero_point, alpha)
|
||||
|
||||
|
||||
def hardsigmoid(input: Tensor, inplace: bool = False) -> Tensor:
|
||||
r"""This is the quantized version of :func:`~torch.nn.functional.hardsigmoid`."""
|
||||
if not input.is_quantized:
|
||||
raise ValueError("Input to 'quantized.hardsigmoid' must be quantized!")
|
||||
if inplace:
|
||||
return torch._C._nn.hardsigmoid_(input) # type: ignore[attr-defined]
|
||||
return torch._C._nn.hardsigmoid(input)
|
||||
|
||||
|
||||
def clamp(input: Tensor, min_: float, max_: float) -> Tensor:
|
||||
r"""float(input, min\_, max\_) -> Tensor
|
||||
|
||||
Applies the clamp function element-wise.
|
||||
See :class:`~torch.ao.nn.quantized.clamp` for more details.
|
||||
|
||||
Args:
|
||||
input: quantized input
|
||||
min_: minimum value for clamping
|
||||
max_: maximum value for clamping
|
||||
"""
|
||||
if not input.is_quantized:
|
||||
raise ValueError("Input to 'quantized.clamp' must be quantized!")
|
||||
return torch.clamp(input, min_, max_)
|
||||
|
||||
|
||||
def upsample(input, size=None, scale_factor=None, mode="nearest", align_corners=None):
|
||||
r"""Upsamples the input to either the given :attr:`size` or the given
|
||||
:attr:`scale_factor`
|
||||
|
||||
.. warning::
|
||||
This function is deprecated in favor of
|
||||
:func:`torch.ao.nn.quantized.functional.interpolate`.
|
||||
This is equivalent with ``nn.quantized.functional.interpolate(...)``.
|
||||
|
||||
See :func:`torch.nn.functional.interpolate` for implementation details.
|
||||
|
||||
The input dimensions are interpreted in the form:
|
||||
`mini-batch x channels x [optional depth] x [optional height] x width`.
|
||||
|
||||
.. note:: The input quantization parameters propagate to the output.
|
||||
|
||||
.. note:: Only 2D input is supported for quantized inputs
|
||||
|
||||
.. note:: Only the following modes are supported for the quantized inputs:
|
||||
|
||||
- `bilinear`
|
||||
- `nearest`
|
||||
|
||||
Args:
|
||||
input (Tensor): quantized input tensor
|
||||
size (int or Tuple[int] or Tuple[int, int] or Tuple[int, int, int]):
|
||||
output spatial size.
|
||||
scale_factor (float or Tuple[float]): multiplier for spatial size. Has to be an integer.
|
||||
mode (str): algorithm used for upsampling:
|
||||
``'nearest'`` | ``'bilinear'``
|
||||
align_corners (bool, optional): Geometrically, we consider the pixels of the
|
||||
input and output as squares rather than points.
|
||||
If set to ``True``, the input and output tensors are aligned by the
|
||||
center points of their corner pixels, preserving the values at the corner pixels.
|
||||
If set to ``False``, the input and output tensors are aligned by the corner
|
||||
points of their corner pixels, and the interpolation uses edge value padding
|
||||
for out-of-boundary values, making this operation *independent* of input size
|
||||
when :attr:`scale_factor` is kept the same. This only has an effect when :attr:`mode`
|
||||
is ``'bilinear'``.
|
||||
Default: ``False``
|
||||
|
||||
.. warning::
|
||||
With ``align_corners = True``, the linearly interpolating modes
|
||||
(`bilinear`) don't proportionally align the
|
||||
output and input pixels, and thus the output values can depend on the
|
||||
input size. This was the default behavior for these modes up to version
|
||||
0.3.1. Since then, the default behavior is ``align_corners = False``.
|
||||
See :class:`~torch.nn.Upsample` for concrete examples on how this
|
||||
affects the outputs.
|
||||
"""
|
||||
warnings.warn(
|
||||
"nn.quantized.functional.upsample is deprecated. Use nn.quantized.functional.interpolate instead.",
|
||||
stacklevel=2,
|
||||
)
|
||||
return interpolate(input, size, scale_factor, mode, align_corners)
|
||||
|
||||
|
||||
def upsample_bilinear(input, size=None, scale_factor=None):
|
||||
r"""Upsamples the input, using bilinear upsampling.
|
||||
|
||||
.. warning::
|
||||
This function is deprecated in favor of
|
||||
:func:`torch.ao.nn.quantized.functional.interpolate`.
|
||||
This is equivalent with
|
||||
``nn.quantized.functional.interpolate(..., mode='bilinear', align_corners=True)``.
|
||||
|
||||
.. note:: The input quantization parameters propagate to the output.
|
||||
|
||||
.. note:: Only 2D inputs are supported
|
||||
|
||||
Args:
|
||||
input (Tensor): quantized input
|
||||
size (int or Tuple[int, int]): output spatial size.
|
||||
scale_factor (int or Tuple[int, int]): multiplier for spatial size
|
||||
"""
|
||||
# DeprecationWarning is ignored by default
|
||||
warnings.warn(
|
||||
"nn.quantized.functional.upsample_bilinear is deprecated. Use nn.quantized.functional.interpolate instead.",
|
||||
stacklevel=2,
|
||||
)
|
||||
return interpolate(input, size, scale_factor, mode="bilinear", align_corners=True)
|
||||
|
||||
|
||||
def upsample_nearest(input, size=None, scale_factor=None):
|
||||
r"""Upsamples the input, using nearest neighbours' pixel values.
|
||||
|
||||
.. warning::
|
||||
This function is deprecated in favor of
|
||||
:func:`torch.ao.nn.quantized.functional.interpolate`.
|
||||
This is equivalent with ``nn.quantized.functional.interpolate(..., mode='nearest')``.
|
||||
|
||||
.. note:: The input quantization parameters propagate to the output.
|
||||
|
||||
.. note:: Only 2D inputs are supported
|
||||
|
||||
Args:
|
||||
input (Tensor): quantized input
|
||||
size (int or Tuple[int, int] or Tuple[int, int, int]): output spatial
|
||||
size.
|
||||
scale_factor (int): multiplier for spatial size. Has to be an integer.
|
||||
"""
|
||||
# DeprecationWarning is ignored by default
|
||||
warnings.warn(
|
||||
"nn.quantized.functional.upsample_nearest is deprecated. Use nn.quantized.functional.interpolate instead.",
|
||||
stacklevel=2,
|
||||
)
|
||||
return interpolate(input, size, scale_factor, mode="nearest")
|
||||
@@ -0,0 +1,165 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import torch
|
||||
|
||||
# The quantized modules use `torch.nn` and `torch.ao.nn.quantizable`
|
||||
# packages. However, the `quantizable` package uses "lazy imports"
|
||||
# to avoid circular dependency.
|
||||
# Hence we need to include it here to make sure it is resolved before
|
||||
# they are used in the modules.
|
||||
import torch.ao.nn.quantizable
|
||||
from torch.nn.modules.pooling import MaxPool2d
|
||||
|
||||
from .activation import (
|
||||
ELU,
|
||||
Hardswish,
|
||||
LeakyReLU,
|
||||
MultiheadAttention,
|
||||
PReLU,
|
||||
ReLU6,
|
||||
Sigmoid,
|
||||
Softmax,
|
||||
)
|
||||
from .batchnorm import BatchNorm2d, BatchNorm3d
|
||||
from .conv import (
|
||||
Conv1d,
|
||||
Conv2d,
|
||||
Conv3d,
|
||||
ConvTranspose1d,
|
||||
ConvTranspose2d,
|
||||
ConvTranspose3d,
|
||||
)
|
||||
from .dropout import Dropout
|
||||
from .embedding_ops import Embedding, EmbeddingBag
|
||||
from .functional_modules import FloatFunctional, FXFloatFunctional, QFunctional
|
||||
from .linear import Linear
|
||||
from .normalization import (
|
||||
GroupNorm,
|
||||
InstanceNorm1d,
|
||||
InstanceNorm2d,
|
||||
InstanceNorm3d,
|
||||
LayerNorm,
|
||||
)
|
||||
from .rnn import LSTM
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BatchNorm2d",
|
||||
"BatchNorm3d",
|
||||
"Conv1d",
|
||||
"Conv2d",
|
||||
"Conv3d",
|
||||
"ConvTranspose1d",
|
||||
"ConvTranspose2d",
|
||||
"ConvTranspose3d",
|
||||
"DeQuantize",
|
||||
"ELU",
|
||||
"Embedding",
|
||||
"EmbeddingBag",
|
||||
"GroupNorm",
|
||||
"Hardswish",
|
||||
"InstanceNorm1d",
|
||||
"InstanceNorm2d",
|
||||
"InstanceNorm3d",
|
||||
"LayerNorm",
|
||||
"LeakyReLU",
|
||||
"Linear",
|
||||
"LSTM",
|
||||
"MultiheadAttention",
|
||||
"Quantize",
|
||||
"ReLU6",
|
||||
"Sigmoid",
|
||||
"Softmax",
|
||||
"Dropout",
|
||||
"PReLU",
|
||||
# Wrapper modules
|
||||
"FloatFunctional",
|
||||
"FXFloatFunctional",
|
||||
"QFunctional",
|
||||
]
|
||||
|
||||
|
||||
class Quantize(torch.nn.Module):
|
||||
r"""Quantizes an incoming tensor
|
||||
|
||||
Args:
|
||||
`scale`: scale of the output Quantized Tensor
|
||||
`zero_point`: zero_point of output Quantized Tensor
|
||||
`dtype`: data type of output Quantized Tensor
|
||||
`factory_kwargs`: Dictionary of kwargs used for configuring initialization
|
||||
of internal buffers. Currently, `device` and `dtype` are supported.
|
||||
Example: `factory_kwargs={'device': 'cuda', 'dtype': torch.float64}`
|
||||
will initialize internal buffers as type `torch.float64` on the current CUDA device.
|
||||
Note that `dtype` only applies to floating-point buffers.
|
||||
|
||||
Examples::
|
||||
>>> t = torch.tensor([[1., -1.], [1., -1.]])
|
||||
>>> scale, zero_point, dtype = 1.0, 2, torch.qint8
|
||||
>>> qm = Quantize(scale, zero_point, dtype)
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> qt = qm(t)
|
||||
>>> print(qt)
|
||||
tensor([[ 1., -1.],
|
||||
[ 1., -1.]], size=(2, 2), dtype=torch.qint8, scale=1.0, zero_point=2)
|
||||
"""
|
||||
|
||||
scale: torch.Tensor
|
||||
zero_point: torch.Tensor
|
||||
|
||||
def __init__(self, scale, zero_point, dtype, factory_kwargs=None):
|
||||
factory_kwargs = torch.nn.factory_kwargs(factory_kwargs)
|
||||
super().__init__()
|
||||
self.register_buffer("scale", torch.tensor([scale], **factory_kwargs))
|
||||
self.register_buffer(
|
||||
"zero_point",
|
||||
torch.tensor(
|
||||
[zero_point],
|
||||
dtype=torch.long,
|
||||
**{k: v for k, v in factory_kwargs.items() if k != "dtype"},
|
||||
),
|
||||
)
|
||||
self.dtype = dtype
|
||||
|
||||
def forward(self, X):
|
||||
return torch.quantize_per_tensor(
|
||||
X, float(self.scale), int(self.zero_point), self.dtype
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def from_float(mod, use_precomputed_fake_quant=False):
|
||||
if not hasattr(mod, "activation_post_process"):
|
||||
raise AssertionError(
|
||||
f"Module {type(mod).__name__} must have activation_post_process attribute"
|
||||
)
|
||||
scale, zero_point = mod.activation_post_process.calculate_qparams()
|
||||
return Quantize(
|
||||
scale.float().item(),
|
||||
zero_point.long().item(),
|
||||
mod.activation_post_process.dtype,
|
||||
)
|
||||
|
||||
def extra_repr(self):
|
||||
return f"scale={self.scale}, zero_point={self.zero_point}, dtype={self.dtype}"
|
||||
|
||||
|
||||
class DeQuantize(torch.nn.Module):
|
||||
r"""Dequantizes an incoming tensor
|
||||
|
||||
Examples::
|
||||
>>> input = torch.tensor([[1., -1.], [1., -1.]])
|
||||
>>> scale, zero_point, dtype = 1.0, 2, torch.qint8
|
||||
>>> qm = Quantize(scale, zero_point, dtype)
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> quantized_input = qm(input)
|
||||
>>> dqm = DeQuantize()
|
||||
>>> dequantized = dqm(quantized_input)
|
||||
>>> print(dequantized)
|
||||
tensor([[ 1., -1.],
|
||||
[ 1., -1.]], dtype=torch.float32)
|
||||
"""
|
||||
|
||||
def forward(self, Xq):
|
||||
return Xq.dequantize()
|
||||
|
||||
@staticmethod
|
||||
def from_float(mod, use_precomputed_fake_quant=False):
|
||||
return DeQuantize()
|
||||
@@ -0,0 +1,346 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from warnings import warn
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ReLU6",
|
||||
"Hardswish",
|
||||
"ELU",
|
||||
"LeakyReLU",
|
||||
"Sigmoid",
|
||||
"Softmax",
|
||||
"MultiheadAttention",
|
||||
"PReLU",
|
||||
]
|
||||
|
||||
|
||||
class ReLU6(torch.nn.ReLU):
|
||||
r"""Applies the element-wise function:
|
||||
|
||||
:math:`\text{ReLU6}(x) = \min(\max(x_0, x), q(6))`, where :math:`x_0` is the
|
||||
zero_point, and :math:`q(6)` is the quantized representation of number 6.
|
||||
|
||||
Args:
|
||||
inplace: can optionally do the operation in-place. Default: ``False``
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, *)` where `*` means, any number of additional
|
||||
dimensions
|
||||
- Output: :math:`(N, *)`, same shape as the input
|
||||
|
||||
.. image:: ../scripts/activation_images/ReLU6.png
|
||||
|
||||
Examples::
|
||||
|
||||
>>> m = nn.quantized.ReLU6()
|
||||
>>> input = torch.randn(2)
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> input = torch.quantize_per_tensor(input, 1.0, 0, dtype=torch.qint32)
|
||||
>>> output = m(input)
|
||||
"""
|
||||
|
||||
def __init__(self, inplace=False):
|
||||
super().__init__(inplace)
|
||||
self.inplace = inplace
|
||||
|
||||
def forward(self, input):
|
||||
return torch.ops.quantized.relu6(input, self.inplace)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedReLU6"
|
||||
|
||||
@staticmethod
|
||||
def from_float(mod, use_precomputed_fake_quant=False):
|
||||
return ReLU6(mod.inplace)
|
||||
|
||||
|
||||
class Hardswish(torch.nn.Hardswish):
|
||||
r"""This is the quantized version of :class:`~torch.nn.Hardswish`.
|
||||
|
||||
Args:
|
||||
scale: quantization scale of the output tensor
|
||||
zero_point: quantization zero point of the output tensor
|
||||
"""
|
||||
|
||||
def __init__(self, scale, zero_point, device=None, dtype=None):
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
self.register_buffer("scale", torch.tensor(scale, **factory_kwargs))
|
||||
self.register_buffer("zero_point", torch.tensor(zero_point, **factory_kwargs))
|
||||
|
||||
def forward(self, input):
|
||||
return torch.ops.quantized.hardswish(input, self.scale, self.zero_point)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedHardswish"
|
||||
|
||||
@staticmethod
|
||||
def from_float(mod, use_precomputed_fake_quant=False):
|
||||
scale, zero_point = mod.activation_post_process.calculate_qparams()
|
||||
return Hardswish(float(scale), int(zero_point))
|
||||
|
||||
@classmethod
|
||||
def from_reference(cls, mod, scale, zero_point):
|
||||
return cls(float(scale), int(zero_point))
|
||||
|
||||
|
||||
class ELU(torch.nn.ELU):
|
||||
r"""This is the quantized equivalent of :class:`~torch.nn.ELU`.
|
||||
|
||||
Args:
|
||||
scale: quantization scale of the output tensor
|
||||
zero_point: quantization zero point of the output tensor
|
||||
alpha: the alpha constant
|
||||
"""
|
||||
|
||||
def __init__(self, scale, zero_point, alpha=1.0):
|
||||
super().__init__(alpha)
|
||||
self.scale = scale
|
||||
self.zero_point = zero_point
|
||||
|
||||
def forward(self, input):
|
||||
return torch.ao.nn.quantized.functional.elu(
|
||||
input, self.scale, self.zero_point, self.alpha
|
||||
)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedELU"
|
||||
|
||||
@staticmethod
|
||||
def from_float(mod, use_precomputed_fake_quant=False):
|
||||
scale, zero_point = mod.activation_post_process.calculate_qparams()
|
||||
return ELU(float(scale), int(zero_point), mod.alpha)
|
||||
|
||||
@classmethod
|
||||
def from_reference(cls, mod, scale, zero_point):
|
||||
return cls(float(scale), int(zero_point), mod.alpha)
|
||||
|
||||
|
||||
class LeakyReLU(torch.nn.LeakyReLU):
|
||||
r"""This is the quantized equivalent of :class:`~torch.nn.LeakyReLU`.
|
||||
|
||||
Args:
|
||||
scale: quantization scale of the output tensor
|
||||
zero_point: quantization zero point of the output tensor
|
||||
negative_slope: Controls the angle of the negative slope. Default: 1e-2
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
scale: float,
|
||||
zero_point: int,
|
||||
negative_slope: float = 1e-2,
|
||||
inplace: bool = False,
|
||||
device=None,
|
||||
dtype=None,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__(negative_slope, inplace)
|
||||
self.register_buffer("scale", torch.tensor(scale, **factory_kwargs))
|
||||
self.register_buffer("zero_point", torch.tensor(zero_point, **factory_kwargs))
|
||||
|
||||
def forward(self, input):
|
||||
return torch.ops.quantized.leaky_relu(
|
||||
input, self.negative_slope, self.inplace, self.scale, self.zero_point
|
||||
)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedLeakyReLU"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False):
|
||||
scale, zero_point = mod.activation_post_process.calculate_qparams()
|
||||
return cls(float(scale), int(zero_point), mod.negative_slope, mod.inplace)
|
||||
|
||||
@classmethod
|
||||
def from_reference(cls, mod, scale, zero_point):
|
||||
return cls(float(scale), int(zero_point), mod.negative_slope, mod.inplace)
|
||||
|
||||
|
||||
class Sigmoid(torch.nn.Sigmoid):
|
||||
r"""This is the quantized equivalent of :class:`~torch.nn.Sigmoid`.
|
||||
|
||||
Args:
|
||||
scale: quantization scale of the output tensor
|
||||
zero_point: quantization zero point of the output tensor
|
||||
"""
|
||||
|
||||
def __init__(self, output_scale: float, output_zero_point: int):
|
||||
super().__init__()
|
||||
self.output_scale = output_scale
|
||||
self.output_zero_point = output_zero_point
|
||||
|
||||
def forward(self, input):
|
||||
return torch.ops.quantized.sigmoid(
|
||||
input, self.output_scale, self.output_zero_point
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False):
|
||||
(
|
||||
output_scale,
|
||||
output_zero_point,
|
||||
) = mod.activation_post_process.calculate_qparams()
|
||||
return cls(float(output_scale), int(output_zero_point))
|
||||
|
||||
|
||||
class Softmax(torch.nn.Softmax):
|
||||
r"""This is the quantized version of :class:`~torch.nn.Softmax`.
|
||||
|
||||
Args:
|
||||
dim: A dimension along which Softmax will be computed (so every slice along dim will sum to 1).
|
||||
scale: quantization scale of the output tensor
|
||||
zero_point: quantization zero point of the output tensor
|
||||
"""
|
||||
|
||||
def __init__(self, dim=None, scale=1.0, zero_point=0):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.scale = scale
|
||||
self.zero_point = zero_point
|
||||
|
||||
def forward(self, input):
|
||||
dim = self.dim
|
||||
if dim is None:
|
||||
stacklevel = 3
|
||||
# Note: adding the mypy ignore on _get_softmax_dim seems less bad
|
||||
# than making `_get_softmax_dim` an official API.
|
||||
dim = torch.nn.functional._get_softmax_dim( # type: ignore[attr-defined]
|
||||
"softmax", input.dim(), stacklevel
|
||||
)
|
||||
return torch.ops.quantized.softmax(input, dim, self.scale, self.zero_point)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedSoftmax"
|
||||
|
||||
@staticmethod
|
||||
def from_float(mod, use_precomputed_fake_quant=False):
|
||||
scale, zero_point = mod.activation_post_process.calculate_qparams()
|
||||
return Softmax(mod.dim, float(scale), int(zero_point))
|
||||
|
||||
@classmethod
|
||||
def from_reference(cls, mod, scale, zero_point):
|
||||
return cls(mod.dim, float(scale), int(zero_point))
|
||||
|
||||
|
||||
class MultiheadAttention(torch.ao.nn.quantizable.MultiheadAttention):
|
||||
_FLOAT_MODULE = torch.ao.nn.quantizable.MultiheadAttention
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedMultiheadAttention"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, other):
|
||||
# The whole flow is float -> observed -> quantized
|
||||
# This class does observed -> quantized only
|
||||
raise NotImplementedError(
|
||||
"It looks like you are trying to convert a "
|
||||
"non-observed MHA module. Please, see "
|
||||
"the examples on quantizable MHAs."
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_observed(cls, other):
|
||||
converted = torch.ao.quantization.convert(
|
||||
other,
|
||||
mapping=None,
|
||||
inplace=False,
|
||||
remove_qconfig=True,
|
||||
convert_custom_config_dict=None,
|
||||
)
|
||||
converted.__class__ = cls
|
||||
# Remove the parameters for the bias_k and bias_v to quantize them
|
||||
# TODO: This is a potential source of accuracy drop.
|
||||
# quantized cat takes the scale and zp of the first
|
||||
# element, which might lose the precision in the bias_k
|
||||
# and the bias_v (which are cat'ed with k/v being first).
|
||||
if converted.bias_k is not None:
|
||||
bias_k = converted._parameters.pop("bias_k")
|
||||
sc, zp = torch._choose_qparams_per_tensor(bias_k, reduce_range=False)
|
||||
bias_k = torch.quantize_per_tensor(bias_k, sc, zp, torch.quint8)
|
||||
setattr(converted, "bias_k", bias_k) # noqa: B010
|
||||
|
||||
if converted.bias_v is not None:
|
||||
bias_v = converted._parameters.pop("bias_v")
|
||||
sc, zp = torch._choose_qparams_per_tensor(
|
||||
bias_k, # type: ignore[possibly-undefined]
|
||||
reduce_range=False,
|
||||
)
|
||||
bias_v = torch.quantize_per_tensor(bias_v, sc, zp, torch.quint8)
|
||||
setattr(converted, "bias_v", bias_v) # noqa: B010
|
||||
|
||||
del converted.in_proj_weight
|
||||
del converted.in_proj_bias
|
||||
|
||||
return converted
|
||||
|
||||
|
||||
class PReLU(torch.nn.Module):
|
||||
r"""This is the quantized equivalent of :class:`~torch.nn.PReLU`.
|
||||
|
||||
Args:
|
||||
scale: quantization scale of the output tensor
|
||||
zero_point: quantization zero point of the output tensor
|
||||
num_parameters: number of parameters: 1, or the number of channels at input. Default: 1
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, output_scale: float, output_zero_point: int, num_parameters: int = 1
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.num_parameters = num_parameters
|
||||
self.scale = output_scale
|
||||
self.zero_point = output_zero_point
|
||||
w = torch.randn(num_parameters, dtype=torch.float)
|
||||
qw = torch.quantize_per_tensor(w, scale=1.0, zero_point=0, dtype=torch.quint8)
|
||||
self.set_weight(qw)
|
||||
|
||||
def set_weight(self, w: torch.Tensor) -> None:
|
||||
self.weight = w
|
||||
|
||||
def forward(self, input: torch.Tensor) -> torch.Tensor:
|
||||
return torch.ops.quantized.prelu(
|
||||
input, self.weight, self.scale, self.zero_point
|
||||
)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedPReLU"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False):
|
||||
scale, zero_point = mod.activation_post_process.calculate_qparams()
|
||||
qprelu = cls(float(scale), int(zero_point), mod.num_parameters)
|
||||
float_wt = mod.weight.float()
|
||||
observer = mod.qconfig.weight()
|
||||
observer(float_wt)
|
||||
if observer.dtype != torch.quint8:
|
||||
warn(
|
||||
f"PReLU's weight observer should have dtype quint8 but got {observer.dtype}",
|
||||
stacklevel=2,
|
||||
)
|
||||
wt_scale, wt_zp = observer.calculate_qparams()
|
||||
qweight = torch.quantize_per_tensor(
|
||||
float_wt, float(wt_scale), int(wt_zp), torch.quint8
|
||||
)
|
||||
qprelu.set_weight(qweight)
|
||||
return qprelu
|
||||
|
||||
@classmethod
|
||||
def from_reference(cls, mod, scale, zero_point):
|
||||
qprelu = cls(float(scale), int(zero_point), mod.num_parameters)
|
||||
float_wt = mod.weight.float()
|
||||
observer = mod.qconfig.weight()
|
||||
observer(float_wt)
|
||||
if observer.dtype != torch.quint8:
|
||||
warn(
|
||||
f"PReLU's weight observer should have dtype quint8 but got {observer.dtype}",
|
||||
stacklevel=2,
|
||||
)
|
||||
wt_scale, wt_zp = observer.calculate_qparams()
|
||||
qweight = torch.quantize_per_tensor(
|
||||
float_wt, float(wt_scale), int(wt_zp), torch.quint8
|
||||
)
|
||||
qprelu.set_weight(qweight)
|
||||
return qprelu
|
||||
@@ -0,0 +1,128 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import torch
|
||||
import torch.ao.nn.intrinsic as nni
|
||||
|
||||
|
||||
__all__ = ["BatchNorm2d", "BatchNorm3d"]
|
||||
|
||||
|
||||
class _BatchNorm(torch.nn.modules.batchnorm._BatchNorm):
|
||||
def __init__(
|
||||
self, num_features, eps=1e-5, momentum=0.1, device=None, dtype=None
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__(num_features, eps, momentum, True, True, **factory_kwargs)
|
||||
self.register_buffer("scale", torch.tensor(1.0, **factory_kwargs))
|
||||
self.register_buffer("zero_point", torch.tensor(0, **factory_kwargs))
|
||||
|
||||
@staticmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False):
|
||||
activation_post_process = mod.activation_post_process
|
||||
if type(mod) is cls._NNI_BN_RELU_MODULE:
|
||||
mod = mod[0]
|
||||
scale, zero_point = activation_post_process.calculate_qparams()
|
||||
new_mod = cls(mod.num_features, mod.eps)
|
||||
new_mod.weight = mod.weight
|
||||
new_mod.bias = mod.bias
|
||||
new_mod.running_mean = mod.running_mean
|
||||
new_mod.running_var = mod.running_var
|
||||
new_mod.scale = scale
|
||||
new_mod.zero_point = zero_point
|
||||
return new_mod
|
||||
|
||||
@classmethod
|
||||
def from_reference(cls, bn, output_scale, output_zero_point):
|
||||
qbn = cls(
|
||||
bn.num_features,
|
||||
bn.eps,
|
||||
bn.momentum,
|
||||
device=bn.weight.device,
|
||||
dtype=bn.weight.dtype,
|
||||
)
|
||||
qbn.weight = bn.weight
|
||||
qbn.bias = bn.bias
|
||||
qbn.running_mean = bn.running_mean
|
||||
qbn.running_var = bn.running_var
|
||||
qbn.scale = output_scale
|
||||
qbn.zero_point = output_zero_point
|
||||
return qbn
|
||||
|
||||
|
||||
class BatchNorm2d(_BatchNorm):
|
||||
r"""This is the quantized version of :class:`~torch.nn.BatchNorm2d`."""
|
||||
|
||||
_NNI_BN_RELU_MODULE = nni.BNReLU2d
|
||||
|
||||
def __init__(
|
||||
self, num_features, eps=1e-5, momentum=0.1, device=None, dtype=None
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__(num_features, eps, momentum, **factory_kwargs)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedBatchNorm2d"
|
||||
|
||||
def _check_input_dim(self, input):
|
||||
# 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)`!")
|
||||
|
||||
def forward(self, input: torch.Tensor) -> torch.Tensor:
|
||||
# disabling this since this is not symbolically traceable
|
||||
# self._check_input_dim(input)
|
||||
return torch.ops.quantized.batch_norm2d(
|
||||
input,
|
||||
self.weight,
|
||||
self.bias,
|
||||
self.running_mean,
|
||||
self.running_var,
|
||||
self.eps,
|
||||
self.scale,
|
||||
self.zero_point,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False): # type: ignore[override]
|
||||
return _BatchNorm.from_float(
|
||||
cls, mod, use_precomputed_fake_quant=use_precomputed_fake_quant
|
||||
)
|
||||
|
||||
|
||||
class BatchNorm3d(_BatchNorm):
|
||||
r"""This is the quantized version of :class:`~torch.nn.BatchNorm3d`."""
|
||||
|
||||
_NNI_BN_RELU_MODULE = nni.BNReLU3d
|
||||
|
||||
def __init__(self, num_features, eps=1e-5, momentum=0.1, device=None, dtype=None):
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__(num_features, eps, momentum, **factory_kwargs)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedBatchNorm3d"
|
||||
|
||||
def _check_input_dim(self, input):
|
||||
# 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, H, W)`!")
|
||||
|
||||
def forward(self, input: torch.Tensor) -> torch.Tensor:
|
||||
# disabling this since this is not symbolically traceable
|
||||
# self._check_input_dim(input)
|
||||
return torch.ops.quantized.batch_norm3d(
|
||||
input,
|
||||
self.weight,
|
||||
self.bias,
|
||||
self.running_mean,
|
||||
self.running_var,
|
||||
self.eps,
|
||||
self.scale,
|
||||
self.zero_point,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False): # type: ignore[override]
|
||||
return _BatchNorm.from_float(
|
||||
cls, mod, use_precomputed_fake_quant=use_precomputed_fake_quant
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,30 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import torch
|
||||
|
||||
|
||||
__all__ = ["Dropout"]
|
||||
|
||||
|
||||
class Dropout(torch.nn.Dropout):
|
||||
r"""This is the quantized equivalent of :class:`~torch.nn.Dropout`.
|
||||
And this is a placeholder to enable models where fp32 tensors
|
||||
had dropout to work with quantized tensors in train and eval mode.
|
||||
|
||||
Args:
|
||||
p: probability of an element to be zeroed
|
||||
inplace: can optionally do the operation in-place. Default: ``False``
|
||||
"""
|
||||
|
||||
def forward(self, input):
|
||||
return input
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedDropout"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False):
|
||||
return cls(mod.p, mod.inplace)
|
||||
|
||||
@classmethod
|
||||
def from_reference(cls, mod, scale, zero_point):
|
||||
return cls(mod.p, mod.inplace)
|
||||
+424
@@ -0,0 +1,424 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch import Tensor # noqa: F401
|
||||
from torch._jit_internal import List, Optional # noqa: F401
|
||||
|
||||
from .utils import _hide_packed_params_repr, _quantize_weight
|
||||
|
||||
|
||||
__all__ = ["EmbeddingPackedParams", "Embedding", "EmbeddingBag"]
|
||||
|
||||
|
||||
class EmbeddingPackedParams(torch.nn.Module):
|
||||
_version = 1
|
||||
|
||||
def __init__(self, num_embeddings, embedding_dim, dtype=torch.quint8):
|
||||
super().__init__()
|
||||
self.dtype = dtype
|
||||
if self.dtype in [torch.quint8, torch.quint4x2]:
|
||||
scales = torch.ones(num_embeddings, dtype=torch.float)
|
||||
zero_points = torch.zeros(num_embeddings, dtype=torch.float)
|
||||
wq = torch._empty_per_channel_affine_quantized(
|
||||
[num_embeddings, embedding_dim],
|
||||
scales=scales,
|
||||
zero_points=zero_points,
|
||||
axis=0,
|
||||
dtype=self.dtype,
|
||||
)
|
||||
self.set_weight(wq)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Unsupported dtype on quantized embedding! Supports quint8 and quint4x2. Got dtype: {dtype}"
|
||||
)
|
||||
|
||||
@torch.jit.export
|
||||
def set_weight(self, weight: torch.Tensor) -> None:
|
||||
if self.dtype in [torch.quint8, torch.quint4x2]:
|
||||
self._packed_weight = torch.ops.quantized.embedding_bag_prepack(weight)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
"Unsupported dtype for quantized embedding prepack! Supports quint8 and quint4x2."
|
||||
)
|
||||
|
||||
@torch.jit.export
|
||||
def _weight(self):
|
||||
if self.dtype in [torch.quint8, torch.quint4x2]:
|
||||
return torch.ops.quantized.embedding_bag_unpack(self._packed_weight)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
"Unsupported dtype for quantized embedding unpack! Supports quint8 and quint4x2."
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return x
|
||||
|
||||
# Version 1
|
||||
# self
|
||||
# |--- _packed_weight : Tensor representing weight of EmbeddingPackedParamsBase
|
||||
# |--- dtype : torch.dtype
|
||||
|
||||
def _save_to_state_dict(self, destination, prefix, keep_vars):
|
||||
super()._save_to_state_dict(destination, prefix, keep_vars)
|
||||
destination[prefix + "dtype"] = self.dtype
|
||||
destination[prefix + "_packed_weight"] = self._weight()
|
||||
|
||||
def _load_from_state_dict(
|
||||
self,
|
||||
state_dict,
|
||||
prefix,
|
||||
local_metadata,
|
||||
strict,
|
||||
missing_keys,
|
||||
unexpected_keys,
|
||||
error_msgs,
|
||||
):
|
||||
self.dtype = state_dict[prefix + "dtype"]
|
||||
state_dict.pop(prefix + "dtype")
|
||||
|
||||
weight = state_dict[prefix + "_packed_weight"]
|
||||
state_dict.pop(prefix + "_packed_weight")
|
||||
self.set_weight(weight)
|
||||
|
||||
super()._load_from_state_dict(
|
||||
state_dict,
|
||||
prefix,
|
||||
local_metadata,
|
||||
False,
|
||||
missing_keys,
|
||||
unexpected_keys,
|
||||
error_msgs,
|
||||
)
|
||||
|
||||
def __repr__(self):
|
||||
return self._weight().__repr__()
|
||||
|
||||
|
||||
class Embedding(torch.nn.Module):
|
||||
r"""
|
||||
A quantized Embedding module with quantized packed weights as inputs.
|
||||
We adopt the same interface as `torch.nn.Embedding`, please see
|
||||
https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html for documentation.
|
||||
|
||||
Similar to :class:`~torch.nn.Embedding`, attributes will be randomly
|
||||
initialized at module creation time and will be overwritten later
|
||||
|
||||
Attributes:
|
||||
weight (Tensor): the non-learnable quantized weights of the module of
|
||||
shape :math:`(\text{num\_embeddings}, \text{embedding\_dim})`.
|
||||
|
||||
Examples::
|
||||
>>> m = nn.quantized.Embedding(num_embeddings=10, embedding_dim=12)
|
||||
>>> indices = torch.tensor([9, 6, 5, 7, 8, 8, 9, 2, 8])
|
||||
>>> output = m(indices)
|
||||
>>> print(output.size())
|
||||
torch.Size([9, 12])
|
||||
|
||||
"""
|
||||
|
||||
_version = 1
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_embeddings: int,
|
||||
embedding_dim: int,
|
||||
padding_idx: Optional[int] = None,
|
||||
max_norm: Optional[float] = None,
|
||||
norm_type: float = 2.0,
|
||||
scale_grad_by_freq: bool = False,
|
||||
sparse: bool = False,
|
||||
_weight: Optional[Tensor] = None,
|
||||
dtype=torch.quint8,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.num_embeddings = num_embeddings
|
||||
self.embedding_dim = embedding_dim
|
||||
self.dtype = dtype
|
||||
|
||||
if _weight is None:
|
||||
scales = torch.ones(num_embeddings, dtype=torch.float)
|
||||
zero_points = torch.zeros(num_embeddings, dtype=torch.float)
|
||||
qweight = torch._empty_per_channel_affine_quantized(
|
||||
[num_embeddings, embedding_dim],
|
||||
scales=scales,
|
||||
zero_points=zero_points,
|
||||
axis=0,
|
||||
dtype=torch.quint8,
|
||||
)
|
||||
else:
|
||||
expected_shape = [num_embeddings, embedding_dim]
|
||||
if list(_weight.shape) != expected_shape:
|
||||
raise AssertionError(
|
||||
f"Shape of weight does not match num_embeddings and embedding_dim: "
|
||||
f"expected {expected_shape}, got {list(_weight.shape)}"
|
||||
)
|
||||
qweight = _weight
|
||||
|
||||
self._packed_params = EmbeddingPackedParams(
|
||||
num_embeddings, embedding_dim, dtype
|
||||
)
|
||||
self._packed_params.set_weight(qweight)
|
||||
|
||||
def forward(self, indices: Tensor) -> Tensor:
|
||||
if self.dtype == torch.quint4x2:
|
||||
return torch.ops.quantized.embedding_4bit(
|
||||
self._packed_params._packed_weight, indices
|
||||
)
|
||||
else:
|
||||
return torch.ops.quantized.embedding_byte(
|
||||
self._packed_params._packed_weight, indices
|
||||
)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedEmbedding"
|
||||
|
||||
def __repr__(self):
|
||||
return _hide_packed_params_repr(self, EmbeddingPackedParams)
|
||||
|
||||
def extra_repr(self):
|
||||
extra_repr_str = (
|
||||
f"num_embeddings={self.num_embeddings}, embedding_dim={self.embedding_dim}, "
|
||||
f"dtype={self._packed_params.dtype}, qscheme={self.weight().qscheme()}"
|
||||
)
|
||||
|
||||
return extra_repr_str
|
||||
|
||||
def set_weight(self, w: torch.Tensor) -> None:
|
||||
self._packed_params.set_weight(w)
|
||||
|
||||
def weight(self):
|
||||
return self._packed_params._weight()
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False):
|
||||
r"""Create a quantized embedding module from a float module
|
||||
|
||||
Args:
|
||||
mod (Module): a float module, either produced by torch.ao.quantization
|
||||
utilities or provided by user
|
||||
"""
|
||||
if hasattr(mod, "weight_fake_quant"):
|
||||
if type(mod) is not torch.ao.nn.qat.Embedding:
|
||||
raise AssertionError(
|
||||
"nnq."
|
||||
+ cls.__name__
|
||||
+ ".from_float "
|
||||
+ "with fake quant only works for "
|
||||
+ torch.ao.nn.qat.Embedding.__name__
|
||||
)
|
||||
weight_observer = mod.weight_fake_quant
|
||||
else:
|
||||
if type(mod) is not nn.Embedding:
|
||||
raise AssertionError(
|
||||
"nnq."
|
||||
+ cls.__name__
|
||||
+ ".from_float only works for "
|
||||
+ nn.Embedding.__name__
|
||||
)
|
||||
if not hasattr(mod, "qconfig"):
|
||||
raise AssertionError(
|
||||
"Embedding input float module must have qconfig defined"
|
||||
)
|
||||
from torch.ao.quantization import float_qparams_weight_only_qconfig
|
||||
|
||||
if mod.qconfig is not None and mod.qconfig.weight is not None: # type: ignore[union-attr]
|
||||
weight_observer = mod.qconfig.weight() # type: ignore[union-attr, operator]
|
||||
else:
|
||||
weight_observer = float_qparams_weight_only_qconfig.weight()
|
||||
|
||||
dtype = weight_observer.dtype
|
||||
is_float_qparams_qconfig = (
|
||||
weight_observer.qscheme == torch.per_channel_affine_float_qparams
|
||||
)
|
||||
if not is_float_qparams_qconfig:
|
||||
raise AssertionError(
|
||||
"Embedding quantization is only supported with float_qparams_weight_only_qconfig."
|
||||
)
|
||||
|
||||
if dtype != torch.quint8 and dtype != torch.quint4x2:
|
||||
raise AssertionError(
|
||||
f"The only supported dtype for nnq.Embedding is torch.quint8 and torch.quint4x2, got {dtype}"
|
||||
)
|
||||
|
||||
# Run the observer to calculate qparams.
|
||||
weight_observer(mod.weight)
|
||||
qweight = _quantize_weight(mod.weight.float(), weight_observer)
|
||||
|
||||
# Create quantized Embedding module and pass in the quantized weight
|
||||
qembedding = Embedding(mod.num_embeddings, mod.embedding_dim)
|
||||
qembedding.set_weight(qweight)
|
||||
return qembedding
|
||||
|
||||
@classmethod
|
||||
def from_reference(cls, ref_embedding):
|
||||
qembedding = cls(
|
||||
ref_embedding.num_embeddings,
|
||||
ref_embedding.embedding_dim,
|
||||
ref_embedding.padding_idx,
|
||||
ref_embedding.max_norm,
|
||||
ref_embedding.norm_type,
|
||||
ref_embedding.scale_grad_by_freq,
|
||||
ref_embedding.sparse,
|
||||
ref_embedding.get_quantized_weight(),
|
||||
ref_embedding.weight_dtype,
|
||||
)
|
||||
return qembedding
|
||||
|
||||
|
||||
class EmbeddingBag(Embedding):
|
||||
r"""
|
||||
A quantized EmbeddingBag module with quantized packed weights as inputs.
|
||||
We adopt the same interface as `torch.nn.EmbeddingBag`, please see
|
||||
https://pytorch.org/docs/stable/generated/torch.nn.EmbeddingBag.html for documentation.
|
||||
|
||||
Similar to :class:`~torch.nn.EmbeddingBag`, attributes will be randomly
|
||||
initialized at module creation time and will be overwritten later
|
||||
|
||||
Attributes:
|
||||
weight (Tensor): the non-learnable quantized weights of the module of
|
||||
shape :math:`(\text{num\_embeddings}, \text{embedding\_dim})`.
|
||||
|
||||
Examples::
|
||||
>>> m = nn.quantized.EmbeddingBag(num_embeddings=10, embedding_dim=12, include_last_offset=True, mode='sum')
|
||||
>>> indices = torch.tensor([9, 6, 5, 7, 8, 8, 9, 2, 8, 6, 6, 9, 1, 6, 8, 8, 3, 2, 3, 6, 3, 6, 5, 7, 0, 8, 4, 6, 5, 8, 2, 3])
|
||||
>>> offsets = torch.tensor([0, 19, 20, 28, 28, 32])
|
||||
>>> output = m(indices, offsets)
|
||||
>>> print(output.size())
|
||||
torch.Size([5, 12])
|
||||
|
||||
"""
|
||||
|
||||
_version = 1
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_embeddings: int,
|
||||
embedding_dim: int,
|
||||
max_norm: Optional[float] = None,
|
||||
norm_type: float = 2.0,
|
||||
scale_grad_by_freq: bool = False,
|
||||
mode: str = "sum",
|
||||
sparse: bool = False,
|
||||
_weight: Optional[Tensor] = None,
|
||||
include_last_offset: bool = False,
|
||||
dtype=torch.quint8,
|
||||
) -> None:
|
||||
super().__init__(num_embeddings, embedding_dim, _weight=_weight, dtype=dtype)
|
||||
|
||||
self.mode = mode
|
||||
self.pruned_weights = False
|
||||
self.include_last_offset = include_last_offset
|
||||
self.dtype = dtype
|
||||
|
||||
def forward(
|
||||
self,
|
||||
indices: Tensor,
|
||||
offsets: Optional[Tensor] = None,
|
||||
per_sample_weights: Optional[Tensor] = None,
|
||||
compressed_indices_mapping: Optional[Tensor] = None,
|
||||
) -> Tensor:
|
||||
if self.dtype == torch.quint4x2:
|
||||
return torch.ops.quantized.embedding_bag_4bit(
|
||||
self._packed_params._packed_weight,
|
||||
indices,
|
||||
offsets,
|
||||
False,
|
||||
0,
|
||||
self.pruned_weights,
|
||||
per_sample_weights,
|
||||
compressed_indices_mapping,
|
||||
self.include_last_offset,
|
||||
)
|
||||
else:
|
||||
return torch.ops.quantized.embedding_bag_byte(
|
||||
self._packed_params._packed_weight,
|
||||
indices,
|
||||
offsets,
|
||||
False,
|
||||
0,
|
||||
self.pruned_weights,
|
||||
per_sample_weights,
|
||||
compressed_indices_mapping,
|
||||
self.include_last_offset,
|
||||
)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedEmbeddingBag"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False):
|
||||
r"""Create a quantized embedding_bag module from a float module
|
||||
|
||||
Args:
|
||||
mod (Module): a float module, either produced by torch.ao.quantization
|
||||
utilities or provided by user
|
||||
"""
|
||||
if hasattr(mod, "weight_fake_quant"):
|
||||
weight_observer = mod.weight_fake_quant
|
||||
else:
|
||||
if type(mod) is not nn.EmbeddingBag:
|
||||
raise AssertionError(
|
||||
"nnq."
|
||||
+ cls.__name__
|
||||
+ ".from_float only works for "
|
||||
+ nn.EmbeddingBag.__name__
|
||||
)
|
||||
if not hasattr(mod, "qconfig"):
|
||||
raise AssertionError(
|
||||
"EmbeddingBag input float module must have qconfig defined"
|
||||
)
|
||||
from torch.ao.quantization.qconfig import float_qparams_weight_only_qconfig
|
||||
|
||||
if mod.qconfig is not None and mod.qconfig.weight is not None: # type: ignore[union-attr]
|
||||
weight_observer = mod.qconfig.weight() # type: ignore[union-attr, operator]
|
||||
else:
|
||||
weight_observer = float_qparams_weight_only_qconfig.weight()
|
||||
|
||||
dtype = weight_observer.dtype
|
||||
is_float_qparams_qconfig = (
|
||||
weight_observer.qscheme == torch.per_channel_affine_float_qparams
|
||||
)
|
||||
if not is_float_qparams_qconfig:
|
||||
raise AssertionError(
|
||||
"EmbeddingBag quantization is only supported with float_qparams_weight_only_qconfig."
|
||||
)
|
||||
|
||||
if dtype != torch.quint8 and dtype != torch.quint4x2:
|
||||
raise AssertionError(
|
||||
f"The only supported dtype for nnq.EmbeddingBag is torch.quint8 and torch.quint4x2, got {dtype}"
|
||||
)
|
||||
|
||||
# Run the observer to calculate qparams.
|
||||
weight_observer(mod.weight)
|
||||
qweight = _quantize_weight(mod.weight.float(), weight_observer)
|
||||
|
||||
# Create quantized EmbeddingBag module and pass in the quantized weight
|
||||
qembedding_bag = EmbeddingBag(
|
||||
mod.num_embeddings,
|
||||
mod.embedding_dim,
|
||||
max_norm=mod.max_norm,
|
||||
norm_type=mod.norm_type,
|
||||
scale_grad_by_freq=mod.scale_grad_by_freq,
|
||||
mode=mod.mode,
|
||||
sparse=mod.sparse,
|
||||
include_last_offset=mod.include_last_offset,
|
||||
dtype=dtype,
|
||||
)
|
||||
qembedding_bag.set_weight(qweight)
|
||||
return qembedding_bag
|
||||
|
||||
@classmethod
|
||||
def from_reference(cls, ref_embedding_bag):
|
||||
qembedding_bag = cls(
|
||||
ref_embedding_bag.num_embeddings,
|
||||
ref_embedding_bag.embedding_dim,
|
||||
ref_embedding_bag.max_norm,
|
||||
ref_embedding_bag.norm_type,
|
||||
ref_embedding_bag.scale_grad_by_freq,
|
||||
ref_embedding_bag.mode,
|
||||
ref_embedding_bag.sparse,
|
||||
ref_embedding_bag.get_quantized_weight(),
|
||||
ref_embedding_bag.include_last_offset,
|
||||
ref_embedding_bag.weight_dtype,
|
||||
)
|
||||
return qembedding_bag
|
||||
+300
@@ -0,0 +1,300 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch._ops import ops
|
||||
|
||||
|
||||
__all__ = ["FloatFunctional", "FXFloatFunctional", "QFunctional"]
|
||||
|
||||
|
||||
class FloatFunctional(torch.nn.Module):
|
||||
r"""State collector class for float operations.
|
||||
|
||||
The instance of this class can be used instead of the ``torch.`` prefix for
|
||||
some operations. See example usage below.
|
||||
|
||||
.. note::
|
||||
|
||||
This class does not provide a ``forward`` hook. Instead, you must use
|
||||
one of the underlying functions (e.g. ``add``).
|
||||
|
||||
Examples::
|
||||
|
||||
>>> f_add = FloatFunctional()
|
||||
>>> a = torch.tensor(3.0)
|
||||
>>> b = torch.tensor(4.0)
|
||||
>>> f_add.add(a, b) # Equivalent to ``torch.add(a, b)``
|
||||
|
||||
Valid operation names:
|
||||
- add
|
||||
- cat
|
||||
- mul
|
||||
- add_relu
|
||||
- add_scalar
|
||||
- mul_scalar
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.activation_post_process = torch.nn.Identity()
|
||||
|
||||
def forward(self, x):
|
||||
raise RuntimeError(
|
||||
"FloatFunctional is not intended to use the "
|
||||
+ "'forward'. Please use the underlying operation"
|
||||
)
|
||||
|
||||
r"""Operation equivalent to ``torch.add(Tensor, Tensor)``"""
|
||||
|
||||
def add(self, x: Tensor, y: Tensor) -> Tensor:
|
||||
r = torch.add(x, y)
|
||||
r = self.activation_post_process(r)
|
||||
return r
|
||||
|
||||
r"""Operation equivalent to ``torch.add(Tensor, float)``"""
|
||||
|
||||
def add_scalar(self, x: Tensor, y: float) -> Tensor:
|
||||
r = torch.add(x, y)
|
||||
# Note: this operation is not observed because the observation is not
|
||||
# needed for the quantized op.
|
||||
return r
|
||||
|
||||
r"""Operation equivalent to ``torch.mul(Tensor, Tensor)``"""
|
||||
|
||||
def mul(self, x: Tensor, y: Tensor) -> Tensor:
|
||||
r = torch.mul(x, y)
|
||||
r = self.activation_post_process(r)
|
||||
return r
|
||||
|
||||
r"""Operation equivalent to ``torch.mul(Tensor, float)``"""
|
||||
|
||||
def mul_scalar(self, x: Tensor, y: float) -> Tensor:
|
||||
r = torch.mul(x, y)
|
||||
# Note: this operation is not observed because the observation is not
|
||||
# needed for the quantized op.
|
||||
return r
|
||||
|
||||
r"""Operation equivalent to ``torch.cat``"""
|
||||
|
||||
def cat(self, x: list[Tensor], dim: int = 0) -> Tensor:
|
||||
r = torch.cat(x, dim=dim)
|
||||
r = self.activation_post_process(r)
|
||||
return r
|
||||
|
||||
r"""Operation equivalent to ``relu(torch.add(x,y))``"""
|
||||
|
||||
def add_relu(self, x: Tensor, y: Tensor) -> Tensor:
|
||||
r = torch.add(x, y)
|
||||
r = torch.nn.functional.relu(r)
|
||||
r = self.activation_post_process(r)
|
||||
return r
|
||||
|
||||
r"""Operation equivalent to ``torch.matmul(Tensor, Tensor)``"""
|
||||
|
||||
def matmul(self, x: Tensor, y: Tensor) -> Tensor:
|
||||
r = torch.matmul(x, y)
|
||||
r = self.activation_post_process(r)
|
||||
return r
|
||||
|
||||
|
||||
class FXFloatFunctional(torch.nn.Module):
|
||||
r"""module to replace FloatFunctional module before FX graph mode quantization,
|
||||
since activation_post_process will be inserted in top level module directly
|
||||
|
||||
Valid operation names:
|
||||
- add
|
||||
- cat
|
||||
- mul
|
||||
- add_relu
|
||||
- add_scalar
|
||||
- mul_scalar
|
||||
"""
|
||||
|
||||
def forward(self, x):
|
||||
raise RuntimeError(
|
||||
"FloatFunctional is not intended to use the "
|
||||
+ "'forward'. Please use the underlying operation"
|
||||
)
|
||||
|
||||
r"""Operation equivalent to ``torch.add(Tensor, Tensor)``"""
|
||||
|
||||
def add(self, x: Tensor, y: Tensor) -> Tensor:
|
||||
r = torch.add(x, y)
|
||||
return r
|
||||
|
||||
r"""Operation equivalent to ``torch.add(Tensor, float)``"""
|
||||
|
||||
def add_scalar(self, x: Tensor, y: float) -> Tensor:
|
||||
r = torch.add(x, y)
|
||||
return r
|
||||
|
||||
r"""Operation equivalent to ``torch.mul(Tensor, Tensor)``"""
|
||||
|
||||
def mul(self, x: Tensor, y: Tensor) -> Tensor:
|
||||
r = torch.mul(x, y)
|
||||
return r
|
||||
|
||||
r"""Operation equivalent to ``torch.mul(Tensor, float)``"""
|
||||
|
||||
def mul_scalar(self, x: Tensor, y: float) -> Tensor:
|
||||
r = torch.mul(x, y)
|
||||
return r
|
||||
|
||||
r"""Operation equivalent to ``torch.cat``"""
|
||||
|
||||
def cat(self, x: list[Tensor], dim: int = 0) -> Tensor:
|
||||
r = torch.cat(x, dim=dim)
|
||||
return r
|
||||
|
||||
r"""Operation equivalent to ``relu(torch.add(x,y))``"""
|
||||
|
||||
def add_relu(self, x: Tensor, y: Tensor) -> Tensor:
|
||||
r = torch.add(x, y)
|
||||
r = torch.nn.functional.relu(r)
|
||||
return r
|
||||
|
||||
r"""Operation equivalent to ``torch.matmul(Tensor, Tensor)``"""
|
||||
|
||||
def matmul(self, x: Tensor, y: Tensor) -> Tensor:
|
||||
r = torch.matmul(x, y)
|
||||
return r
|
||||
|
||||
|
||||
class QFunctional(torch.nn.Module):
|
||||
r"""Wrapper class for quantized operations.
|
||||
|
||||
The instance of this class can be used instead of the
|
||||
``torch.ops.quantized`` prefix. See example usage below.
|
||||
|
||||
.. note::
|
||||
|
||||
This class does not provide a ``forward`` hook. Instead, you must use
|
||||
one of the underlying functions (e.g. ``add``).
|
||||
|
||||
Examples::
|
||||
|
||||
>>> q_add = QFunctional()
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> a = torch.quantize_per_tensor(torch.tensor(3.0), 1.0, 0, torch.qint32)
|
||||
>>> b = torch.quantize_per_tensor(torch.tensor(4.0), 1.0, 0, torch.qint32)
|
||||
>>> q_add.add(a, b) # Equivalent to ``torch.ops.quantized.add(a, b, 1.0, 0)``
|
||||
|
||||
Valid operation names:
|
||||
- add
|
||||
- cat
|
||||
- mul
|
||||
- add_relu
|
||||
- add_scalar
|
||||
- mul_scalar
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.scale = 1.0
|
||||
self.zero_point = 0
|
||||
self.activation_post_process = torch.nn.Identity()
|
||||
|
||||
def _save_to_state_dict(self, destination, prefix, keep_vars):
|
||||
super()._save_to_state_dict(destination, prefix, keep_vars)
|
||||
destination[prefix + "scale"] = torch.tensor(self.scale)
|
||||
destination[prefix + "zero_point"] = torch.tensor(self.zero_point)
|
||||
|
||||
def _load_from_state_dict(
|
||||
self,
|
||||
state_dict,
|
||||
prefix,
|
||||
local_metadata,
|
||||
strict,
|
||||
missing_keys,
|
||||
unexpected_keys,
|
||||
error_msgs,
|
||||
):
|
||||
self.scale = float(state_dict.pop(prefix + "scale"))
|
||||
self.zero_point = int(state_dict.pop(prefix + "zero_point"))
|
||||
super()._load_from_state_dict(
|
||||
state_dict,
|
||||
prefix,
|
||||
local_metadata,
|
||||
False,
|
||||
missing_keys,
|
||||
unexpected_keys,
|
||||
error_msgs,
|
||||
)
|
||||
|
||||
def _get_name(self):
|
||||
return "QFunctional"
|
||||
|
||||
def extra_repr(self):
|
||||
return f"scale={self.scale}, zero_point={self.zero_point}"
|
||||
|
||||
def forward(self, x):
|
||||
raise RuntimeError(
|
||||
"Functional is not intended to use the "
|
||||
+ "'forward'. Please use the underlying operation"
|
||||
)
|
||||
|
||||
r"""Operation equivalent to ``torch.ops.quantized.add``"""
|
||||
|
||||
def add(self, x: Tensor, y: Tensor) -> Tensor:
|
||||
r = ops.quantized.add(x, y, scale=self.scale, zero_point=self.zero_point)
|
||||
r = self.activation_post_process(r)
|
||||
return r
|
||||
|
||||
r"""Operation equivalent to ``torch.ops.quantized.add(Tensor, float)``"""
|
||||
|
||||
def add_scalar(self, x: Tensor, y: float) -> Tensor:
|
||||
r = ops.quantized.add_scalar(x, y)
|
||||
# Note: this operation is not observed because the observation is not
|
||||
# needed for the quantized op.
|
||||
return r
|
||||
|
||||
r"""Operation equivalent to ``torch.ops.quantized.mul(Tensor, Tensor)``"""
|
||||
|
||||
def mul(self, x: Tensor, y: Tensor) -> Tensor:
|
||||
r = ops.quantized.mul(x, y, scale=self.scale, zero_point=self.zero_point)
|
||||
r = self.activation_post_process(r)
|
||||
return r
|
||||
|
||||
r"""Operation equivalent to ``torch.ops.quantized.mul(Tensor, float)``"""
|
||||
|
||||
def mul_scalar(self, x: Tensor, y: float) -> Tensor:
|
||||
r = ops.quantized.mul_scalar(x, y)
|
||||
# Note: this operation is not observed because the observation is not
|
||||
# needed for the quantized op.
|
||||
return r
|
||||
|
||||
r"""Operation equivalent to ``torch.ops.quantized.cat``"""
|
||||
|
||||
def cat(self, x: list[Tensor], dim: int = 0) -> Tensor:
|
||||
r = ops.quantized.cat(x, scale=self.scale, zero_point=self.zero_point, dim=dim)
|
||||
r = self.activation_post_process(r)
|
||||
return r
|
||||
|
||||
r"""Operation equivalent to ``torch.ops.quantized.add_relu``"""
|
||||
|
||||
def add_relu(self, x: Tensor, y: Tensor) -> Tensor:
|
||||
r = ops.quantized.add_relu(x, y, scale=self.scale, zero_point=self.zero_point)
|
||||
r = self.activation_post_process(r)
|
||||
return r
|
||||
|
||||
r"""Operation equivalent to ``torch.ops.quantized.matmul(Tensor, Tensor)``"""
|
||||
|
||||
def matmul(self, x: Tensor, y: Tensor) -> Tensor:
|
||||
r = ops.quantized.matmul(x, y, scale=self.scale, zero_point=self.zero_point)
|
||||
# Note: this operation is not observed because the observation is not
|
||||
# needed for the quantized op.
|
||||
return r
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False):
|
||||
if type(mod) is not FloatFunctional:
|
||||
raise AssertionError(
|
||||
f"QFunctional.from_float expects an instance of FloatFunctional, "
|
||||
f"got {type(mod).__name__}"
|
||||
)
|
||||
scale, zero_point = mod.activation_post_process.calculate_qparams() # type: ignore[operator]
|
||||
new_mod = QFunctional()
|
||||
new_mod.scale = float(scale)
|
||||
new_mod.zero_point = int(zero_point)
|
||||
return new_mod
|
||||
@@ -0,0 +1,362 @@
|
||||
# mypy: allow-untyped-decorators
|
||||
# mypy: allow-untyped-defs
|
||||
from collections.abc import Iterable
|
||||
|
||||
import torch
|
||||
import torch.ao.nn.intrinsic as nni
|
||||
import torch.ao.nn.intrinsic.qat as nniqat
|
||||
import torch.nn as nn
|
||||
from torch.nn.utils.fusion import fuse_linear_bn_weights
|
||||
from torch.nn.utils.parametrize import type_before_parametrizations
|
||||
|
||||
from .utils import _hide_packed_params_repr, _quantize_weight, WeightedQuantizedModule
|
||||
|
||||
|
||||
__all__ = ["LinearPackedParams", "Linear"]
|
||||
|
||||
|
||||
class LinearPackedParams(torch.nn.Module):
|
||||
_version = 3
|
||||
|
||||
def __init__(self, dtype=torch.qint8):
|
||||
super().__init__()
|
||||
self.dtype = dtype
|
||||
if self.dtype == torch.qint8:
|
||||
wq = torch._empty_affine_quantized(
|
||||
[1, 1], scale=1.0, zero_point=0, dtype=torch.qint8
|
||||
)
|
||||
elif self.dtype == torch.float16:
|
||||
wq = torch.zeros([1, 1], dtype=torch.float)
|
||||
self.set_weight_bias(wq, None) # type: ignore[possibly-undefined]
|
||||
|
||||
@torch.jit.export
|
||||
def set_weight_bias(self, weight: torch.Tensor, bias: torch.Tensor | None) -> None:
|
||||
if self.dtype == torch.qint8:
|
||||
self._packed_params = torch.ops.quantized.linear_prepack(weight, bias)
|
||||
elif self.dtype == torch.float16:
|
||||
self._packed_params = torch.ops.quantized.linear_prepack_fp16(weight, bias)
|
||||
else:
|
||||
raise RuntimeError("Unsupported dtype on dynamic quantized linear!")
|
||||
|
||||
@torch.jit.export
|
||||
def _weight_bias(self):
|
||||
if self.dtype == torch.qint8:
|
||||
return torch.ops.quantized.linear_unpack(self._packed_params)
|
||||
elif self.dtype == torch.float16:
|
||||
return torch.ops.quantized.linear_unpack_fp16(self._packed_params)
|
||||
else:
|
||||
raise RuntimeError("Unsupported dtype on dynamic quantized linear!")
|
||||
|
||||
def forward(self, x):
|
||||
return x
|
||||
|
||||
# Version 1
|
||||
# self
|
||||
# |--- weight : Tensor
|
||||
# |--- bias : Tensor
|
||||
#
|
||||
# Version 2
|
||||
# self
|
||||
# |--- weight : Tensor
|
||||
# |--- bias : Tensor
|
||||
# |--- dtype : torch.dtype
|
||||
#
|
||||
# Version 3
|
||||
# self
|
||||
# |--- _packed_params : (Tensor, Tensor) representing (weight, bias)
|
||||
# of LinearPackedParams
|
||||
# |--- dtype : torch.dtype
|
||||
def _save_to_state_dict(self, destination, prefix, keep_vars):
|
||||
super()._save_to_state_dict(destination, prefix, keep_vars)
|
||||
destination[prefix + "dtype"] = self.dtype
|
||||
destination[prefix + "_packed_params"] = self._weight_bias()
|
||||
|
||||
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 < 2:
|
||||
self.dtype = torch.qint8
|
||||
else:
|
||||
self.dtype = state_dict[prefix + "dtype"]
|
||||
state_dict.pop(prefix + "dtype")
|
||||
|
||||
if version is None or version < 3:
|
||||
self.set_weight_bias(
|
||||
state_dict[prefix + "weight"], state_dict[prefix + "bias"]
|
||||
)
|
||||
state_dict.pop(prefix + "weight")
|
||||
state_dict.pop(prefix + "bias")
|
||||
|
||||
if version == 3:
|
||||
weight, bias = state_dict[prefix + "_packed_params"]
|
||||
state_dict.pop(prefix + "_packed_params")
|
||||
self.set_weight_bias(weight, bias)
|
||||
|
||||
super()._load_from_state_dict(
|
||||
state_dict,
|
||||
prefix,
|
||||
local_metadata,
|
||||
False,
|
||||
missing_keys,
|
||||
unexpected_keys,
|
||||
error_msgs,
|
||||
)
|
||||
|
||||
def __repr__(self):
|
||||
return self._weight_bias().__repr__()
|
||||
|
||||
|
||||
class Linear(WeightedQuantizedModule):
|
||||
r"""
|
||||
A quantized linear module with quantized tensor as inputs and outputs.
|
||||
We adopt the same interface as `torch.nn.Linear`, please see
|
||||
https://pytorch.org/docs/stable/nn.html#torch.nn.Linear for documentation.
|
||||
|
||||
Similar to :class:`~torch.nn.Linear`, attributes will be randomly
|
||||
initialized at module creation time and will be overwritten later
|
||||
|
||||
Attributes:
|
||||
weight (Tensor): the non-learnable quantized weights of the module of
|
||||
shape :math:`(\text{out\_features}, \text{in\_features})`.
|
||||
bias (Tensor): the non-learnable bias of the module of shape :math:`(\text{out\_features})`.
|
||||
If :attr:`bias` is ``True``, the values are initialized to zero.
|
||||
scale: `scale` parameter of output Quantized Tensor, type: double
|
||||
zero_point: `zero_point` parameter for output Quantized Tensor, type: long
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_QENGINE)
|
||||
>>> m = nn.quantized.Linear(20, 30)
|
||||
>>> input = torch.randn(128, 20)
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> input = torch.quantize_per_tensor(input, 1.0, 0, torch.quint8)
|
||||
>>> output = m(input)
|
||||
>>> print(output.size())
|
||||
torch.Size([128, 30])
|
||||
"""
|
||||
|
||||
_version = 3
|
||||
_FLOAT_MODULE = (nn.Linear, nn.modules.linear.NonDynamicallyQuantizableLinear)
|
||||
|
||||
def __init__(self, in_features, out_features, bias_=True, dtype=torch.qint8):
|
||||
super().__init__()
|
||||
# We don't muck around with buffers or attributes or anything here
|
||||
# to keep the module simple. *everything* is simply a Python attribute.
|
||||
# Serialization logic is explicitly handled in the below serialization and
|
||||
# deserialization modules
|
||||
self.in_features = in_features
|
||||
self.out_features = out_features
|
||||
bias = None
|
||||
if bias_:
|
||||
bias = torch.zeros(out_features, dtype=torch.float)
|
||||
|
||||
if dtype == torch.qint8:
|
||||
qweight = torch._empty_affine_quantized(
|
||||
[out_features, in_features], scale=1, zero_point=0, dtype=torch.qint8
|
||||
)
|
||||
elif dtype == torch.float16:
|
||||
qweight = torch.zeros([out_features, in_features], dtype=torch.float)
|
||||
else:
|
||||
raise RuntimeError("Unsupported dtype specified for quantized Linear!")
|
||||
|
||||
self._packed_params = LinearPackedParams(dtype)
|
||||
self._packed_params.set_weight_bias(qweight, bias)
|
||||
self.scale = 1.0
|
||||
self.zero_point = 0
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedLinear"
|
||||
|
||||
def extra_repr(self):
|
||||
return (
|
||||
f"in_features={self.in_features}, out_features={self.out_features}, scale={self.scale}, "
|
||||
f"zero_point={self.zero_point}, qscheme={self.weight().qscheme()}"
|
||||
)
|
||||
|
||||
def __repr__(self):
|
||||
return _hide_packed_params_repr(self, LinearPackedParams)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return torch.ops.quantized.linear(
|
||||
x, self._packed_params._packed_params, self.scale, self.zero_point
|
||||
)
|
||||
|
||||
# ===== Serialization methods =====
|
||||
# The special consideration here is that we have to unpack the weights into their
|
||||
# regular QTensor form for serialization. Packed weights should not live
|
||||
# outside the process in which they were created, rather they should be derived
|
||||
# from the QTensor weight.
|
||||
#
|
||||
# Version 1
|
||||
# self
|
||||
# |--- scale : float
|
||||
# |--- zero_point : int
|
||||
# |--- weight : Tensor
|
||||
# |--- bias : Tensor
|
||||
#
|
||||
# Version 2
|
||||
# self
|
||||
# |--- scale : float
|
||||
# |--- zero_point : int
|
||||
# |--- _packed_params : Module
|
||||
# |--- weight : Tensor
|
||||
# |--- bias : Tensor
|
||||
#
|
||||
# Version 3
|
||||
# self
|
||||
# |--- scale : float
|
||||
# |--- zero_point : int
|
||||
# |--- _packed_params : Module
|
||||
# |--- _packed_params : (Tensor, Tensor) representing weight, bias
|
||||
# of LinearPackedParams C++ struct
|
||||
#
|
||||
def _save_to_state_dict(self, destination, prefix, keep_vars):
|
||||
super()._save_to_state_dict(destination, prefix, keep_vars)
|
||||
destination[prefix + "scale"] = torch.tensor(self.scale)
|
||||
destination[prefix + "zero_point"] = torch.tensor(self.zero_point)
|
||||
|
||||
# ===== Deserialization methods =====
|
||||
# Counterpart to the serialization methods, we must pack the serialized QTensor
|
||||
# weight into its packed format for use by the FBGEMM ops.
|
||||
def _load_from_state_dict(
|
||||
self,
|
||||
state_dict,
|
||||
prefix,
|
||||
local_metadata,
|
||||
strict,
|
||||
missing_keys,
|
||||
unexpected_keys,
|
||||
error_msgs,
|
||||
):
|
||||
self.scale = float(state_dict[prefix + "scale"])
|
||||
state_dict.pop(prefix + "scale")
|
||||
|
||||
self.zero_point = int(state_dict[prefix + "zero_point"])
|
||||
state_dict.pop(prefix + "zero_point")
|
||||
|
||||
version = local_metadata.get("version", None)
|
||||
|
||||
if version is None or version == 1:
|
||||
# We moved the parameters into a LinearPackedParameters submodule
|
||||
weight = state_dict.pop(prefix + "weight")
|
||||
bias = state_dict.pop(prefix + "bias")
|
||||
state_dict.update(
|
||||
{
|
||||
prefix + "_packed_params.weight": weight,
|
||||
prefix + "_packed_params.bias": bias,
|
||||
}
|
||||
)
|
||||
|
||||
super()._load_from_state_dict(
|
||||
state_dict,
|
||||
prefix,
|
||||
local_metadata,
|
||||
False,
|
||||
missing_keys,
|
||||
unexpected_keys,
|
||||
error_msgs,
|
||||
)
|
||||
|
||||
# Function rather than property to make sure that JIT serialization doesn't
|
||||
# register this as an attribute
|
||||
def _weight_bias(self):
|
||||
return self._packed_params._weight_bias()
|
||||
|
||||
def weight(self):
|
||||
return self._weight_bias()[0]
|
||||
|
||||
def bias(self):
|
||||
return self._weight_bias()[1]
|
||||
|
||||
def set_weight_bias(self, w: torch.Tensor, b: torch.Tensor | None) -> None:
|
||||
self._packed_params.set_weight_bias(w, b)
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False):
|
||||
r"""Create a quantized module from an observed float module
|
||||
|
||||
Args:
|
||||
mod (Module): a float module, either produced by torch.ao.quantization
|
||||
utilities or provided by the user
|
||||
use_precomputed_fake_quant (bool): if True, the module will reuse min/max
|
||||
values from the precomputed fake quant module.
|
||||
"""
|
||||
if hasattr(mod, "weight_fake_quant"):
|
||||
if type_before_parametrizations(mod) == nniqat.LinearBn1d:
|
||||
mod.weight, mod.bias = fuse_linear_bn_weights(
|
||||
mod.weight,
|
||||
mod.bias,
|
||||
mod.bn.running_mean,
|
||||
mod.bn.running_var,
|
||||
mod.bn.eps,
|
||||
mod.bn.weight,
|
||||
mod.bn.bias,
|
||||
)
|
||||
weight_post_process = mod.weight_fake_quant
|
||||
activation_post_process = mod.activation_post_process
|
||||
else:
|
||||
# This function does not participate in JIT, so it is OK to ignore
|
||||
# the type mismatch in assignment. Also, mypy has an issue with
|
||||
# iterables not being implemented, so we are ignoring those too.
|
||||
if not isinstance(cls._FLOAT_MODULE, Iterable):
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
cls._FLOAT_MODULE = [cls._FLOAT_MODULE]
|
||||
supported_modules = ", ".join(
|
||||
[float_mod.__name__ for float_mod in cls._FLOAT_MODULE]
|
||||
)
|
||||
error_msg = f"nnq.{cls.__name__}.from_float only works for {supported_modules}, but got: {type(mod)}"
|
||||
if type_before_parametrizations(mod) not in cls._FLOAT_MODULE:
|
||||
raise AssertionError(error_msg)
|
||||
if not hasattr(mod, "qconfig"):
|
||||
raise AssertionError("Input float module must have qconfig defined")
|
||||
activation_post_process = mod.activation_post_process
|
||||
if type_before_parametrizations(mod) == nni.LinearReLU:
|
||||
mod = mod[0]
|
||||
weight_post_process = (
|
||||
mod.qconfig.weight()
|
||||
if not hasattr(mod, "weight_fake_quant")
|
||||
else mod.weight_fake_quant
|
||||
)
|
||||
|
||||
if not use_precomputed_fake_quant:
|
||||
# Observer may not have been called yet
|
||||
# Observer might have been called in the previous stage via PTQ algorithm e.g. AdaRound
|
||||
weight_post_process(mod.weight)
|
||||
dtype = weight_post_process.dtype
|
||||
act_scale, act_zp = activation_post_process.calculate_qparams()
|
||||
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 = cls(mod.in_features, mod.out_features, dtype=dtype)
|
||||
qlinear.set_weight_bias(qweight, mod.bias)
|
||||
qlinear.scale = float(act_scale)
|
||||
qlinear.zero_point = int(act_zp)
|
||||
return qlinear
|
||||
|
||||
@classmethod
|
||||
def from_reference(cls, ref_qlinear, output_scale, output_zero_point):
|
||||
r"""Create a (fbgemm/qnnpack) quantized module from a reference quantized module
|
||||
|
||||
Args:
|
||||
ref_qlinear (Module): a reference quantized linear module, either produced by torch.ao.quantization
|
||||
utilities or provided by the user
|
||||
output_scale (float): scale for output Tensor
|
||||
output_zero_point (int): zero point for output Tensor
|
||||
"""
|
||||
qlinear = cls(ref_qlinear.in_features, ref_qlinear.out_features)
|
||||
qweight = ref_qlinear.get_quantized_weight()
|
||||
qlinear.set_weight_bias(qweight, ref_qlinear.bias)
|
||||
|
||||
qlinear.scale = float(output_scale)
|
||||
qlinear.zero_point = int(output_zero_point)
|
||||
return qlinear
|
||||
+347
@@ -0,0 +1,347 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import torch
|
||||
|
||||
|
||||
__all__ = [
|
||||
"LayerNorm",
|
||||
"GroupNorm",
|
||||
"InstanceNorm1d",
|
||||
"InstanceNorm2d",
|
||||
"InstanceNorm3d",
|
||||
]
|
||||
|
||||
|
||||
class LayerNorm(torch.nn.LayerNorm):
|
||||
r"""This is the quantized version of :class:`~torch.nn.LayerNorm`.
|
||||
|
||||
Additional args:
|
||||
* **scale** - quantization scale of the output, type: double.
|
||||
* **zero_point** - quantization zero point of the output, type: long.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
normalized_shape,
|
||||
weight,
|
||||
bias,
|
||||
scale,
|
||||
zero_point,
|
||||
eps=1e-5,
|
||||
elementwise_affine=True,
|
||||
device=None,
|
||||
dtype=None,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__(
|
||||
normalized_shape,
|
||||
eps=eps,
|
||||
elementwise_affine=elementwise_affine,
|
||||
**factory_kwargs,
|
||||
)
|
||||
self.weight = weight
|
||||
self.bias = bias
|
||||
self.register_buffer("scale", torch.tensor(scale, **factory_kwargs))
|
||||
self.register_buffer("zero_point", torch.tensor(zero_point, **factory_kwargs))
|
||||
|
||||
def forward(self, input):
|
||||
return torch.ops.quantized.layer_norm(
|
||||
input,
|
||||
self.normalized_shape,
|
||||
weight=self.weight,
|
||||
bias=self.bias,
|
||||
eps=self.eps,
|
||||
output_scale=self.scale,
|
||||
output_zero_point=self.zero_point,
|
||||
)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedLayerNorm"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False):
|
||||
scale, zero_point = mod.activation_post_process.calculate_qparams()
|
||||
new_mod = cls(
|
||||
mod.normalized_shape,
|
||||
mod.weight,
|
||||
mod.bias,
|
||||
float(scale),
|
||||
int(zero_point),
|
||||
mod.eps,
|
||||
mod.elementwise_affine,
|
||||
)
|
||||
return new_mod
|
||||
|
||||
@classmethod
|
||||
def from_reference(cls, mod, scale, zero_point):
|
||||
return cls(
|
||||
mod.normalized_shape,
|
||||
mod.weight,
|
||||
mod.bias,
|
||||
float(scale),
|
||||
int(zero_point),
|
||||
mod.eps,
|
||||
mod.elementwise_affine,
|
||||
)
|
||||
|
||||
|
||||
class GroupNorm(torch.nn.GroupNorm):
|
||||
r"""This is the quantized version of :class:`~torch.nn.GroupNorm`.
|
||||
|
||||
Additional args:
|
||||
* **scale** - quantization scale of the output, type: double.
|
||||
* **zero_point** - quantization zero point of the output, type: long.
|
||||
|
||||
"""
|
||||
|
||||
__constants__ = ["num_groups", "num_channels", "eps", "affine"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_groups,
|
||||
num_channels,
|
||||
weight,
|
||||
bias,
|
||||
scale,
|
||||
zero_point,
|
||||
eps=1e-5,
|
||||
affine=True,
|
||||
device=None,
|
||||
dtype=None,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__(num_groups, num_channels, eps, affine, **factory_kwargs)
|
||||
self.weight = weight
|
||||
self.bias = bias
|
||||
self.register_buffer("scale", torch.tensor(scale, **factory_kwargs))
|
||||
self.register_buffer("zero_point", torch.tensor(zero_point, **factory_kwargs))
|
||||
|
||||
def forward(self, input):
|
||||
return torch.ops.quantized.group_norm(
|
||||
input,
|
||||
self.num_groups,
|
||||
self.weight,
|
||||
self.bias,
|
||||
self.eps,
|
||||
self.scale,
|
||||
self.zero_point,
|
||||
)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedGroupNorm"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False):
|
||||
scale, zero_point = mod.activation_post_process.calculate_qparams()
|
||||
new_mod = cls(
|
||||
mod.num_groups,
|
||||
mod.num_channels,
|
||||
mod.weight,
|
||||
mod.bias,
|
||||
float(scale),
|
||||
int(zero_point),
|
||||
mod.eps,
|
||||
mod.affine,
|
||||
)
|
||||
return new_mod
|
||||
|
||||
|
||||
class InstanceNorm1d(torch.nn.InstanceNorm1d):
|
||||
r"""This is the quantized version of :class:`~torch.nn.InstanceNorm1d`.
|
||||
|
||||
Additional args:
|
||||
* **scale** - quantization scale of the output, type: double.
|
||||
* **zero_point** - quantization zero point of the output, type: long.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_features,
|
||||
weight,
|
||||
bias,
|
||||
scale,
|
||||
zero_point,
|
||||
eps=1e-5,
|
||||
momentum=0.1,
|
||||
affine=False,
|
||||
track_running_stats=False,
|
||||
device=None,
|
||||
dtype=None,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__(
|
||||
num_features, eps, momentum, affine, track_running_stats, **factory_kwargs
|
||||
)
|
||||
self.weight = weight
|
||||
self.bias = bias
|
||||
self.register_buffer("scale", torch.tensor(scale, **factory_kwargs))
|
||||
self.register_buffer("zero_point", torch.tensor(zero_point, **factory_kwargs))
|
||||
|
||||
def forward(self, input):
|
||||
return torch.ops.quantized.instance_norm(
|
||||
input, self.weight, self.bias, self.eps, self.scale, self.zero_point
|
||||
)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedInstanceNorm1d"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False):
|
||||
scale, zero_point = mod.activation_post_process.calculate_qparams()
|
||||
new_mod = cls(
|
||||
mod.num_features,
|
||||
mod.weight,
|
||||
mod.bias,
|
||||
float(scale),
|
||||
int(zero_point),
|
||||
mod.eps,
|
||||
mod.affine,
|
||||
)
|
||||
return new_mod
|
||||
|
||||
@classmethod
|
||||
def from_reference(cls, mod, scale, zero_point):
|
||||
return cls(
|
||||
mod.num_features,
|
||||
mod.weight,
|
||||
mod.bias,
|
||||
float(scale),
|
||||
int(zero_point),
|
||||
mod.eps,
|
||||
mod.affine,
|
||||
)
|
||||
|
||||
|
||||
class InstanceNorm2d(torch.nn.InstanceNorm2d):
|
||||
r"""This is the quantized version of :class:`~torch.nn.InstanceNorm2d`.
|
||||
|
||||
Additional args:
|
||||
* **scale** - quantization scale of the output, type: double.
|
||||
* **zero_point** - quantization zero point of the output, type: long.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_features,
|
||||
weight,
|
||||
bias,
|
||||
scale,
|
||||
zero_point,
|
||||
eps=1e-5,
|
||||
momentum=0.1,
|
||||
affine=False,
|
||||
track_running_stats=False,
|
||||
device=None,
|
||||
dtype=None,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__(
|
||||
num_features, eps, momentum, affine, track_running_stats, **factory_kwargs
|
||||
)
|
||||
self.weight = weight
|
||||
self.bias = bias
|
||||
self.register_buffer("scale", torch.tensor(scale, **factory_kwargs))
|
||||
self.register_buffer("zero_point", torch.tensor(zero_point, **factory_kwargs))
|
||||
|
||||
def forward(self, input):
|
||||
return torch.ops.quantized.instance_norm(
|
||||
input, self.weight, self.bias, self.eps, self.scale, self.zero_point
|
||||
)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedInstanceNorm2d"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False):
|
||||
scale, zero_point = mod.activation_post_process.calculate_qparams()
|
||||
new_mod = cls(
|
||||
mod.num_features,
|
||||
mod.weight,
|
||||
mod.bias,
|
||||
float(scale),
|
||||
int(zero_point),
|
||||
mod.eps,
|
||||
mod.affine,
|
||||
)
|
||||
return new_mod
|
||||
|
||||
@classmethod
|
||||
def from_reference(cls, mod, scale, zero_point):
|
||||
return cls(
|
||||
mod.num_features,
|
||||
mod.weight,
|
||||
mod.bias,
|
||||
float(scale),
|
||||
int(zero_point),
|
||||
mod.eps,
|
||||
mod.affine,
|
||||
)
|
||||
|
||||
|
||||
class InstanceNorm3d(torch.nn.InstanceNorm3d):
|
||||
r"""This is the quantized version of :class:`~torch.nn.InstanceNorm3d`.
|
||||
|
||||
Additional args:
|
||||
* **scale** - quantization scale of the output, type: double.
|
||||
* **zero_point** - quantization zero point of the output, type: long.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_features,
|
||||
weight,
|
||||
bias,
|
||||
scale,
|
||||
zero_point,
|
||||
eps=1e-5,
|
||||
momentum=0.1,
|
||||
affine=False,
|
||||
track_running_stats=False,
|
||||
device=None,
|
||||
dtype=None,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__(
|
||||
num_features, eps, momentum, affine, track_running_stats, **factory_kwargs
|
||||
)
|
||||
self.weight = weight
|
||||
self.bias = bias
|
||||
self.register_buffer("scale", torch.tensor(scale, **factory_kwargs))
|
||||
self.register_buffer("zero_point", torch.tensor(zero_point, **factory_kwargs))
|
||||
|
||||
def forward(self, input):
|
||||
return torch.ops.quantized.instance_norm(
|
||||
input, self.weight, self.bias, self.eps, self.scale, self.zero_point
|
||||
)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedInstanceNorm3d"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False):
|
||||
scale, zero_point = mod.activation_post_process.calculate_qparams()
|
||||
new_mod = cls(
|
||||
mod.num_features,
|
||||
mod.weight,
|
||||
mod.bias,
|
||||
float(scale),
|
||||
int(zero_point),
|
||||
mod.eps,
|
||||
mod.affine,
|
||||
)
|
||||
return new_mod
|
||||
|
||||
@classmethod
|
||||
def from_reference(cls, mod, scale, zero_point):
|
||||
return cls(
|
||||
mod.num_features,
|
||||
mod.weight,
|
||||
mod.bias,
|
||||
float(scale),
|
||||
int(zero_point),
|
||||
mod.eps,
|
||||
mod.affine,
|
||||
)
|
||||
@@ -0,0 +1,62 @@
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
__all__ = [
|
||||
"LSTM",
|
||||
]
|
||||
|
||||
|
||||
class LSTM(torch.ao.nn.quantizable.LSTM):
|
||||
r"""A quantized long short-term memory (LSTM).
|
||||
|
||||
For the description and the argument types, please, refer to :class:`~torch.nn.LSTM`
|
||||
|
||||
Attributes:
|
||||
layers : instances of the `_LSTMLayer`
|
||||
|
||||
.. note::
|
||||
To access the weights and biases, you need to access them per layer.
|
||||
See examples in :class:`~torch.ao.nn.quantizable.LSTM`
|
||||
|
||||
Examples::
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> custom_module_config = {
|
||||
... 'float_to_observed_custom_module_class': {
|
||||
... nn.LSTM: nn.quantizable.LSTM,
|
||||
... },
|
||||
... 'observed_to_quantized_custom_module_class': {
|
||||
... nn.quantizable.LSTM: nn.quantized.LSTM,
|
||||
... }
|
||||
... }
|
||||
>>> tq.prepare(model, prepare_custom_module_class=custom_module_config)
|
||||
>>> tq.convert(model, convert_custom_module_class=custom_module_config)
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE = torch.ao.nn.quantizable.LSTM # type: ignore[assignment]
|
||||
|
||||
def _get_name(self) -> str:
|
||||
return "QuantizedLSTM"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, *args: Any, **kwargs: Any) -> None:
|
||||
# The whole flow is float -> observed -> quantized
|
||||
# This class does observed -> quantized only
|
||||
raise NotImplementedError(
|
||||
"It looks like you are trying to convert a "
|
||||
"non-observed LSTM module. Please, see "
|
||||
"the examples on quantizable LSTMs."
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_observed(cls: type["LSTM"], other: torch.ao.nn.quantizable.LSTM) -> "LSTM":
|
||||
if not isinstance(other, cls._FLOAT_MODULE): # type: ignore[has-type]
|
||||
raise AssertionError(
|
||||
f"Expected module type {cls._FLOAT_MODULE}, got {type(other)}"
|
||||
)
|
||||
converted = torch.ao.quantization.convert(
|
||||
other, inplace=False, remove_qconfig=True
|
||||
)
|
||||
converted.__class__ = cls
|
||||
return converted
|
||||
@@ -0,0 +1,144 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import abc
|
||||
import collections
|
||||
import itertools
|
||||
|
||||
import torch
|
||||
from torch.nn.modules.module import _addindent
|
||||
|
||||
|
||||
__all__ = [
|
||||
"WeightedQuantizedModule",
|
||||
]
|
||||
|
||||
|
||||
class WeightedQuantizedModule(torch.nn.Module, metaclass=abc.ABCMeta):
|
||||
"""Wrapper for quantized modules than can be lowered from reference modules."""
|
||||
|
||||
@classmethod
|
||||
@abc.abstractmethod
|
||||
def from_reference(cls, ref_module, output_scale, output_zero_point):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
def _get_weight_observer(observer):
|
||||
# FakeQuantize observer
|
||||
if hasattr(observer, "activation_post_process"):
|
||||
observer = observer.activation_post_process
|
||||
# UniformQuantizationObserverBase observer
|
||||
return observer
|
||||
|
||||
|
||||
def _needs_weight_clamping(observer, dtype):
|
||||
observer = _get_weight_observer(observer)
|
||||
if dtype in [torch.qint8, torch.quint8, torch.qint32]:
|
||||
info = torch.iinfo(dtype)
|
||||
return observer.quant_min > info.min or observer.quant_max < info.max
|
||||
return False
|
||||
|
||||
|
||||
def _clamp_weights(qweight, observer, scale, zp):
|
||||
if not _needs_weight_clamping(observer, qweight.dtype):
|
||||
return qweight
|
||||
|
||||
observer = _get_weight_observer(observer)
|
||||
min_, max_ = observer.quant_min, observer.quant_max
|
||||
|
||||
# Doing this because can't use torch.ops.quantized.clamp() with per_channel qscheme yet.
|
||||
qw_int_max = torch.clone(qweight.int_repr()).fill_(max_)
|
||||
qw_int_min = torch.clone(qweight.int_repr()).fill_(min_)
|
||||
qw_int = torch.minimum(torch.maximum(qweight.int_repr(), qw_int_min), qw_int_max)
|
||||
|
||||
if observer.qscheme in [torch.per_tensor_symmetric, torch.per_tensor_affine]:
|
||||
qweight = torch._make_per_tensor_quantized_tensor(
|
||||
qw_int, scale.item(), zp.item()
|
||||
)
|
||||
elif observer.qscheme in [
|
||||
torch.per_channel_symmetric,
|
||||
torch.per_channel_affine,
|
||||
torch.per_channel_affine_float_qparams,
|
||||
]:
|
||||
qweight = torch._make_per_channel_quantized_tensor(
|
||||
qw_int, scale, zp, axis=observer.ch_axis
|
||||
)
|
||||
else:
|
||||
raise ValueError("Unexpected qscheme " + observer.qscheme)
|
||||
return qweight
|
||||
|
||||
|
||||
def _quantize_weight(float_wt, observer):
|
||||
wt_scale, wt_zp = observer.calculate_qparams()
|
||||
if observer.qscheme in [torch.per_tensor_symmetric, torch.per_tensor_affine]:
|
||||
qweight = torch.quantize_per_tensor(
|
||||
float_wt, float(wt_scale), int(wt_zp), torch.qint8
|
||||
)
|
||||
qweight = _clamp_weights(qweight, observer, wt_scale, wt_zp)
|
||||
elif observer.qscheme in [torch.per_channel_symmetric, torch.per_channel_affine]:
|
||||
wt_axis = observer.ch_axis
|
||||
qweight = torch.quantize_per_channel(
|
||||
float_wt,
|
||||
wt_scale.to(torch.double),
|
||||
wt_zp.to(torch.int64),
|
||||
wt_axis,
|
||||
torch.qint8,
|
||||
)
|
||||
qweight = _clamp_weights(qweight, observer, wt_scale, wt_zp)
|
||||
elif observer.qscheme == torch.per_channel_affine_float_qparams:
|
||||
qweight = torch.quantize_per_channel(
|
||||
float_wt,
|
||||
wt_scale.to(torch.float),
|
||||
wt_zp.to(torch.float),
|
||||
observer.ch_axis,
|
||||
observer.dtype,
|
||||
)
|
||||
qweight = _clamp_weights(qweight, observer, wt_scale, wt_zp)
|
||||
else:
|
||||
raise ValueError("Unexpected qscheme " + observer.qscheme)
|
||||
return qweight
|
||||
|
||||
|
||||
def _ntuple_from_first(n):
|
||||
"""Converts the argument to a tuple of size n
|
||||
with the first element repeated."""
|
||||
|
||||
def parse(x):
|
||||
while isinstance(x, collections.abc.Sequence):
|
||||
if len(x) == n:
|
||||
break
|
||||
x = x[0]
|
||||
return tuple(itertools.repeat(x, n))
|
||||
|
||||
return parse
|
||||
|
||||
|
||||
def _hide_packed_params_repr(self, params):
|
||||
# We don't want to show `PackedParams` children, hence custom
|
||||
# `__repr__`. This is the same as nn.Module.__repr__, except the check
|
||||
# for the `params module`.
|
||||
extra_lines = []
|
||||
extra_repr = self.extra_repr()
|
||||
# empty string will be split into list ['']
|
||||
if extra_repr:
|
||||
extra_lines = extra_repr.split("\n")
|
||||
child_lines = []
|
||||
for key, module in self._modules.items():
|
||||
if isinstance(module, params):
|
||||
continue
|
||||
mod_str = repr(module)
|
||||
mod_str = _addindent(mod_str, 2)
|
||||
child_lines.append("(" + key + "): " + mod_str)
|
||||
lines = extra_lines + child_lines
|
||||
|
||||
main_str = self._get_name() + "("
|
||||
if lines:
|
||||
# simple one-liner info, which most builtin Modules will use
|
||||
if len(extra_lines) == 1 and not child_lines:
|
||||
main_str += extra_lines[0]
|
||||
else:
|
||||
main_str += "\n " + "\n ".join(lines) + "\n"
|
||||
|
||||
main_str += ")"
|
||||
return main_str
|
||||
|
||||
|
||||
_pair_from_first = _ntuple_from_first(2)
|
||||
@@ -0,0 +1,19 @@
|
||||
from .modules import * # noqa: F403
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Linear",
|
||||
"Conv1d",
|
||||
"Conv2d",
|
||||
"Conv3d",
|
||||
"ConvTranspose1d",
|
||||
"ConvTranspose2d",
|
||||
"ConvTranspose3d",
|
||||
"RNNCell",
|
||||
"LSTMCell",
|
||||
"GRUCell",
|
||||
"LSTM",
|
||||
"GRU",
|
||||
"Embedding",
|
||||
"EmbeddingBag",
|
||||
]
|
||||
+29
@@ -0,0 +1,29 @@
|
||||
from .conv import (
|
||||
Conv1d,
|
||||
Conv2d,
|
||||
Conv3d,
|
||||
ConvTranspose1d,
|
||||
ConvTranspose2d,
|
||||
ConvTranspose3d,
|
||||
)
|
||||
from .linear import Linear
|
||||
from .rnn import GRU, GRUCell, LSTM, LSTMCell, RNNCell
|
||||
from .sparse import Embedding, EmbeddingBag
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Linear",
|
||||
"Conv1d",
|
||||
"Conv2d",
|
||||
"Conv3d",
|
||||
"ConvTranspose1d",
|
||||
"ConvTranspose2d",
|
||||
"ConvTranspose3d",
|
||||
"RNNCell",
|
||||
"LSTMCell",
|
||||
"GRUCell",
|
||||
"LSTM",
|
||||
"GRU",
|
||||
"Embedding",
|
||||
"EmbeddingBag",
|
||||
]
|
||||
+527
@@ -0,0 +1,527 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from typing import Any, Literal
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.nn.common_types import _size_1_t
|
||||
|
||||
from .utils import ReferenceQuantizedModule
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Conv1d",
|
||||
"Conv2d",
|
||||
"Conv3d",
|
||||
"ConvTranspose1d",
|
||||
"ConvTranspose2d",
|
||||
"ConvTranspose3d",
|
||||
]
|
||||
|
||||
|
||||
class _ConvNd(torch.nn.modules.conv._ConvNd, ReferenceQuantizedModule):
|
||||
"""A reference version of nn.quantized.Conv2d
|
||||
we will not pack the parameters in this module, since weight packing is an
|
||||
optimization for quantized backends supported in PyTorch (fbgemm/qnnpack),
|
||||
this is useful when user want to use this module in other backends like Glow.
|
||||
"""
|
||||
|
||||
__annotations__ = {"bias": torch.Tensor | None}
|
||||
_IS_REFERENCE = True
|
||||
|
||||
@staticmethod
|
||||
def from_float(cls, float_conv, weight_qparams):
|
||||
qref_conv = cls(
|
||||
float_conv.in_channels,
|
||||
float_conv.out_channels,
|
||||
float_conv.kernel_size, # type: ignore[arg-type]
|
||||
float_conv.stride, # type: ignore[arg-type]
|
||||
float_conv.padding, # type: ignore[arg-type]
|
||||
float_conv.dilation, # type: ignore[arg-type]
|
||||
float_conv.groups,
|
||||
float_conv.bias is not None, # type: ignore[arg-type]
|
||||
float_conv.padding_mode,
|
||||
device=float_conv.weight.device,
|
||||
dtype=float_conv.weight.dtype,
|
||||
weight_qparams=weight_qparams,
|
||||
)
|
||||
qref_conv.weight = torch.nn.Parameter(float_conv.weight.detach())
|
||||
if float_conv.bias is not None:
|
||||
qref_conv.bias = torch.nn.Parameter(float_conv.bias.detach())
|
||||
return qref_conv
|
||||
|
||||
|
||||
class Conv1d(_ConvNd, nn.Conv1d):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size: _size_1_t,
|
||||
stride: _size_1_t = 1,
|
||||
padding: _size_1_t = 0,
|
||||
dilation: _size_1_t = 1,
|
||||
groups: int = 1,
|
||||
bias: bool = True,
|
||||
padding_mode: Literal["zeros", "reflect", "replicate", "circular"] = "zeros",
|
||||
device=None,
|
||||
dtype=None,
|
||||
weight_qparams: dict[str, Any] | None = None,
|
||||
):
|
||||
nn.Conv1d.__init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
dilation,
|
||||
groups,
|
||||
bias,
|
||||
padding_mode,
|
||||
device,
|
||||
dtype,
|
||||
)
|
||||
self._init_weight_qparams(weight_qparams, device)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
we have:
|
||||
w(float) -- quant - dequant \
|
||||
x(float) ------------- F.conv1d ---
|
||||
|
||||
In the full model, we will see
|
||||
w(float) -- quant - *dequant \
|
||||
x -- quant --- *dequant -- *F.conv1d --- *quant - dequant
|
||||
and the backend should be able to fuse the ops with `*` into a quantized conv1d
|
||||
"""
|
||||
weight_quant_dequant = self.get_weight()
|
||||
|
||||
result = F.conv1d(
|
||||
x,
|
||||
weight_quant_dequant,
|
||||
self.bias,
|
||||
self.stride,
|
||||
self.padding,
|
||||
self.dilation,
|
||||
self.groups,
|
||||
)
|
||||
return result
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedConv1d(Reference)"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, float_conv, weight_qparams): # type: ignore[override]
|
||||
return _ConvNd.from_float(cls, float_conv, weight_qparams)
|
||||
|
||||
|
||||
class Conv2d(_ConvNd, nn.Conv2d):
|
||||
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,
|
||||
weight_qparams: dict[str, Any] | None = None,
|
||||
):
|
||||
nn.Conv2d.__init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
dilation,
|
||||
groups,
|
||||
bias,
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
padding_mode,
|
||||
device,
|
||||
dtype,
|
||||
)
|
||||
self._init_weight_qparams(weight_qparams, device)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
we have:
|
||||
w(float) -- quant - dequant \
|
||||
x(float) ------------- F.conv2d ---
|
||||
|
||||
In the full model, we will see
|
||||
w(float) -- quant - *dequant \
|
||||
x -- quant --- *dequant -- *F.conv2d --- *quant - dequant
|
||||
and the backend should be able to fuse the ops with `*` into a quantized conv2d
|
||||
"""
|
||||
weight_quant_dequant = self.get_weight()
|
||||
|
||||
result = F.conv2d(
|
||||
x,
|
||||
weight_quant_dequant,
|
||||
self.bias,
|
||||
self.stride,
|
||||
self.padding,
|
||||
self.dilation,
|
||||
self.groups,
|
||||
)
|
||||
return result
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedConv2d(Reference)"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, float_conv, weight_qparams): # type: ignore[override]
|
||||
return _ConvNd.from_float(cls, float_conv, weight_qparams)
|
||||
|
||||
|
||||
class Conv3d(_ConvNd, nn.Conv3d):
|
||||
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,
|
||||
weight_qparams: dict[str, Any] | None = None,
|
||||
):
|
||||
nn.Conv3d.__init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
dilation,
|
||||
groups,
|
||||
bias,
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
padding_mode,
|
||||
device,
|
||||
dtype,
|
||||
)
|
||||
self._init_weight_qparams(weight_qparams, device)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
we have:
|
||||
w(float) -- quant - dequant \
|
||||
x(float) ------------- F.conv3d ---
|
||||
|
||||
In the full model, we will see
|
||||
w(float) -- quant - *dequant \
|
||||
x -- quant --- *dequant -- *F.conv3d --- *quant - dequant
|
||||
and the backend should be able to fuse the ops with `*` into a quantized conv3d
|
||||
"""
|
||||
weight_quant_dequant = self.get_weight()
|
||||
|
||||
result = F.conv3d(
|
||||
x,
|
||||
weight_quant_dequant,
|
||||
self.bias,
|
||||
self.stride,
|
||||
self.padding,
|
||||
self.dilation,
|
||||
self.groups,
|
||||
)
|
||||
return result
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedConv3d(Reference)"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, float_conv, weight_qparams): # type: ignore[override]
|
||||
return _ConvNd.from_float(cls, float_conv, weight_qparams)
|
||||
|
||||
|
||||
class _ConvTransposeNd(_ConvNd, torch.nn.modules.conv._ConvTransposeNd):
|
||||
"""A reference version of nn.quantized.ConvTranspose2d
|
||||
we will not pack the parameters in this module, since weight packing is an
|
||||
optimization for quantized backends supported in PyTorch (fbgemm/qnnpack),
|
||||
this is useful when user want to use this module in other backends like Glow.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def from_float(cls, float_conv, weight_qparams):
|
||||
qref_conv = cls(
|
||||
float_conv.in_channels,
|
||||
float_conv.out_channels,
|
||||
float_conv.kernel_size, # type: ignore[arg-type]
|
||||
float_conv.stride, # type: ignore[arg-type]
|
||||
float_conv.padding, # type: ignore[arg-type]
|
||||
float_conv.output_padding, # type: ignore[arg-type]
|
||||
float_conv.groups,
|
||||
float_conv.bias is not None, # type: ignore[arg-type]
|
||||
float_conv.dilation, # type: ignore[arg-type]
|
||||
float_conv.padding_mode,
|
||||
device=float_conv.weight.device,
|
||||
dtype=float_conv.weight.dtype,
|
||||
weight_qparams=weight_qparams,
|
||||
)
|
||||
qref_conv.weight = torch.nn.Parameter(float_conv.weight.detach())
|
||||
if float_conv.bias is not None:
|
||||
qref_conv.bias = torch.nn.Parameter(float_conv.bias.detach())
|
||||
return qref_conv
|
||||
|
||||
|
||||
class ConvTranspose1d(_ConvTransposeNd, nn.ConvTranspose1d):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size: _size_1_t,
|
||||
stride: _size_1_t = 1,
|
||||
padding: _size_1_t = 0,
|
||||
output_padding: _size_1_t = 0,
|
||||
groups: int = 1,
|
||||
bias: bool = True,
|
||||
dilation: _size_1_t = 1,
|
||||
padding_mode: Literal["zeros", "reflect", "replicate", "circular"] = "zeros",
|
||||
device=None,
|
||||
dtype=None,
|
||||
weight_qparams: dict[str, Any] | None = None,
|
||||
):
|
||||
nn.ConvTranspose1d.__init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
output_padding,
|
||||
groups,
|
||||
bias,
|
||||
dilation,
|
||||
padding_mode,
|
||||
device,
|
||||
dtype,
|
||||
)
|
||||
self._init_weight_qparams(weight_qparams, device)
|
||||
|
||||
def forward(
|
||||
self, x: torch.Tensor, output_size: list[int] | None = None
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
we have:
|
||||
w(float) -- quant - dequant \
|
||||
x(float) ------------- F.convTranspose1d ---
|
||||
In the full model, we will see
|
||||
w(float) -- quant - *dequant \
|
||||
x -- quant --- *dequant -- *F.convTranspose1d --- *quant - dequant
|
||||
and the backend should be able to fuse the ops with `*` into a quantized conv1d
|
||||
"""
|
||||
|
||||
if not isinstance(self.padding, tuple):
|
||||
raise AssertionError(
|
||||
f"Expected self.padding to be a tuple, got {type(self.padding)}"
|
||||
)
|
||||
# One cannot replace List by Tuple or Sequence in "_output_padding" because
|
||||
# TorchScript does not support `Sequence[T]` or `Tuple[T, ...]`.
|
||||
output_padding = self._output_padding(
|
||||
input, # type: ignore[arg-type]
|
||||
output_size,
|
||||
self.stride, # type: ignore[arg-type]
|
||||
self.padding, # type: ignore[arg-type]
|
||||
self.kernel_size, # type: ignore[arg-type]
|
||||
self.dilation, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
weight_quant_dequant = self.get_weight()
|
||||
result = F.conv_transpose1d(
|
||||
x,
|
||||
weight_quant_dequant,
|
||||
self.bias,
|
||||
self.stride,
|
||||
self.padding,
|
||||
output_padding,
|
||||
self.groups,
|
||||
self.dilation,
|
||||
)
|
||||
return result
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedConvTranspose1d(Reference)"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, float_conv, weight_qparams): # type: ignore[override]
|
||||
return _ConvTransposeNd.from_float(cls, float_conv, weight_qparams)
|
||||
|
||||
|
||||
class ConvTranspose2d(_ConvTransposeNd, nn.ConvTranspose2d):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
output_padding=0,
|
||||
groups=1,
|
||||
bias=True,
|
||||
dilation=1,
|
||||
padding_mode="zeros",
|
||||
device=None,
|
||||
dtype=None,
|
||||
weight_qparams: dict[str, Any] | None = None,
|
||||
):
|
||||
nn.ConvTranspose2d.__init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
output_padding,
|
||||
groups,
|
||||
bias,
|
||||
dilation,
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
padding_mode,
|
||||
device,
|
||||
dtype,
|
||||
)
|
||||
self._init_weight_qparams(weight_qparams, device)
|
||||
|
||||
def forward(
|
||||
self, x: torch.Tensor, output_size: list[int] | None = None
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
we have:
|
||||
w(float) -- quant - dequant \
|
||||
x(float) ------------- F.convTranspose2d ---
|
||||
In the full model, we will see
|
||||
w(float) -- quant - *dequant \
|
||||
x -- quant --- *dequant -- *F.convTranspose2d --- *quant - dequant
|
||||
and the backend should be able to fuse the ops with `*` into a quantized conv2d
|
||||
"""
|
||||
if not isinstance(self.padding, tuple):
|
||||
raise AssertionError(
|
||||
f"Expected self.padding to be a tuple, got {type(self.padding)}"
|
||||
)
|
||||
# One cannot replace List by Tuple or Sequence in "_output_padding" because
|
||||
# TorchScript does not support `Sequence[T]` or `Tuple[T, ...]`.
|
||||
|
||||
output_padding = self._output_padding(
|
||||
input, # type: ignore[arg-type]
|
||||
output_size,
|
||||
self.stride, # type: ignore[arg-type]
|
||||
self.padding, # type: ignore[arg-type]
|
||||
self.kernel_size, # type: ignore[arg-type]
|
||||
self.dilation, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
weight_quant_dequant = self.get_weight()
|
||||
result = F.conv_transpose2d(
|
||||
x,
|
||||
weight_quant_dequant,
|
||||
self.bias,
|
||||
self.stride,
|
||||
self.padding,
|
||||
output_padding,
|
||||
self.groups,
|
||||
self.dilation,
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedConvTranspose2d(Reference)"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, float_conv, weight_qparams): # type: ignore[override]
|
||||
return _ConvTransposeNd.from_float(cls, float_conv, weight_qparams)
|
||||
|
||||
|
||||
class ConvTranspose3d(_ConvTransposeNd, nn.ConvTranspose3d):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
output_padding=0,
|
||||
groups=1,
|
||||
bias=True,
|
||||
dilation=1,
|
||||
padding_mode="zeros",
|
||||
device=None,
|
||||
dtype=None,
|
||||
weight_qparams: dict[str, Any] | None = None,
|
||||
):
|
||||
nn.ConvTranspose3d.__init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
output_padding,
|
||||
groups,
|
||||
bias,
|
||||
dilation,
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
padding_mode,
|
||||
device,
|
||||
dtype,
|
||||
)
|
||||
self._init_weight_qparams(weight_qparams, device)
|
||||
|
||||
def forward(
|
||||
self, x: torch.Tensor, output_size: list[int] | None = None
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
we have:
|
||||
w(float) -- quant - dequant \
|
||||
x(float) ------------- F.convTranspose3d ---
|
||||
In the full model, we will see
|
||||
w(float) -- quant - *dequant \
|
||||
x -- quant --- *dequant -- *F.convTranspose3d --- *quant - dequant
|
||||
and the backend should be able to fuse the ops with `*` into a quantized conv3d
|
||||
"""
|
||||
|
||||
if not isinstance(self.padding, tuple):
|
||||
raise AssertionError(
|
||||
f"Expected self.padding to be a tuple, got {type(self.padding)}"
|
||||
)
|
||||
# One cannot replace List by Tuple or Sequence in "_output_padding" because
|
||||
# TorchScript does not support `Sequence[T]` or `Tuple[T, ...]`.
|
||||
output_padding = self._output_padding(
|
||||
input, # type: ignore[arg-type]
|
||||
output_size,
|
||||
self.stride, # type: ignore[arg-type]
|
||||
self.padding, # type: ignore[arg-type]
|
||||
self.kernel_size, # type: ignore[arg-type]
|
||||
self.dilation, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
weight_quant_dequant = self.get_weight()
|
||||
result = F.conv_transpose3d(
|
||||
x,
|
||||
weight_quant_dequant,
|
||||
self.bias,
|
||||
self.stride,
|
||||
self.padding,
|
||||
output_padding,
|
||||
self.groups,
|
||||
self.dilation,
|
||||
)
|
||||
return result
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedConvTranspose3d(Reference)"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, float_conv, weight_qparams): # type: ignore[override]
|
||||
return _ConvTransposeNd.from_float(cls, float_conv, weight_qparams)
|
||||
+69
@@ -0,0 +1,69 @@
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .utils import ReferenceQuantizedModule
|
||||
|
||||
|
||||
__all__ = ["Linear"]
|
||||
|
||||
|
||||
class Linear(nn.Linear, ReferenceQuantizedModule):
|
||||
"""A reference quantized linear module that fits into the FX
|
||||
Graph Mode Quantization workflow
|
||||
activation will be floating point Tensor, we will store floating
|
||||
point weight as well in the module, but in forward we'll quantize
|
||||
and dequantize the weight before running the floating point functional
|
||||
linear operator.
|
||||
"""
|
||||
|
||||
_IS_REFERENCE = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_features: int,
|
||||
out_features: int,
|
||||
bias_: bool = True,
|
||||
device: torch.device | None = None,
|
||||
dtype: torch.dtype | None = None,
|
||||
weight_qparams: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
super().__init__(in_features, out_features, bias_, device, dtype)
|
||||
self._init_weight_qparams(weight_qparams, device)
|
||||
|
||||
def _get_name(self) -> str:
|
||||
return "QuantizedLinear(Reference)"
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
we have:
|
||||
w(float) -- quant - dequant \
|
||||
x(float) ------------- F.linear ---
|
||||
|
||||
In the full model, we will see
|
||||
w(float) -- quant - *dequant \
|
||||
x -- quant --- *dequant -- *F.linear --- *quant - dequant
|
||||
and the backend should be able to fuse the ops with `*` into a quantized linear
|
||||
"""
|
||||
weight_quant_dequant = self.get_weight()
|
||||
result = F.linear(x, weight_quant_dequant, self.bias)
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
def from_float(
|
||||
cls, float_linear: nn.Linear, weight_qparams: dict[str, Any]
|
||||
) -> "Linear":
|
||||
qref_linear = Linear(
|
||||
float_linear.in_features,
|
||||
float_linear.out_features,
|
||||
float_linear.bias is not None,
|
||||
device=float_linear.weight.device,
|
||||
dtype=float_linear.weight.dtype,
|
||||
weight_qparams=weight_qparams,
|
||||
)
|
||||
qref_linear.weight = torch.nn.Parameter(float_linear.weight.detach())
|
||||
if float_linear.bias is not None:
|
||||
qref_linear.bias = torch.nn.Parameter(float_linear.bias.detach())
|
||||
return qref_linear
|
||||
+856
@@ -0,0 +1,856 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch import _VF, Tensor
|
||||
from torch.nn.utils.rnn import PackedSequence
|
||||
|
||||
from .utils import _quantize_and_dequantize_weight, _quantize_weight
|
||||
|
||||
|
||||
__all__ = [
|
||||
"RNNCellBase",
|
||||
"RNNCell",
|
||||
"LSTMCell",
|
||||
"GRUCell",
|
||||
"RNNBase",
|
||||
"LSTM",
|
||||
"GRU",
|
||||
"get_quantized_weight",
|
||||
]
|
||||
|
||||
|
||||
def _apply_permutation(tensor: Tensor, permutation: Tensor, dim: int = 1) -> Tensor:
|
||||
return tensor.index_select(dim, permutation)
|
||||
|
||||
|
||||
def _get_weight_and_quantization_params(module, wn):
|
||||
weight = getattr(module, wn)
|
||||
params = [weight]
|
||||
for param_name in [
|
||||
wn + n for n in ["_qscheme", "_dtype", "_scale", "_zero_point", "_axis_int"]
|
||||
]:
|
||||
if hasattr(module, param_name):
|
||||
param = getattr(module, param_name)
|
||||
else:
|
||||
param = None
|
||||
params.append(param)
|
||||
return params
|
||||
|
||||
|
||||
def get_quantized_weight(module, wn):
|
||||
if not hasattr(module, wn):
|
||||
return None
|
||||
params = _get_weight_and_quantization_params(module, wn)
|
||||
weight = _quantize_weight(*params)
|
||||
return weight
|
||||
|
||||
|
||||
def _get_quantize_and_dequantized_weight(module, wn):
|
||||
if not hasattr(module, wn):
|
||||
return None
|
||||
params = _get_weight_and_quantization_params(module, wn)
|
||||
weight = _quantize_and_dequantize_weight(*params)
|
||||
return weight
|
||||
|
||||
|
||||
class RNNCellBase(nn.RNNCellBase):
|
||||
def __init__(
|
||||
self,
|
||||
input_size: int,
|
||||
hidden_size: int,
|
||||
bias: bool,
|
||||
num_chunks: int,
|
||||
device=None,
|
||||
dtype=None,
|
||||
weight_qparams_dict=None,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
input_size, hidden_size, bias, num_chunks, device=device, dtype=dtype
|
||||
)
|
||||
# TODO(jerryzh168): maybe make this arg a required arg
|
||||
if weight_qparams_dict is None:
|
||||
weight_qparams = {
|
||||
"qscheme": torch.per_tensor_affine,
|
||||
"dtype": torch.quint8,
|
||||
"scale": 1.0,
|
||||
"zero_point": 0,
|
||||
}
|
||||
weight_qparams_dict = {
|
||||
"weight_ih": weight_qparams,
|
||||
"weight_hh": weight_qparams,
|
||||
"is_decomposed": False,
|
||||
}
|
||||
if len(weight_qparams_dict) != 3:
|
||||
raise AssertionError(
|
||||
f"Expected length for weight_qparams_dict to be 3 for QuantizedRNNCellBase(Reference), "
|
||||
f"got {len(weight_qparams_dict)}"
|
||||
)
|
||||
self._init_weight_qparams_dict(weight_qparams_dict, device)
|
||||
|
||||
def _init_weight_qparams_dict(self, weight_qparams_dict, device):
|
||||
if weight_qparams_dict is None:
|
||||
raise AssertionError("weight_qparams_dict must not be None")
|
||||
self.is_decomposed = weight_qparams_dict["is_decomposed"]
|
||||
for key, weight_qparams in weight_qparams_dict.items():
|
||||
if key == "is_decomposed":
|
||||
continue
|
||||
# TODO: refactor the duplicated code to utils.py
|
||||
weight_qscheme = weight_qparams["qscheme"]
|
||||
weight_dtype = weight_qparams["dtype"]
|
||||
setattr(self, key + "_qscheme", weight_qscheme)
|
||||
setattr(self, key + "_dtype", weight_dtype)
|
||||
if weight_qscheme not in [
|
||||
None,
|
||||
torch.per_tensor_affine,
|
||||
torch.per_channel_affine,
|
||||
]:
|
||||
raise AssertionError(
|
||||
f"qscheme: {weight_qscheme} is not supported in {self._get_name()}"
|
||||
)
|
||||
if weight_qscheme is not None:
|
||||
scale = weight_qparams["scale"]
|
||||
scale_tensor = (
|
||||
scale.detach().clone()
|
||||
if isinstance(scale, torch.Tensor)
|
||||
else torch.tensor(scale, dtype=torch.float, device=device)
|
||||
)
|
||||
self.register_buffer(key + "_scale", scale_tensor)
|
||||
zp = weight_qparams["zero_point"]
|
||||
zp_tensor = (
|
||||
zp.detach().clone()
|
||||
if isinstance(zp, torch.Tensor)
|
||||
else torch.tensor(zp, dtype=torch.int, device=device)
|
||||
)
|
||||
self.register_buffer(key + "_zero_point", zp_tensor)
|
||||
if weight_qscheme == torch.per_channel_affine:
|
||||
axis = weight_qparams["axis"]
|
||||
axis_tensor = (
|
||||
axis.detach().clone()
|
||||
if isinstance(axis, torch.Tensor)
|
||||
else torch.tensor(axis, dtype=torch.int, device=device)
|
||||
)
|
||||
self.register_buffer(key + "_axis", axis_tensor)
|
||||
else:
|
||||
# added for TorchScriptability, not used
|
||||
self.register_buffer(
|
||||
key + "_axis", torch.tensor(0, dtype=torch.int, device=device)
|
||||
)
|
||||
setattr(self, key + "_axis_int", getattr(self, key + "_axis").item())
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedRNNCellBase(Reference)"
|
||||
|
||||
def get_quantized_weight_ih(self):
|
||||
return get_quantized_weight(self, "weight_ih")
|
||||
|
||||
def get_quantized_weight_hh(self):
|
||||
return get_quantized_weight(self, "weight_hh")
|
||||
|
||||
def get_weight_ih(self):
|
||||
return _get_quantize_and_dequantized_weight(self, "weight_ih")
|
||||
|
||||
def get_weight_hh(self):
|
||||
return _get_quantize_and_dequantized_weight(self, "weight_hh")
|
||||
|
||||
|
||||
class RNNCell(RNNCellBase):
|
||||
"""
|
||||
We'll store weight_qparams for all the weights (weight_ih and weight_hh),
|
||||
we need to pass in a `weight_qparams_dict` that maps from weight name,
|
||||
e.g. weight_ih, to the weight_qparams for that weight
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_size: int,
|
||||
hidden_size: int,
|
||||
bias: bool = True,
|
||||
nonlinearity: str = "tanh",
|
||||
device=None,
|
||||
dtype=None,
|
||||
weight_qparams_dict: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
factory_kwargs = {
|
||||
"device": device,
|
||||
"dtype": dtype,
|
||||
"weight_qparams_dict": weight_qparams_dict,
|
||||
}
|
||||
super().__init__(input_size, hidden_size, bias, num_chunks=1, **factory_kwargs)
|
||||
self.nonlinearity = nonlinearity
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedRNNCell(Reference)"
|
||||
|
||||
# TODO: refactor nn.RNNCell to have a _forward that takes weight_ih and weight_hh as input
|
||||
# and remove duplicated code, same for the other two Cell modules
|
||||
def forward(self, input: Tensor, hx: Tensor | None = None) -> Tensor:
|
||||
if input.dim() not in (1, 2):
|
||||
raise AssertionError(
|
||||
f"RNNCell: Expected input to be 1-D or 2-D but received {input.dim()}-D tensor"
|
||||
)
|
||||
is_batched = input.dim() == 2
|
||||
if not is_batched:
|
||||
input = input.unsqueeze(0)
|
||||
|
||||
if hx is None:
|
||||
hx = torch.zeros(
|
||||
input.size(0), self.hidden_size, dtype=input.dtype, device=input.device
|
||||
)
|
||||
else:
|
||||
hx = hx.unsqueeze(0) if not is_batched else hx
|
||||
|
||||
if self.nonlinearity == "tanh":
|
||||
ret = _VF.rnn_tanh_cell(
|
||||
input,
|
||||
hx,
|
||||
self.get_weight_ih(),
|
||||
self.get_weight_hh(),
|
||||
self.bias_ih,
|
||||
self.bias_hh,
|
||||
)
|
||||
elif self.nonlinearity == "relu":
|
||||
ret = _VF.rnn_relu_cell(
|
||||
input,
|
||||
hx,
|
||||
self.get_weight_ih(),
|
||||
self.get_weight_hh(),
|
||||
self.bias_ih,
|
||||
self.bias_hh,
|
||||
)
|
||||
else:
|
||||
ret = input # TODO: remove when jit supports exception flow
|
||||
raise RuntimeError(f"Unknown nonlinearity: {self.nonlinearity}")
|
||||
|
||||
if not is_batched:
|
||||
ret = ret.squeeze(0)
|
||||
|
||||
return ret
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, weight_qparams_dict):
|
||||
ref_mod = cls(
|
||||
mod.input_size,
|
||||
mod.hidden_size,
|
||||
mod.bias,
|
||||
mod.nonlinearity,
|
||||
mod.weight_ih.device,
|
||||
mod.weight_ih.dtype,
|
||||
weight_qparams_dict,
|
||||
)
|
||||
ref_mod.weight_ih = mod.weight_ih
|
||||
ref_mod.weight_hh = mod.weight_hh
|
||||
ref_mod.bias_ih = mod.bias_ih
|
||||
ref_mod.bias_hh = mod.bias_hh
|
||||
return ref_mod
|
||||
|
||||
|
||||
class LSTMCell(RNNCellBase):
|
||||
"""
|
||||
We'll store weight_qparams for all the weights (weight_ih and weight_hh),
|
||||
we need to pass in a `weight_qparams_dict` that maps from weight name,
|
||||
e.g. weight_ih, to the weight_qparams for that weight
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_size: int,
|
||||
hidden_size: int,
|
||||
bias: bool = True,
|
||||
device=None,
|
||||
dtype=None,
|
||||
weight_qparams_dict: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
factory_kwargs = {
|
||||
"device": device,
|
||||
"dtype": dtype,
|
||||
"weight_qparams_dict": weight_qparams_dict,
|
||||
}
|
||||
super().__init__(input_size, hidden_size, bias, num_chunks=4, **factory_kwargs)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedLSTMCell(Reference)"
|
||||
|
||||
def forward(
|
||||
self, input: Tensor, hx: tuple[Tensor, Tensor] | None = None
|
||||
) -> tuple[Tensor, Tensor]:
|
||||
if input.dim() not in (1, 2):
|
||||
raise AssertionError(
|
||||
f"LSTMCell: Expected input to be 1-D or 2-D but received {input.dim()}-D tensor"
|
||||
)
|
||||
is_batched = input.dim() == 2
|
||||
if not is_batched:
|
||||
input = input.unsqueeze(0)
|
||||
|
||||
if hx is None:
|
||||
zeros = torch.zeros(
|
||||
input.size(0), self.hidden_size, dtype=input.dtype, device=input.device
|
||||
)
|
||||
hx = (zeros, zeros)
|
||||
else:
|
||||
hx = (hx[0].unsqueeze(0), hx[1].unsqueeze(0)) if not is_batched else hx
|
||||
|
||||
ret = _VF.lstm_cell(
|
||||
input,
|
||||
hx,
|
||||
self.get_weight_ih(),
|
||||
self.get_weight_hh(),
|
||||
self.bias_ih,
|
||||
self.bias_hh,
|
||||
)
|
||||
|
||||
if not is_batched:
|
||||
ret = (ret[0].squeeze(0), ret[1].squeeze(0))
|
||||
return ret
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, weight_qparams_dict, use_precomputed_fake_quant=False):
|
||||
ref_mod = cls(
|
||||
mod.input_size,
|
||||
mod.hidden_size,
|
||||
mod.bias,
|
||||
mod.weight_ih.device,
|
||||
mod.weight_ih.dtype,
|
||||
weight_qparams_dict,
|
||||
)
|
||||
ref_mod.weight_ih = mod.weight_ih
|
||||
ref_mod.weight_hh = mod.weight_hh
|
||||
ref_mod.bias_ih = mod.bias_ih
|
||||
ref_mod.bias_hh = mod.bias_hh
|
||||
return ref_mod
|
||||
|
||||
|
||||
class GRUCell(RNNCellBase):
|
||||
"""
|
||||
We'll store weight_qparams for all the weights (weight_ih and weight_hh),
|
||||
we need to pass in a `weight_qparams_dict` that maps from weight name,
|
||||
e.g. weight_ih, to the weight_qparams for that weight
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_size: int,
|
||||
hidden_size: int,
|
||||
bias: bool = True,
|
||||
device=None,
|
||||
dtype=None,
|
||||
weight_qparams_dict: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
factory_kwargs = {
|
||||
"device": device,
|
||||
"dtype": dtype,
|
||||
"weight_qparams_dict": weight_qparams_dict,
|
||||
}
|
||||
super().__init__(input_size, hidden_size, bias, num_chunks=3, **factory_kwargs)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedGRUCell(Reference)"
|
||||
|
||||
def forward(self, input: Tensor, hx: Tensor | None = None) -> Tensor:
|
||||
if input.dim() not in (1, 2):
|
||||
raise AssertionError(
|
||||
f"GRUCell: Expected input to be 1-D or 2-D but received {input.dim()}-D tensor"
|
||||
)
|
||||
is_batched = input.dim() == 2
|
||||
if not is_batched:
|
||||
input = input.unsqueeze(0)
|
||||
|
||||
if hx is None:
|
||||
hx = torch.zeros(
|
||||
input.size(0), self.hidden_size, dtype=input.dtype, device=input.device
|
||||
)
|
||||
else:
|
||||
hx = hx.unsqueeze(0) if not is_batched else hx
|
||||
|
||||
ret = _VF.gru_cell(
|
||||
input,
|
||||
hx,
|
||||
self.get_weight_ih(),
|
||||
self.get_weight_hh(),
|
||||
self.bias_ih,
|
||||
self.bias_hh,
|
||||
)
|
||||
|
||||
if not is_batched:
|
||||
ret = ret.squeeze(0)
|
||||
|
||||
return ret
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, weight_qparams_dict):
|
||||
ref_mod = cls(
|
||||
mod.input_size,
|
||||
mod.hidden_size,
|
||||
mod.bias,
|
||||
mod.weight_ih.device,
|
||||
mod.weight_ih.dtype,
|
||||
weight_qparams_dict,
|
||||
)
|
||||
ref_mod.weight_ih = mod.weight_ih
|
||||
ref_mod.weight_hh = mod.weight_hh
|
||||
ref_mod.bias_ih = mod.bias_ih
|
||||
ref_mod.bias_hh = mod.bias_hh
|
||||
return ref_mod
|
||||
|
||||
|
||||
class RNNBase(nn.RNNBase):
|
||||
def __init__(
|
||||
self,
|
||||
mode: str,
|
||||
input_size: int,
|
||||
hidden_size: int,
|
||||
num_layers: int = 1,
|
||||
bias: bool = True,
|
||||
batch_first: bool = False,
|
||||
dropout: float = 0.0,
|
||||
bidirectional: bool = False,
|
||||
proj_size: int = 0,
|
||||
device=None,
|
||||
dtype=None,
|
||||
weight_qparams_dict: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
mode,
|
||||
input_size,
|
||||
hidden_size,
|
||||
num_layers,
|
||||
bias,
|
||||
batch_first,
|
||||
dropout,
|
||||
bidirectional,
|
||||
proj_size,
|
||||
device,
|
||||
dtype,
|
||||
)
|
||||
# TODO(jerryzh168): maybe make this arg a required arg
|
||||
if weight_qparams_dict is None:
|
||||
weight_qparams = {
|
||||
"qscheme": torch.per_tensor_affine,
|
||||
"dtype": torch.quint8,
|
||||
"scale": 1.0,
|
||||
"zero_point": 0,
|
||||
}
|
||||
weight_qparams_dict = {"is_decomposed": False} # type: ignore[dict-item]
|
||||
for wn in self._flat_weights_names:
|
||||
if wn.startswith("weight"):
|
||||
weight_qparams_dict[wn] = weight_qparams
|
||||
self._init_weight_qparams_dict(weight_qparams_dict, device)
|
||||
|
||||
def _init_weight_qparams_dict(self, weight_qparams_dict, device):
|
||||
self.is_decomposed = weight_qparams_dict["is_decomposed"]
|
||||
for key, weight_qparams in weight_qparams_dict.items():
|
||||
if key == "is_decomposed":
|
||||
continue
|
||||
weight_qscheme = weight_qparams["qscheme"]
|
||||
weight_dtype = weight_qparams["dtype"]
|
||||
setattr(self, key + "_qscheme", weight_qscheme)
|
||||
setattr(self, key + "_dtype", weight_dtype)
|
||||
if weight_qscheme not in [
|
||||
None,
|
||||
torch.per_tensor_affine,
|
||||
torch.per_channel_affine,
|
||||
]:
|
||||
raise AssertionError(
|
||||
f"qscheme: {weight_qscheme} is not supported in {self._get_name()}"
|
||||
)
|
||||
if weight_qscheme is not None:
|
||||
self.register_buffer(
|
||||
key + "_scale",
|
||||
torch.tensor(
|
||||
weight_qparams["scale"], dtype=torch.float, device=device
|
||||
),
|
||||
)
|
||||
self.register_buffer(
|
||||
key + "_zero_point",
|
||||
torch.tensor(
|
||||
weight_qparams["zero_point"], dtype=torch.int, device=device
|
||||
),
|
||||
)
|
||||
if weight_qscheme == torch.per_channel_affine:
|
||||
self.register_buffer(
|
||||
key + "_axis",
|
||||
torch.tensor(
|
||||
weight_qparams["axis"], dtype=torch.int, device=device
|
||||
),
|
||||
)
|
||||
else:
|
||||
# added for TorchScriptability, not used
|
||||
self.register_buffer(
|
||||
key + "_axis", torch.tensor(0, dtype=torch.int, device=device)
|
||||
)
|
||||
setattr(self, key + "_axis_int", getattr(self, key + "_axis").item())
|
||||
|
||||
|
||||
class LSTM(RNNBase):
|
||||
"""Reference Quantized LSTM Module
|
||||
We'll store weight_qparams for all the weights in _flat_weights, we need to pass in
|
||||
a `weight_qparams_dict` that maps from weight name, e.g. weight_ih_l0,
|
||||
to the weight_qparams for that weight
|
||||
"""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__("LSTM", *args, **kwargs)
|
||||
|
||||
# Same as above, see torch/nn/modules/module.py::_forward_unimplemented
|
||||
def permute_hidden( # type: ignore[override]
|
||||
self,
|
||||
hx: tuple[Tensor, Tensor],
|
||||
permutation: Tensor | None,
|
||||
) -> tuple[Tensor, Tensor]:
|
||||
if permutation is None:
|
||||
return hx
|
||||
return _apply_permutation(hx[0], permutation), _apply_permutation(
|
||||
hx[1], permutation
|
||||
)
|
||||
|
||||
def get_expected_cell_size(
|
||||
self, input: Tensor, batch_sizes: Tensor | None
|
||||
) -> tuple[int, int, int]:
|
||||
if batch_sizes is not None:
|
||||
mini_batch = int(batch_sizes[0])
|
||||
else:
|
||||
mini_batch = input.size(0) if self.batch_first else input.size(1)
|
||||
num_directions = 2 if self.bidirectional else 1
|
||||
expected_hidden_size = (
|
||||
self.num_layers * num_directions,
|
||||
mini_batch,
|
||||
self.hidden_size,
|
||||
)
|
||||
return expected_hidden_size
|
||||
|
||||
# In the future, we should prevent mypy from applying contravariance rules here.
|
||||
# See torch/nn/modules/module.py::_forward_unimplemented
|
||||
def check_forward_args( # type: ignore[override]
|
||||
self,
|
||||
input: Tensor,
|
||||
hidden: tuple[Tensor, Tensor],
|
||||
batch_sizes: Tensor | None,
|
||||
):
|
||||
self.check_input(input, batch_sizes)
|
||||
self.check_hidden_size(
|
||||
hidden[0],
|
||||
self.get_expected_hidden_size(input, batch_sizes),
|
||||
"Expected hidden[0] size {}, got {}",
|
||||
)
|
||||
self.check_hidden_size(
|
||||
hidden[1],
|
||||
self.get_expected_cell_size(input, batch_sizes),
|
||||
"Expected hidden[1] size {}, got {}",
|
||||
)
|
||||
|
||||
def get_quantized_weight_bias_dict(self):
|
||||
"""dictionary from flat_weight_name to quantized weight or (unquantized) bias
|
||||
e.g.
|
||||
{
|
||||
"weight_ih_l0": quantized_weight,
|
||||
"bias_ih_l0": unquantized_bias,
|
||||
...
|
||||
}
|
||||
"""
|
||||
quantized_weight_bias_dict = {}
|
||||
for wn in self._flat_weights_names:
|
||||
if hasattr(self, wn):
|
||||
if wn.startswith("weight"):
|
||||
weight_or_bias = get_quantized_weight(self, wn)
|
||||
else:
|
||||
weight_or_bias = getattr(self, wn)
|
||||
else:
|
||||
weight_or_bias = None
|
||||
quantized_weight_bias_dict[wn] = weight_or_bias
|
||||
return quantized_weight_bias_dict
|
||||
|
||||
def get_flat_weights(self):
|
||||
flat_weights = []
|
||||
for wn in self._flat_weights_names:
|
||||
if hasattr(self, wn):
|
||||
weight = getattr(self, wn)
|
||||
if wn.startswith("weight"):
|
||||
params = _get_weight_and_quantization_params(self, wn)
|
||||
weight = _quantize_and_dequantize_weight(*params)
|
||||
else:
|
||||
weight = None
|
||||
flat_weights.append(weight)
|
||||
return flat_weights
|
||||
|
||||
def forward(self, input, hx=None): # noqa: F811
|
||||
orig_input = input
|
||||
# xxx: isinstance check needs to be in conditional for TorchScript to compile
|
||||
batch_sizes = None
|
||||
if isinstance(orig_input, PackedSequence):
|
||||
input, batch_sizes, sorted_indices, unsorted_indices = input
|
||||
max_batch_size = int(batch_sizes[0])
|
||||
else:
|
||||
batch_sizes = None
|
||||
is_batched = input.dim() == 3
|
||||
batch_dim = 0 if self.batch_first else 1
|
||||
if not is_batched:
|
||||
input = input.unsqueeze(batch_dim)
|
||||
max_batch_size = input.size(0) if self.batch_first else input.size(1)
|
||||
sorted_indices = None
|
||||
unsorted_indices = None
|
||||
|
||||
if hx is None:
|
||||
num_directions = 2 if self.bidirectional else 1
|
||||
real_hidden_size = (
|
||||
self.proj_size if self.proj_size > 0 else self.hidden_size
|
||||
)
|
||||
h_zeros = torch.zeros(
|
||||
self.num_layers * num_directions,
|
||||
max_batch_size,
|
||||
real_hidden_size,
|
||||
dtype=input.dtype,
|
||||
device=input.device,
|
||||
)
|
||||
c_zeros = torch.zeros(
|
||||
self.num_layers * num_directions,
|
||||
max_batch_size,
|
||||
self.hidden_size,
|
||||
dtype=input.dtype,
|
||||
device=input.device,
|
||||
)
|
||||
hx = (h_zeros, c_zeros)
|
||||
else:
|
||||
if batch_sizes is None: # If not PackedSequence input.
|
||||
if is_batched: # type: ignore[possibly-undefined]
|
||||
if hx[0].dim() != 3 or hx[1].dim() != 3:
|
||||
msg = (
|
||||
"For batched 3-D input, hx and cx should "
|
||||
f"also be 3-D but got ({hx[0].dim()}-D, {hx[1].dim()}-D) tensors"
|
||||
)
|
||||
raise RuntimeError(msg)
|
||||
else:
|
||||
if hx[0].dim() != 2 or hx[1].dim() != 2:
|
||||
msg = (
|
||||
"For unbatched 2-D input, hx and cx should "
|
||||
f"also be 2-D but got ({hx[0].dim()}-D, {hx[1].dim()}-D) tensors"
|
||||
)
|
||||
raise RuntimeError(msg)
|
||||
hx = (hx[0].unsqueeze(1), hx[1].unsqueeze(1))
|
||||
|
||||
# Each batch of the hidden state should match the input sequence that
|
||||
# the user believes he/she is passing in.
|
||||
hx = self.permute_hidden(hx, sorted_indices)
|
||||
|
||||
self.check_forward_args(input, hx, batch_sizes)
|
||||
if batch_sizes is None:
|
||||
result = _VF.lstm(
|
||||
input,
|
||||
hx,
|
||||
self.get_flat_weights(),
|
||||
self.bias,
|
||||
self.num_layers,
|
||||
self.dropout,
|
||||
self.training,
|
||||
self.bidirectional,
|
||||
self.batch_first,
|
||||
)
|
||||
else:
|
||||
result = _VF.lstm(
|
||||
input,
|
||||
batch_sizes,
|
||||
hx,
|
||||
self.get_flat_weights(),
|
||||
self.bias,
|
||||
self.num_layers,
|
||||
self.dropout,
|
||||
self.training,
|
||||
self.bidirectional,
|
||||
)
|
||||
output = result[0]
|
||||
hidden = result[1:]
|
||||
# xxx: isinstance check needs to be in conditional for TorchScript to compile
|
||||
if isinstance(orig_input, PackedSequence):
|
||||
output_packed = PackedSequence(
|
||||
output,
|
||||
batch_sizes,
|
||||
sorted_indices,
|
||||
unsorted_indices,
|
||||
)
|
||||
return output_packed, self.permute_hidden(hidden, unsorted_indices)
|
||||
else:
|
||||
if not is_batched: # type: ignore[possibly-undefined]
|
||||
output = output.squeeze(batch_dim) # type: ignore[possibly-undefined]
|
||||
hidden = (hidden[0].squeeze(1), hidden[1].squeeze(1))
|
||||
return output, self.permute_hidden(hidden, unsorted_indices)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedLSTM(Reference)"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, weight_qparams_dict):
|
||||
ref_mod = cls(
|
||||
mod.input_size,
|
||||
mod.hidden_size,
|
||||
mod.num_layers,
|
||||
mod.bias,
|
||||
mod.batch_first,
|
||||
mod.dropout,
|
||||
mod.bidirectional,
|
||||
weight_qparams_dict=weight_qparams_dict,
|
||||
)
|
||||
for wn in mod._flat_weights_names:
|
||||
setattr(ref_mod, wn, getattr(mod, wn))
|
||||
return ref_mod
|
||||
|
||||
|
||||
class GRU(RNNBase):
|
||||
"""Reference Quantized GRU Module
|
||||
We'll store weight_qparams for all the weights in _flat_weights, we need to pass in
|
||||
a `weight_qparams_dict` that maps from weight name, e.g. weight_ih_l0,
|
||||
to the weight_qparams for that weight
|
||||
"""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
if "proj_size" in kwargs:
|
||||
raise ValueError(
|
||||
"proj_size argument is only supported for LSTM, not RNN or GRU"
|
||||
)
|
||||
super().__init__("GRU", *args, **kwargs)
|
||||
|
||||
def get_quantized_weight_bias_dict(self):
|
||||
"""dictionary from flat_weight_name to quantized weight or (unquantized) bias
|
||||
e.g.
|
||||
{
|
||||
"weight_ih_l0": quantized_weight,
|
||||
"bias_ih_l0": unquantized_bias,
|
||||
...
|
||||
}
|
||||
"""
|
||||
quantized_weight_bias_dict = {}
|
||||
for wn in self._flat_weights_names:
|
||||
if hasattr(self, wn):
|
||||
if wn.startswith("weight"):
|
||||
weight_or_bias = get_quantized_weight(self, wn)
|
||||
else:
|
||||
weight_or_bias = getattr(self, wn)
|
||||
else:
|
||||
weight_or_bias = None
|
||||
quantized_weight_bias_dict[wn] = weight_or_bias
|
||||
return quantized_weight_bias_dict
|
||||
|
||||
def get_flat_weights(self):
|
||||
flat_weights = []
|
||||
for wn in self._flat_weights_names:
|
||||
if hasattr(self, wn):
|
||||
weight = getattr(self, wn)
|
||||
if wn.startswith("weight"):
|
||||
params = _get_weight_and_quantization_params(self, wn)
|
||||
weight = _quantize_and_dequantize_weight(*params)
|
||||
else:
|
||||
weight = None
|
||||
flat_weights.append(weight)
|
||||
return flat_weights
|
||||
|
||||
def forward(self, input, hx=None): # noqa: F811
|
||||
# Note: this is copied from the forward of GRU in https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/rnn.py
|
||||
# only changed self._flat_weights to self.get_flat_weights()
|
||||
# TODO: maybe we can try inheriting from that class and define get_flat_weights
|
||||
# as a @property? this might interfere with TorchScript, if we remove that
|
||||
# requirement in the future we should be able to do this
|
||||
orig_input = input
|
||||
# xxx: isinstance check needs to be in conditional for TorchScript to compile
|
||||
if isinstance(orig_input, PackedSequence):
|
||||
input, batch_sizes, sorted_indices, unsorted_indices = input
|
||||
max_batch_size = int(batch_sizes[0])
|
||||
else:
|
||||
batch_sizes = None
|
||||
if input.dim() not in (2, 3):
|
||||
raise AssertionError(
|
||||
f"GRU: Expected input to be 2-D or 3-D but received {input.dim()}-D tensor"
|
||||
)
|
||||
is_batched = input.dim() == 3
|
||||
batch_dim = 0 if self.batch_first else 1
|
||||
if not is_batched:
|
||||
input = input.unsqueeze(batch_dim)
|
||||
if hx is not None:
|
||||
if hx.dim() != 2:
|
||||
raise RuntimeError(
|
||||
f"For unbatched 2-D input, hx should also be 2-D but got {hx.dim()}-D tensor"
|
||||
)
|
||||
hx = hx.unsqueeze(1)
|
||||
else:
|
||||
if hx is not None and hx.dim() != 3:
|
||||
raise RuntimeError(
|
||||
f"For batched 3-D input, hx should also be 3-D but got {hx.dim()}-D tensor"
|
||||
)
|
||||
max_batch_size = input.size(0) if self.batch_first else input.size(1)
|
||||
sorted_indices = None
|
||||
unsorted_indices = None
|
||||
|
||||
if hx is None:
|
||||
num_directions = 2 if self.bidirectional else 1
|
||||
hx = torch.zeros(
|
||||
self.num_layers * num_directions,
|
||||
max_batch_size,
|
||||
self.hidden_size,
|
||||
dtype=input.dtype,
|
||||
device=input.device,
|
||||
)
|
||||
else:
|
||||
# Each batch of the hidden state should match the input sequence that
|
||||
# the user believes he/she is passing in.
|
||||
hx = self.permute_hidden(hx, sorted_indices)
|
||||
|
||||
self.check_forward_args(input, hx, batch_sizes)
|
||||
if batch_sizes is None:
|
||||
result = _VF.gru(
|
||||
input,
|
||||
hx,
|
||||
self.get_flat_weights(),
|
||||
self.bias,
|
||||
self.num_layers,
|
||||
self.dropout,
|
||||
self.training,
|
||||
self.bidirectional,
|
||||
self.batch_first,
|
||||
)
|
||||
else:
|
||||
result = _VF.gru(
|
||||
input,
|
||||
batch_sizes,
|
||||
hx,
|
||||
self.get_flat_weights(),
|
||||
self.bias,
|
||||
self.num_layers,
|
||||
self.dropout,
|
||||
self.training,
|
||||
self.bidirectional,
|
||||
)
|
||||
output = result[0]
|
||||
hidden = result[1]
|
||||
|
||||
# xxx: isinstance check needs to be in conditional for TorchScript to compile
|
||||
if isinstance(orig_input, PackedSequence):
|
||||
output_packed = PackedSequence(
|
||||
output,
|
||||
batch_sizes,
|
||||
sorted_indices,
|
||||
unsorted_indices,
|
||||
)
|
||||
return output_packed, self.permute_hidden(hidden, unsorted_indices)
|
||||
else:
|
||||
if not is_batched: # type: ignore[possibly-undefined]
|
||||
output = output.squeeze(batch_dim) # type: ignore[possibly-undefined]
|
||||
hidden = hidden.squeeze(1)
|
||||
|
||||
return output, self.permute_hidden(hidden, unsorted_indices)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedGRU(Reference)"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, weight_qparams_dict):
|
||||
ref_mod = cls(
|
||||
mod.input_size,
|
||||
mod.hidden_size,
|
||||
mod.num_layers,
|
||||
mod.bias,
|
||||
mod.batch_first,
|
||||
mod.dropout,
|
||||
mod.bidirectional,
|
||||
weight_qparams_dict=weight_qparams_dict,
|
||||
)
|
||||
for wn in mod._flat_weights_names:
|
||||
setattr(ref_mod, wn, getattr(mod, wn))
|
||||
return ref_mod
|
||||
+163
@@ -0,0 +1,163 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from typing import Any
|
||||
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
|
||||
from .utils import ReferenceQuantizedModule
|
||||
|
||||
|
||||
__all__ = ["Embedding", "EmbeddingBag"]
|
||||
|
||||
|
||||
class Embedding(nn.Embedding, ReferenceQuantizedModule):
|
||||
"""A reference quantized Embedding module that fits into the
|
||||
FX Graph Mode Quantization workflow, activation will be floating point Tensor,
|
||||
we will store floating point weight as well in the module, but in forward we'll
|
||||
quantize and dequantize the weight before running the floating point functional
|
||||
embedding operator.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_embeddings: int,
|
||||
embedding_dim: int,
|
||||
padding_idx: int | None = None,
|
||||
max_norm: float | None = None,
|
||||
norm_type: float = 2.0,
|
||||
scale_grad_by_freq: bool = False,
|
||||
sparse: bool = False,
|
||||
_weight: Tensor | None = None,
|
||||
device=None,
|
||||
dtype=None,
|
||||
weight_qparams: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
num_embeddings,
|
||||
embedding_dim,
|
||||
padding_idx,
|
||||
max_norm,
|
||||
norm_type,
|
||||
scale_grad_by_freq,
|
||||
sparse,
|
||||
_weight,
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
device,
|
||||
dtype,
|
||||
)
|
||||
self._init_weight_qparams(weight_qparams, device)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedEmbedding(Reference)"
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
weight_quant_dequant = self.get_weight()
|
||||
return F.embedding(
|
||||
input,
|
||||
weight_quant_dequant,
|
||||
self.padding_idx,
|
||||
self.max_norm,
|
||||
self.norm_type,
|
||||
self.scale_grad_by_freq,
|
||||
self.sparse,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, weight_qparams):
|
||||
return cls(
|
||||
mod.num_embeddings,
|
||||
mod.embedding_dim,
|
||||
mod.padding_idx,
|
||||
mod.max_norm,
|
||||
mod.norm_type,
|
||||
mod.scale_grad_by_freq,
|
||||
mod.sparse,
|
||||
mod.weight,
|
||||
mod.weight.device,
|
||||
mod.weight.dtype,
|
||||
weight_qparams,
|
||||
)
|
||||
|
||||
|
||||
class EmbeddingBag(nn.EmbeddingBag, ReferenceQuantizedModule):
|
||||
"""A reference quantized EmbeddingBag module that fits into the
|
||||
FX Graph Mode Quantization workflow, activation will be floating point Tensor,
|
||||
we will store floating point weight as well in the module, but in forward we'll
|
||||
quantize and dequantize the weight before running the floating point functional
|
||||
embedding operator.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_embeddings: int,
|
||||
embedding_dim: int,
|
||||
max_norm: float | None = None,
|
||||
norm_type: float = 2.0,
|
||||
scale_grad_by_freq: bool = False,
|
||||
mode: str = "mean",
|
||||
sparse: bool = False,
|
||||
_weight: Tensor | None = None,
|
||||
include_last_offset: bool = False,
|
||||
padding_idx: int | None = None,
|
||||
device=None,
|
||||
dtype=None,
|
||||
weight_qparams: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
num_embeddings,
|
||||
embedding_dim,
|
||||
max_norm,
|
||||
norm_type,
|
||||
scale_grad_by_freq,
|
||||
mode,
|
||||
sparse,
|
||||
_weight,
|
||||
include_last_offset,
|
||||
padding_idx,
|
||||
device,
|
||||
dtype,
|
||||
)
|
||||
self._init_weight_qparams(weight_qparams, device)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizedEmbedding(Reference)"
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input: Tensor,
|
||||
offsets: Tensor | None = None,
|
||||
per_sample_weights: Tensor | None = None,
|
||||
) -> Tensor:
|
||||
weight_quant_dequant = self.get_weight()
|
||||
return F.embedding_bag(
|
||||
input,
|
||||
weight_quant_dequant,
|
||||
offsets,
|
||||
self.max_norm,
|
||||
self.norm_type,
|
||||
self.scale_grad_by_freq,
|
||||
self.mode,
|
||||
self.sparse,
|
||||
per_sample_weights,
|
||||
self.include_last_offset,
|
||||
self.padding_idx,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, weight_qparams, use_precomputed_fake_quant=False):
|
||||
return cls(
|
||||
mod.num_embeddings,
|
||||
mod.embedding_dim,
|
||||
mod.max_norm,
|
||||
mod.norm_type,
|
||||
mod.scale_grad_by_freq,
|
||||
mod.mode,
|
||||
mod.sparse,
|
||||
mod.weight,
|
||||
mod.include_last_offset,
|
||||
mod.padding_idx,
|
||||
mod.weight.device,
|
||||
mod.weight.dtype,
|
||||
weight_qparams,
|
||||
)
|
||||
+439
@@ -0,0 +1,439 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import typing
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ReferenceQuantizedModule",
|
||||
]
|
||||
|
||||
|
||||
class ReferenceQuantizedModule(torch.nn.Module):
|
||||
def _init_weight_qparams(self, weight_qparams, device):
|
||||
if weight_qparams is None:
|
||||
weight_qparams = {
|
||||
"qscheme": torch.per_tensor_affine,
|
||||
"dtype": torch.quint8,
|
||||
"scale": 1.0,
|
||||
"zero_point": 0,
|
||||
}
|
||||
|
||||
self.weight_qscheme: torch.qscheme = weight_qparams["qscheme"]
|
||||
self.weight_dtype = weight_qparams["dtype"]
|
||||
if self.weight_qscheme not in [
|
||||
None,
|
||||
torch.per_tensor_affine,
|
||||
torch.per_channel_affine,
|
||||
torch.per_channel_affine_float_qparams,
|
||||
]:
|
||||
raise AssertionError(
|
||||
f"qscheme: {self.weight_qscheme} is not supported in reference quantized {self._get_name()}"
|
||||
)
|
||||
if self.weight_dtype in [
|
||||
torch.quint8,
|
||||
torch.qint8,
|
||||
torch.quint4x2,
|
||||
torch.qint32,
|
||||
]:
|
||||
zero_point_dtype = (
|
||||
weight_qparams["zero_point"].dtype
|
||||
if isinstance(weight_qparams["zero_point"], torch.Tensor)
|
||||
else torch.int
|
||||
)
|
||||
w_scale = weight_qparams["scale"]
|
||||
w_scale_tensor = (
|
||||
w_scale.detach().clone()
|
||||
if isinstance(w_scale, torch.Tensor)
|
||||
else torch.tensor(w_scale, dtype=torch.float, device=device)
|
||||
)
|
||||
self.register_buffer("weight_scale", w_scale_tensor)
|
||||
w_zp = weight_qparams["zero_point"]
|
||||
w_zp_tensor = (
|
||||
w_zp.detach().clone()
|
||||
if isinstance(w_zp, torch.Tensor)
|
||||
else torch.tensor(w_zp, dtype=zero_point_dtype, device=device)
|
||||
)
|
||||
self.register_buffer("weight_zero_point", w_zp_tensor)
|
||||
if self.weight_qscheme in [
|
||||
torch.per_channel_affine,
|
||||
torch.per_channel_affine_float_qparams,
|
||||
]:
|
||||
w_axis = weight_qparams["axis"]
|
||||
w_axis_tensor = (
|
||||
w_axis.detach().clone()
|
||||
if isinstance(w_axis, torch.Tensor)
|
||||
else torch.tensor(w_axis, dtype=torch.int, device=device)
|
||||
)
|
||||
self.register_buffer("weight_axis", w_axis_tensor)
|
||||
else:
|
||||
# added for TorchScriptability, not used
|
||||
self.register_buffer(
|
||||
"weight_axis", torch.tensor(0, dtype=torch.int, device=device)
|
||||
)
|
||||
else:
|
||||
# added for TorchScriptability, and for torch.float
|
||||
self.register_buffer(
|
||||
"weight_scale", torch.tensor(1.0, dtype=torch.float, device=device)
|
||||
)
|
||||
self.register_buffer(
|
||||
"weight_zero_point", torch.tensor(0, dtype=torch.int, device=device)
|
||||
)
|
||||
self.register_buffer(
|
||||
"weight_axis", torch.tensor(0, dtype=torch.int, device=device)
|
||||
)
|
||||
|
||||
self.is_decomposed: bool = weight_qparams.get("is_decomposed", False)
|
||||
# store weight_axis as weight_axis_int due to some constraints of torchdynamo.export
|
||||
# for capturing `.item` operations
|
||||
self.weight_axis_int: int = self.weight_axis.item() # type: ignore[operator, assignment]
|
||||
|
||||
self.weight_quant_min: int | None = weight_qparams.get("quant_min")
|
||||
|
||||
self.weight_quant_max: int | None = weight_qparams.get("quant_max")
|
||||
|
||||
def get_weight(self):
|
||||
"""
|
||||
Fake quantize (quantize and dequantize) the weight with
|
||||
the quantization parameters for weight, this is used to
|
||||
simulate the numerics for the quantized weight in a quantized
|
||||
model
|
||||
"""
|
||||
# suppress mypy warning
|
||||
if not isinstance(self.weight_scale, torch.Tensor):
|
||||
raise AssertionError("weight_scale must be a Tensor")
|
||||
if not isinstance(self.weight_zero_point, torch.Tensor):
|
||||
raise AssertionError("weight_zero_point must be a Tensor")
|
||||
if self.is_decomposed:
|
||||
return _quantize_and_dequantize_weight_decomposed(
|
||||
self.weight, # type: ignore[arg-type]
|
||||
self.weight_qscheme,
|
||||
self.weight_dtype,
|
||||
self.weight_scale,
|
||||
self.weight_zero_point,
|
||||
self.weight_axis_int,
|
||||
self.weight_quant_min,
|
||||
self.weight_quant_max,
|
||||
)
|
||||
else:
|
||||
return _quantize_and_dequantize_weight(
|
||||
self.weight, # type: ignore[arg-type]
|
||||
self.weight_qscheme,
|
||||
self.weight_dtype,
|
||||
self.weight_scale,
|
||||
self.weight_zero_point,
|
||||
self.weight_axis_int,
|
||||
)
|
||||
|
||||
def get_quantized_weight(self):
|
||||
# suppress mypy warning
|
||||
if not isinstance(self.weight_scale, torch.Tensor):
|
||||
raise AssertionError("weight_scale must be a Tensor")
|
||||
if not isinstance(self.weight_zero_point, torch.Tensor):
|
||||
raise AssertionError("weight_zero_point must be a Tensor")
|
||||
# assert isinstance(self.weight_axis, torch.Tensor)
|
||||
if self.is_decomposed:
|
||||
return _quantize_weight_decomposed(
|
||||
self.weight, # type: ignore[arg-type]
|
||||
self.weight_qscheme,
|
||||
self.weight_dtype,
|
||||
self.weight_scale,
|
||||
self.weight_zero_point,
|
||||
self.weight_axis_int,
|
||||
self.weight_quant_min,
|
||||
self.weight_quant_max,
|
||||
)
|
||||
else:
|
||||
return _quantize_weight(
|
||||
self.weight, # type: ignore[arg-type]
|
||||
self.weight_qscheme,
|
||||
self.weight_dtype,
|
||||
self.weight_scale,
|
||||
self.weight_zero_point,
|
||||
self.weight_axis_int,
|
||||
)
|
||||
|
||||
def _save_to_state_dict(self, destination, prefix, keep_vars):
|
||||
super()._save_to_state_dict(destination, prefix, keep_vars)
|
||||
_save_weight_qparams(
|
||||
destination,
|
||||
prefix,
|
||||
self.weight_qscheme,
|
||||
self.weight_dtype,
|
||||
self.weight_scale,
|
||||
self.weight_zero_point,
|
||||
self.weight_axis,
|
||||
)
|
||||
|
||||
def _load_from_state_dict(
|
||||
self,
|
||||
state_dict,
|
||||
prefix,
|
||||
local_metadata,
|
||||
strict,
|
||||
missing_keys,
|
||||
unexpected_keys,
|
||||
error_msgs,
|
||||
):
|
||||
for key in _get_weight_qparam_keys(state_dict, prefix):
|
||||
setattr(self, key, state_dict[prefix + key])
|
||||
state_dict.pop(prefix + key)
|
||||
|
||||
super()._load_from_state_dict(
|
||||
state_dict,
|
||||
prefix,
|
||||
local_metadata,
|
||||
False,
|
||||
missing_keys,
|
||||
unexpected_keys,
|
||||
error_msgs,
|
||||
)
|
||||
|
||||
|
||||
def _quantize_weight_decomposed(
|
||||
weight: torch.Tensor,
|
||||
weight_qscheme: torch.qscheme,
|
||||
weight_dtype: torch.dtype,
|
||||
weight_scale: torch.Tensor,
|
||||
weight_zero_point: torch.Tensor,
|
||||
weight_axis: int,
|
||||
weight_quant_min: int | None,
|
||||
weight_quant_max: int | None,
|
||||
) -> torch.Tensor:
|
||||
_DTYPE_TO_QVALUE_BOUNDS: dict[torch.dtype, tuple[int, int]] = {
|
||||
torch.uint8: (0, 255),
|
||||
torch.int8: (-128, 127),
|
||||
torch.int32: (-2147483648, 2147483647), # torch.jit interprets 2**31 as a float
|
||||
}
|
||||
|
||||
# TODO: add an util function for converting qdtype to dtype
|
||||
_QDTYPE_TO_UNDERLYING_INT_REPR_DTYPE = {
|
||||
torch.quint8: torch.uint8,
|
||||
torch.qint8: torch.int8,
|
||||
torch.qint32: torch.int32,
|
||||
}
|
||||
if weight_qscheme == torch.per_tensor_affine:
|
||||
if weight_dtype in [torch.quint8, torch.qint8, torch.qint32]:
|
||||
weight_dtype_ = _QDTYPE_TO_UNDERLYING_INT_REPR_DTYPE[weight_dtype]
|
||||
if weight_quant_min is None or weight_quant_max is None:
|
||||
weight_quant_min, weight_quant_max = _DTYPE_TO_QVALUE_BOUNDS[
|
||||
weight_dtype_
|
||||
]
|
||||
weight = torch.ops.quantized_decomposed.quantize_per_tensor(
|
||||
weight,
|
||||
weight_scale,
|
||||
weight_zero_point,
|
||||
weight_quant_min,
|
||||
weight_quant_max,
|
||||
weight_dtype_,
|
||||
)
|
||||
return weight
|
||||
elif weight_qscheme in [
|
||||
torch.per_channel_affine,
|
||||
torch.per_channel_affine_float_qparams,
|
||||
]:
|
||||
# TODO: torch.quint4x2 is not supported
|
||||
if weight_dtype in [torch.quint8, torch.qint8, torch.qint32]:
|
||||
weight_dtype_ = _QDTYPE_TO_UNDERLYING_INT_REPR_DTYPE[weight_dtype]
|
||||
if weight_quant_min is None or weight_quant_max is None:
|
||||
weight_quant_min, weight_quant_max = _DTYPE_TO_QVALUE_BOUNDS[
|
||||
weight_dtype_
|
||||
]
|
||||
weight = torch.ops.quantized_decomposed.quantize_per_channel(
|
||||
weight,
|
||||
weight_scale,
|
||||
weight_zero_point,
|
||||
weight_axis,
|
||||
weight_quant_min,
|
||||
weight_quant_max,
|
||||
weight_dtype_,
|
||||
) # type: ignore[arg-type]
|
||||
return weight
|
||||
raise ValueError(f"Unsupported dtype and qscheme: {weight_dtype}, {weight_qscheme}")
|
||||
|
||||
|
||||
def _dequantize_weight_decomposed(
|
||||
weight: torch.Tensor,
|
||||
weight_qscheme: torch.qscheme,
|
||||
weight_dtype: torch.dtype,
|
||||
weight_scale: torch.Tensor,
|
||||
weight_zero_point: torch.Tensor,
|
||||
weight_axis: int,
|
||||
weight_quant_min: int | None,
|
||||
weight_quant_max: int | None,
|
||||
) -> torch.Tensor:
|
||||
# TODO: get the quant_min and quant_max from activation_post_process
|
||||
_DTYPE_TO_QVALUE_BOUNDS: dict[torch.dtype, tuple[int, int]] = {
|
||||
torch.uint8: (0, 255),
|
||||
torch.int8: (-128, 127),
|
||||
torch.int32: (-2147483648, 2147483647), # torch.jit interprets 2**31 as a float
|
||||
}
|
||||
# TODO: add an util function for converting qdtype to dtype
|
||||
_QDTYPE_TO_UNDERLYING_INT_REPR_DTYPE = {
|
||||
torch.quint8: torch.uint8,
|
||||
torch.qint8: torch.int8,
|
||||
torch.qint32: torch.int32,
|
||||
}
|
||||
weight_dtype_ = _QDTYPE_TO_UNDERLYING_INT_REPR_DTYPE[weight_dtype]
|
||||
if weight_quant_min is None or weight_quant_max is None:
|
||||
weight_quant_min, weight_quant_max = _DTYPE_TO_QVALUE_BOUNDS[weight_dtype_]
|
||||
if weight_qscheme == torch.per_tensor_affine:
|
||||
if weight_dtype in [torch.quint8, torch.qint8, torch.qint32]:
|
||||
weight = torch.ops.quantized_decomposed.dequantize_per_tensor(
|
||||
weight,
|
||||
weight_scale,
|
||||
weight_zero_point,
|
||||
weight_quant_min,
|
||||
weight_quant_max,
|
||||
weight_dtype_,
|
||||
)
|
||||
return weight
|
||||
elif weight_qscheme in [
|
||||
torch.per_channel_affine,
|
||||
torch.per_channel_affine_float_qparams,
|
||||
]:
|
||||
# TODO: torch.quint4x2 is not supported
|
||||
if weight_dtype in [torch.quint8, torch.qint8, torch.qint32]:
|
||||
weight = torch.ops.quantized_decomposed.dequantize_per_channel(
|
||||
weight,
|
||||
weight_scale,
|
||||
weight_zero_point,
|
||||
weight_axis,
|
||||
weight_quant_min,
|
||||
weight_quant_max,
|
||||
weight_dtype_,
|
||||
) # type: ignore[arg-type]
|
||||
return weight
|
||||
raise ValueError(f"Unsupported dtype and qscheme: {weight_dtype}, {weight_qscheme}")
|
||||
|
||||
|
||||
def _quantize_weight(
|
||||
weight: torch.Tensor,
|
||||
weight_qscheme: torch.qscheme,
|
||||
weight_dtype: torch.dtype,
|
||||
weight_scale: torch.Tensor,
|
||||
weight_zero_point: torch.Tensor,
|
||||
weight_axis_int: int,
|
||||
) -> torch.Tensor:
|
||||
if weight_dtype == torch.float16:
|
||||
weight = weight.to(weight_dtype)
|
||||
return weight
|
||||
|
||||
if weight_qscheme == torch.per_tensor_affine:
|
||||
if weight_dtype in [torch.quint8, torch.qint8, torch.qint32]:
|
||||
weight = torch.quantize_per_tensor(
|
||||
weight, weight_scale, weight_zero_point, weight_dtype
|
||||
)
|
||||
return weight
|
||||
elif weight_qscheme in [
|
||||
torch.per_channel_affine,
|
||||
torch.per_channel_affine_float_qparams,
|
||||
]:
|
||||
if weight_dtype in [torch.quint8, torch.qint8, torch.quint4x2, torch.qint32]:
|
||||
weight = torch.quantize_per_channel(
|
||||
weight, weight_scale, weight_zero_point, weight_axis_int, weight_dtype
|
||||
) # type: ignore[arg-type]
|
||||
return weight
|
||||
raise ValueError(f"Unsupported dtype and qscheme: {weight_dtype}, {weight_qscheme}")
|
||||
|
||||
|
||||
def _quantize_and_dequantize_weight_decomposed(
|
||||
weight: torch.Tensor,
|
||||
weight_qscheme: torch.qscheme,
|
||||
weight_dtype: torch.dtype,
|
||||
weight_scale: torch.Tensor,
|
||||
weight_zero_point: torch.Tensor,
|
||||
weight_axis_int: int,
|
||||
weight_quant_min: int | None,
|
||||
weight_quant_max: int | None,
|
||||
) -> torch.Tensor:
|
||||
"""Quantize and then dequantize the weight based on
|
||||
the quantization parameters
|
||||
"""
|
||||
if weight_qscheme in [
|
||||
torch.per_tensor_affine,
|
||||
torch.per_channel_affine,
|
||||
torch.per_channel_affine_float_qparams,
|
||||
]:
|
||||
weight_quant = _quantize_weight_decomposed(
|
||||
weight,
|
||||
weight_qscheme,
|
||||
weight_dtype,
|
||||
weight_scale,
|
||||
weight_zero_point,
|
||||
weight_axis_int,
|
||||
weight_quant_min,
|
||||
weight_quant_max,
|
||||
)
|
||||
weight_dequant = _dequantize_weight_decomposed(
|
||||
weight_quant,
|
||||
weight_qscheme,
|
||||
weight_dtype,
|
||||
weight_scale,
|
||||
weight_zero_point,
|
||||
weight_axis_int,
|
||||
weight_quant_min,
|
||||
weight_quant_max,
|
||||
)
|
||||
else:
|
||||
weight_dequant = weight
|
||||
return weight_dequant
|
||||
|
||||
|
||||
def _quantize_and_dequantize_weight(
|
||||
weight: torch.Tensor,
|
||||
weight_qscheme: torch.qscheme,
|
||||
weight_dtype: torch.dtype,
|
||||
weight_scale: torch.Tensor,
|
||||
weight_zero_point: torch.Tensor,
|
||||
weight_axis_int: int,
|
||||
) -> torch.Tensor:
|
||||
"""Quantize and then dequantize the weight based on
|
||||
the quantization parameters
|
||||
"""
|
||||
if weight_qscheme in [
|
||||
torch.per_tensor_affine,
|
||||
torch.per_channel_affine,
|
||||
torch.per_channel_affine_float_qparams,
|
||||
]:
|
||||
weight_quant = _quantize_weight(
|
||||
weight,
|
||||
weight_qscheme,
|
||||
weight_dtype,
|
||||
weight_scale,
|
||||
weight_zero_point,
|
||||
weight_axis_int,
|
||||
)
|
||||
weight_dequant = weight_quant.dequantize()
|
||||
else:
|
||||
weight_dequant = weight
|
||||
return weight_dequant
|
||||
|
||||
|
||||
def _save_weight_qparams(
|
||||
destination,
|
||||
prefix,
|
||||
weight_qscheme,
|
||||
weight_dtype,
|
||||
weight_scale,
|
||||
weight_zero_point,
|
||||
weight_axis,
|
||||
):
|
||||
destination[prefix + "weight_qscheme"] = weight_qscheme
|
||||
destination[prefix + "weight_dtype"] = weight_dtype
|
||||
if weight_qscheme is not None:
|
||||
destination[prefix + "weight_scale"] = weight_scale
|
||||
destination[prefix + "weight_zero_point"] = weight_zero_point
|
||||
if weight_qscheme == torch.per_channel_affine:
|
||||
destination[prefix + "weight_axis"] = weight_axis
|
||||
|
||||
|
||||
def _get_weight_qparam_keys(state_dict: dict[str, typing.Any], prefix: str):
|
||||
keys = ["weight_qscheme", "weight_dtype"]
|
||||
weight_qscheme = state_dict[prefix + "weight_qscheme"]
|
||||
if weight_qscheme is not None:
|
||||
keys.append("weight_scale")
|
||||
keys.append("weight_zero_point")
|
||||
if weight_qscheme == torch.quantize_per_channel:
|
||||
keys.append("weight_axis")
|
||||
return keys
|
||||
@@ -0,0 +1 @@
|
||||
from . import quantized
|
||||
@@ -0,0 +1,10 @@
|
||||
from torch.ao.nn.sparse.quantized import dynamic
|
||||
|
||||
from .linear import Linear, LinearPackedParams
|
||||
|
||||
|
||||
__all__ = [
|
||||
"dynamic",
|
||||
"Linear",
|
||||
"LinearPackedParams",
|
||||
]
|
||||
+6
@@ -0,0 +1,6 @@
|
||||
from .linear import Linear
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Linear",
|
||||
]
|
||||
+203
@@ -0,0 +1,203 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import torch
|
||||
import torch.ao.nn.intrinsic as nni
|
||||
from torch.ao.nn.quantized.modules.utils import (
|
||||
_hide_packed_params_repr,
|
||||
_quantize_weight,
|
||||
)
|
||||
from torch.ao.nn.sparse.quantized import linear
|
||||
from torch.ao.nn.sparse.quantized.utils import LinearBlockSparsePattern
|
||||
|
||||
|
||||
__all__ = ["Linear"]
|
||||
|
||||
|
||||
class Linear(torch.nn.Module):
|
||||
r"""
|
||||
A dynamically quantized sparse linear module with float tensor as inputs and outputs.
|
||||
"""
|
||||
|
||||
_version = 1
|
||||
_op_type = "sparse_dynamic"
|
||||
_FLOAT_MODULE = torch.nn.Linear
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_features,
|
||||
out_features,
|
||||
row_block_size,
|
||||
col_block_size,
|
||||
bias=True,
|
||||
dtype=torch.qint8,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
if dtype != torch.qint8:
|
||||
raise NotImplementedError(
|
||||
"Only QINT8 is supported for Sparse Quantized Linear Dynamic"
|
||||
)
|
||||
|
||||
self.in_features = in_features
|
||||
self.out_features = out_features
|
||||
|
||||
if bias:
|
||||
bias = torch.zeros(self.out_features, dtype=torch.float)
|
||||
else:
|
||||
bias = None
|
||||
|
||||
qweight = torch._empty_affine_quantized(
|
||||
[out_features, in_features], scale=1, zero_point=0, dtype=torch.qint8
|
||||
)
|
||||
self._packed_params = linear.LinearPackedParams(
|
||||
row_block_size=row_block_size, col_block_size=col_block_size, dtype=dtype
|
||||
)
|
||||
self._packed_params.set_weight_bias(
|
||||
qweight, bias, row_block_size, col_block_size
|
||||
)
|
||||
|
||||
def _get_name(self):
|
||||
return "SparseQuantizedDynamicLinear"
|
||||
|
||||
def extra_repr(self):
|
||||
return f"in_features={self.in_features}, out_features={self.out_features}, qscheme={self.weight().qscheme()}"
|
||||
|
||||
def __repr__(self):
|
||||
return _hide_packed_params_repr(self, linear.LinearPackedParams)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return torch.ops.sparse.qlinear_dynamic(x, self._packed_params._packed_params)
|
||||
|
||||
def _save_to_state_dict(self, destination, prefix, keep_vars):
|
||||
super()._save_to_state_dict(destination, prefix, keep_vars)
|
||||
destination[prefix + "op_type"] = self._op_type
|
||||
|
||||
def _load_from_state_dict(
|
||||
self,
|
||||
state_dict,
|
||||
prefix,
|
||||
local_metadata,
|
||||
strict,
|
||||
missing_keys,
|
||||
unexpected_keys,
|
||||
error_msgs,
|
||||
):
|
||||
op_type = int(state_dict[prefix + "op_type"])
|
||||
if op_type != "sparse":
|
||||
raise AssertionError(
|
||||
f"Cannot load from op_type [{op_type}], expecting [{self._op_type}]"
|
||||
)
|
||||
state_dict.pop(prefix + "op_type")
|
||||
|
||||
version = local_metadata.get("version", None)
|
||||
if version is not None and version > self._version:
|
||||
raise AssertionError(f"version {version} > self._version {self._version}")
|
||||
|
||||
# Is this code valid? In old quantization it seemed to be used to load
|
||||
# older model
|
||||
weight = state_dict.pop(prefix + "weight")
|
||||
bias = state_dict.pop(prefix + "bias")
|
||||
state_dict.update(
|
||||
{
|
||||
prefix + "_packed_params.weight": weight,
|
||||
prefix + "_packed_params.bias": bias,
|
||||
}
|
||||
)
|
||||
|
||||
super()._load_from_state_dict(
|
||||
state_dict,
|
||||
prefix,
|
||||
local_metadata,
|
||||
False,
|
||||
missing_keys,
|
||||
unexpected_keys,
|
||||
error_msgs,
|
||||
)
|
||||
|
||||
def _weight_bias(self):
|
||||
return self._packed_params._weight_bias()
|
||||
|
||||
def weight(self):
|
||||
return self._weight_bias()[0]
|
||||
|
||||
def bias(self):
|
||||
return self._weight_bias()[1]
|
||||
|
||||
def set_weight_bias(
|
||||
self,
|
||||
w: torch.Tensor,
|
||||
b: torch.Tensor | None,
|
||||
row_block_size: int | None,
|
||||
col_block_size: int | None,
|
||||
) -> None:
|
||||
if row_block_size is None or col_block_size is None:
|
||||
raise AssertionError(
|
||||
f"row_block_size and col_block_size must not be None, got {row_block_size=}, {col_block_size=}"
|
||||
)
|
||||
self.out_features = w.shape[0]
|
||||
self.in_features = w.shape[1]
|
||||
self._packed_params.set_weight_bias(w, b, row_block_size, col_block_size)
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False):
|
||||
r"""Create a quantized sparse dynamic module from a float module.
|
||||
|
||||
We only care about the convert at this stage, no need for observers just yet.
|
||||
"""
|
||||
if type(mod) is not cls._FLOAT_MODULE:
|
||||
raise AssertionError(
|
||||
" nnq."
|
||||
+ cls.__name__
|
||||
+ ".from_float only works for "
|
||||
+ cls._FLOAT_MODULE.__name__
|
||||
)
|
||||
# TODO: Need to add options to qconfig to avoid the calibration.
|
||||
# TODO: Add calibration for the sparsity
|
||||
if not hasattr(mod, "qconfig"):
|
||||
raise AssertionError("Input float module must have qconfig defined")
|
||||
if type(mod) is nni.LinearReLU:
|
||||
mod = mod[0]
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
if mod.qconfig is not None and mod.qconfig.weight is not None:
|
||||
# pyrefly: ignore [not-callable]
|
||||
weight_observer = mod.qconfig.weight()
|
||||
else:
|
||||
# We have the circular import issues if we import the qconfig in the beginning of this file:
|
||||
# https://github.com/pytorch/pytorch/pull/24231. The current workaround is to postpone the
|
||||
# import until we need it.
|
||||
from torch.ao.quantization.qconfig import default_dynamic_qconfig
|
||||
|
||||
weight_observer = default_dynamic_qconfig.weight()
|
||||
|
||||
# It is important to multiply by the mask BEFORE calling the `weight_observer`
|
||||
# TODO (zaf): Mask might not be part of the qconfig (T83295194)
|
||||
weight = mod.weight
|
||||
if getattr(mod.qconfig, "mask", False):
|
||||
weight = mod.qconfig.mask * mod.weight
|
||||
|
||||
weight_observer(weight)
|
||||
dtype = weight_observer.dtype
|
||||
if dtype != torch.qint8:
|
||||
raise AssertionError(
|
||||
f"Weight observer must have dtype torch.qint8, got {dtype}"
|
||||
)
|
||||
_w_sc, w_zp = weight_observer.calculate_qparams()
|
||||
if isinstance(w_zp, torch.Tensor):
|
||||
if torch.any(w_zp.bool()):
|
||||
raise AssertionError("All weight zero points must map to 0")
|
||||
else:
|
||||
if w_zp != 0:
|
||||
raise AssertionError(f"Weight zero point must map to 0, got {w_zp}")
|
||||
qweight = _quantize_weight(weight.float(), weight_observer)
|
||||
|
||||
row_block_size, col_block_size = LinearBlockSparsePattern.block_size()
|
||||
qlinear = cls(
|
||||
mod.in_features,
|
||||
mod.out_features,
|
||||
row_block_size,
|
||||
col_block_size,
|
||||
dtype=dtype,
|
||||
)
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
qlinear.set_weight_bias(qweight, mod.bias, row_block_size, col_block_size)
|
||||
return qlinear
|
||||
@@ -0,0 +1,300 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import torch
|
||||
from torch.ao.nn.quantized.modules.utils import (
|
||||
_hide_packed_params_repr,
|
||||
_quantize_weight,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["LinearPackedParams", "Linear"]
|
||||
|
||||
|
||||
# TODO (zaf): Inherit from `quantized.LinearPackedParams` (T83294430)
|
||||
class LinearPackedParams(torch.nn.Module):
|
||||
_version = 1
|
||||
|
||||
def __init__(self, row_block_size=1, col_block_size=4, dtype=torch.qint8):
|
||||
super().__init__()
|
||||
|
||||
if dtype != torch.qint8:
|
||||
raise NotImplementedError("Linear prepacking only supports QINT8")
|
||||
self.dtype = dtype
|
||||
wq = torch._empty_affine_quantized(
|
||||
[1, 1], scale=1.0, zero_point=0, dtype=torch.qint8
|
||||
)
|
||||
self.set_weight_bias(wq, None, row_block_size, col_block_size)
|
||||
|
||||
def _get_name(self):
|
||||
return "SparseQuantizedLinearPackedParams"
|
||||
|
||||
@torch.jit.export
|
||||
def set_weight_bias(
|
||||
self,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor | None,
|
||||
row_block_size: int | None,
|
||||
col_block_size: int | None,
|
||||
) -> None:
|
||||
if row_block_size is None or col_block_size is None:
|
||||
raise AssertionError(
|
||||
"row_block_size and col_block_size must not be None, got "
|
||||
f"row_block_size={row_block_size}, col_block_size={col_block_size}"
|
||||
)
|
||||
self._packed_params = torch.ops.sparse.qlinear_prepack(
|
||||
weight, bias, row_block_size, col_block_size
|
||||
)
|
||||
|
||||
@torch.jit.export
|
||||
def _weight_bias(self):
|
||||
(weight, bias, block_sizes) = torch.ops.sparse.qlinear_unpack(
|
||||
self._packed_params
|
||||
)
|
||||
return (weight, bias, block_sizes[0], block_sizes[1])
|
||||
|
||||
def forward(self, x):
|
||||
return x
|
||||
|
||||
def _save_to_state_dict(self, destination, prefix, keep_vars):
|
||||
super()._save_to_state_dict(destination, prefix, keep_vars)
|
||||
destination[prefix + "dtype"] = self.dtype
|
||||
destination[prefix + "_packed_params"] = self._weight_bias()
|
||||
|
||||
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 not None and version > self._version:
|
||||
raise AssertionError(f"version {version} > self._version {self._version}")
|
||||
|
||||
self.dtype = state_dict.pop(prefix + "dtype")
|
||||
weight, bias, row_block_size, col_block_size = state_dict.pop(
|
||||
prefix + "_packed_params"
|
||||
)
|
||||
self.set_weight_bias(weight, bias, row_block_size, col_block_size)
|
||||
|
||||
super()._load_from_state_dict(
|
||||
state_dict,
|
||||
prefix,
|
||||
local_metadata,
|
||||
False,
|
||||
missing_keys,
|
||||
unexpected_keys,
|
||||
error_msgs,
|
||||
)
|
||||
|
||||
@torch.jit.export
|
||||
def __getstate__(self):
|
||||
return self._packed_params, self.training, self.dtype
|
||||
|
||||
@torch.jit.export
|
||||
def __setstate__(self, state):
|
||||
(self._packed_params, self.training, self.dtype) = state
|
||||
|
||||
def __repr__(self):
|
||||
return self._weight_bias().__repr__()
|
||||
|
||||
|
||||
# TODO (zaf): Inherit from `quantized.Linear` (T83294430)
|
||||
class Linear(torch.nn.Module):
|
||||
r"""
|
||||
A quantized sparse linear module with quantized tensor as inputs and outputs.
|
||||
"""
|
||||
|
||||
_version = 1
|
||||
_FLOAT_MODULE = torch.nn.Linear
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_features,
|
||||
out_features,
|
||||
row_block_size,
|
||||
col_block_size,
|
||||
bias=True,
|
||||
dtype=torch.qint8,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
if dtype != torch.qint8:
|
||||
raise NotImplementedError(
|
||||
"Only QINT8 is supported for Sparse Quantized Linear"
|
||||
)
|
||||
|
||||
self.in_features = in_features
|
||||
self.out_features = out_features
|
||||
|
||||
if bias:
|
||||
bias = torch.zeros(self.out_features, dtype=torch.float)
|
||||
else:
|
||||
bias = None
|
||||
|
||||
qweight = torch._empty_affine_quantized(
|
||||
[out_features, in_features], scale=1, zero_point=0, dtype=torch.qint8
|
||||
)
|
||||
self._packed_params = LinearPackedParams(
|
||||
row_block_size=row_block_size, col_block_size=col_block_size, dtype=dtype
|
||||
)
|
||||
self._packed_params.set_weight_bias(
|
||||
qweight, bias, row_block_size, col_block_size
|
||||
)
|
||||
self.scale = 1.0
|
||||
self.zero_point = 0
|
||||
|
||||
@classmethod
|
||||
def _get_name(cls):
|
||||
return "SparseQuantizedLinear"
|
||||
|
||||
def extra_repr(self):
|
||||
return (
|
||||
f"in_features={self.in_features}, out_features={self.out_features}, scale={self.scale}, "
|
||||
f"zero_point={self.zero_point}, qscheme={self.weight().qscheme()}"
|
||||
)
|
||||
|
||||
def __repr__(self):
|
||||
return _hide_packed_params_repr(self, LinearPackedParams)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return torch.ops.sparse.qlinear(
|
||||
x, self._packed_params._packed_params, self.scale, self.zero_point
|
||||
)
|
||||
|
||||
def _save_to_state_dict(self, destination, prefix, keep_vars):
|
||||
super()._save_to_state_dict(destination, prefix, keep_vars)
|
||||
destination[prefix + "scale"] = torch.tensor(self.scale)
|
||||
destination[prefix + "zero_point"] = torch.tensor(self.zero_point)
|
||||
|
||||
def _load_from_state_dict(
|
||||
self,
|
||||
state_dict,
|
||||
prefix,
|
||||
local_metadata,
|
||||
strict,
|
||||
missing_keys,
|
||||
unexpected_keys,
|
||||
error_msgs,
|
||||
):
|
||||
self.scale = float(state_dict[prefix + "scale"])
|
||||
state_dict.pop(prefix + "scale")
|
||||
|
||||
self.zero_point = int(state_dict[prefix + "zero_point"])
|
||||
state_dict.pop(prefix + "zero_point")
|
||||
|
||||
state_dict.pop(prefix + "op_type")
|
||||
|
||||
version = local_metadata.get("version", None)
|
||||
if version is not None and version > self._version:
|
||||
raise AssertionError(f"version {version} > self._version {self._version}")
|
||||
|
||||
super()._load_from_state_dict(
|
||||
state_dict,
|
||||
prefix,
|
||||
local_metadata,
|
||||
False,
|
||||
missing_keys,
|
||||
unexpected_keys,
|
||||
error_msgs,
|
||||
)
|
||||
|
||||
def _weight_bias(self):
|
||||
return self._packed_params._weight_bias()
|
||||
|
||||
def weight(self):
|
||||
return self._weight_bias()[0]
|
||||
|
||||
def bias(self):
|
||||
return self._weight_bias()[1]
|
||||
|
||||
def set_weight_bias(
|
||||
self,
|
||||
w: torch.Tensor,
|
||||
b: torch.Tensor | None,
|
||||
row_block_size: int | None,
|
||||
col_block_size: int | None,
|
||||
) -> None:
|
||||
if row_block_size is None or col_block_size is None:
|
||||
raise AssertionError(
|
||||
"row_block_size and col_block_size must not be None, "
|
||||
f"got row_block_size={row_block_size}, col_block_size={col_block_size}"
|
||||
)
|
||||
self._packed_params.set_weight_bias(w, b, row_block_size, col_block_size)
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, mod, use_precomputed_fake_quant=False):
|
||||
r"""Create a quantized sparse module from a float module.
|
||||
|
||||
We only care about the convert at this stage, no need for observers just yet.
|
||||
|
||||
TODO(zaf): Need to add the sparse params to the qconfig
|
||||
"""
|
||||
if type(mod) is not cls._FLOAT_MODULE:
|
||||
raise AssertionError(
|
||||
cls._get_name()
|
||||
+ ".from_float only works for "
|
||||
+ cls._FLOAT_MODULE.__name__
|
||||
)
|
||||
if not hasattr(mod, "sparse_params"):
|
||||
raise AssertionError(
|
||||
"Expecting the Linear to have `sparse_params`. Make sure you have provided arguments "
|
||||
'in the `sparsifier.squash_mask(params_to_save=("sparse_block_shape",))` method.'
|
||||
)
|
||||
sparse_block_shape = mod.sparse_params.get("sparse_block_shape", None) # type: ignore[operator, union-attr]
|
||||
if not isinstance(sparse_block_shape, (tuple, list)):
|
||||
raise AssertionError(
|
||||
f"sparse_block_shape must be tuple or list, got {type(sparse_block_shape)}"
|
||||
)
|
||||
if len(sparse_block_shape) != 2:
|
||||
raise AssertionError(
|
||||
f"sparse_block_shape must have length 2, got {len(sparse_block_shape)}"
|
||||
)
|
||||
# TODO: Need to add options to qconfig to avoid the calibration.
|
||||
# TODO: Add calibration for the sparsity
|
||||
if not hasattr(mod, "qconfig"):
|
||||
raise AssertionError("Input float module must have qconfig defined")
|
||||
activation_post_process = mod.activation_post_process
|
||||
weight_post_process = mod.qconfig.weight() # type: ignore[operator, union-attr]
|
||||
|
||||
# Assumption is that the weight is already sparsified by the
|
||||
# `sparsifier.convert`
|
||||
weight = mod.weight
|
||||
|
||||
weight_post_process(weight)
|
||||
dtype = weight_post_process.dtype
|
||||
act_scale, act_zp = activation_post_process.calculate_qparams() # type: ignore[operator, union-attr]
|
||||
if dtype != torch.qint8:
|
||||
raise AssertionError(
|
||||
f"Weight observer must have dtype torch.qint8, got {dtype}"
|
||||
)
|
||||
w_sc, w_zp = weight_post_process.calculate_qparams()
|
||||
if isinstance(w_zp, torch.Tensor):
|
||||
if torch.any(w_zp.bool()):
|
||||
raise AssertionError("All weight zero points must map to 0")
|
||||
else:
|
||||
if w_zp != 0:
|
||||
raise AssertionError(f"Weight zero point must map to 0, got {w_zp}")
|
||||
qweight = _quantize_weight(weight.float(), weight_post_process)
|
||||
|
||||
row_block_size = mod.sparse_params["sparse_block_shape"][0] # type: ignore[index]
|
||||
col_block_size = mod.sparse_params["sparse_block_shape"][1] # type: ignore[index]
|
||||
qlinear = cls(
|
||||
mod.in_features,
|
||||
mod.out_features,
|
||||
row_block_size,
|
||||
col_block_size,
|
||||
dtype=dtype,
|
||||
)
|
||||
qlinear.set_weight_bias(
|
||||
qweight,
|
||||
mod.bias,
|
||||
row_block_size, # type: ignore[arg-type]
|
||||
col_block_size, # type: ignore[arg-type]
|
||||
)
|
||||
qlinear.scale = float(act_scale)
|
||||
qlinear.zero_point = int(act_zp)
|
||||
return qlinear
|
||||
@@ -0,0 +1,66 @@
|
||||
import threading
|
||||
|
||||
|
||||
__all__ = ["LinearBlockSparsePattern"]
|
||||
|
||||
|
||||
def _is_valid_linear_block_sparse_pattern(
|
||||
row_block_size: int, col_block_size: int
|
||||
) -> bool:
|
||||
return (row_block_size == 1 and col_block_size == 4) or (
|
||||
row_block_size == 8 and col_block_size == 1
|
||||
)
|
||||
|
||||
|
||||
# This is a stop-gap measure as current flow does not allow module
|
||||
# specific block sparse pattern.
|
||||
# In fact there is no way to convey sparse pattern via module config
|
||||
# of quantization flow. Thus using the global context to convey
|
||||
# sparsity pattern.
|
||||
# Once the flow supports it, this should be removed.
|
||||
class LinearBlockSparsePattern:
|
||||
rlock = threading.RLock()
|
||||
row_block_size: int = 1
|
||||
col_block_size: int = 4
|
||||
prev_row_block_size: int = 1
|
||||
prev_col_block_size: int = 4
|
||||
|
||||
def __init__(self, row_block_size: int = 1, col_block_size: int = 4):
|
||||
if not _is_valid_linear_block_sparse_pattern(row_block_size, col_block_size):
|
||||
raise AssertionError(
|
||||
f"Invalid linear block sparse pattern: "
|
||||
f"row_block_size={row_block_size}, col_block_size={col_block_size}"
|
||||
)
|
||||
LinearBlockSparsePattern.rlock.acquire()
|
||||
LinearBlockSparsePattern.prev_row_block_size = (
|
||||
LinearBlockSparsePattern.row_block_size
|
||||
)
|
||||
LinearBlockSparsePattern.prev_col_block_size = (
|
||||
LinearBlockSparsePattern.col_block_size
|
||||
)
|
||||
LinearBlockSparsePattern.row_block_size = row_block_size
|
||||
LinearBlockSparsePattern.col_block_size = col_block_size
|
||||
|
||||
def __enter__(self) -> None:
|
||||
pass
|
||||
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_value: BaseException | None,
|
||||
backtrace: object | None,
|
||||
) -> None:
|
||||
LinearBlockSparsePattern.row_block_size = (
|
||||
LinearBlockSparsePattern.prev_row_block_size
|
||||
)
|
||||
LinearBlockSparsePattern.col_block_size = (
|
||||
LinearBlockSparsePattern.prev_col_block_size
|
||||
)
|
||||
LinearBlockSparsePattern.rlock.release()
|
||||
|
||||
@staticmethod
|
||||
def block_size() -> tuple[int, int]:
|
||||
return (
|
||||
LinearBlockSparsePattern.row_block_size,
|
||||
LinearBlockSparsePattern.col_block_size,
|
||||
)
|
||||
@@ -0,0 +1,568 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.ao.nn.quantized as nnq
|
||||
import torch.ao.nn.quantized.dynamic as nnqd
|
||||
import torch.nn as nn
|
||||
from torch.ao.quantization import prepare
|
||||
from torch.ao.quantization.quantization_mappings import (
|
||||
get_default_compare_output_module_list,
|
||||
)
|
||||
|
||||
|
||||
NON_LEAF_MODULE_TO_ADD_OBSERVER_ALLOW_LIST = {
|
||||
nnqd.Linear,
|
||||
nnq.Linear,
|
||||
nnqd.LSTM,
|
||||
nn.LSTM,
|
||||
}
|
||||
|
||||
|
||||
def _find_match(
|
||||
str_list: dict[str, Any] | list[str],
|
||||
key_str: str,
|
||||
postfix: str,
|
||||
) -> str | None:
|
||||
split_str = key_str.split(".")
|
||||
if split_str[-1] == postfix:
|
||||
match_string = "".join(key_str.split(".")[0:-1])
|
||||
for s2 in str_list:
|
||||
pattern1 = "".join(s2.split(".")[0:-1])
|
||||
pattern2 = "".join(s2.split(".")[0:-2])
|
||||
if match_string == pattern1:
|
||||
return s2
|
||||
if match_string == pattern2:
|
||||
return s2
|
||||
|
||||
# For matching "fc.weight" and "fc._packed_params._packed_params"
|
||||
if postfix == "_packed_params":
|
||||
match_string = "".join(key_str.split(".")[0:-2])
|
||||
if len(match_string) == 0:
|
||||
return None
|
||||
for s2 in str_list:
|
||||
pattern1 = "".join(s2.split(".")[0:-1])
|
||||
pattern2 = "".join(s2.split(".")[0:-2])
|
||||
if match_string == pattern1:
|
||||
return s2
|
||||
if match_string == pattern2:
|
||||
return s2
|
||||
return None
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def compare_weights(
|
||||
float_dict: dict[str, Any], quantized_dict: dict[str, Any]
|
||||
) -> dict[str, dict[str, torch.Tensor]]:
|
||||
r"""Compare the weights of the float module with its corresponding quantized
|
||||
module. Return a dict with key corresponding to module names and each entry being
|
||||
a dictionary with two keys 'float' and 'quantized', containing the float and
|
||||
quantized weights. This dict can be used to compare and compute the quantization
|
||||
error of the weights of float and quantized models.
|
||||
|
||||
Example usage::
|
||||
|
||||
wt_compare_dict = compare_weights(float_model.state_dict(), qmodel.state_dict())
|
||||
for key in wt_compare_dict:
|
||||
print(
|
||||
key,
|
||||
compute_error(
|
||||
wt_compare_dict[key]["float"],
|
||||
wt_compare_dict[key]["quantized"].dequantize(),
|
||||
),
|
||||
)
|
||||
|
||||
Args:
|
||||
float_dict: state dict of the float model
|
||||
quantized_dict: state dict of the quantized model
|
||||
|
||||
Return:
|
||||
weight_dict: dict with key corresponding to module names and each entry being
|
||||
a dictionary with two keys 'float' and 'quantized', containing the float and
|
||||
quantized weights
|
||||
"""
|
||||
torch._C._log_api_usage_once("quantization_api._numeric_suite.compare_weights")
|
||||
weight_dict: dict[str, dict] = {}
|
||||
for key in quantized_dict:
|
||||
match_key = _find_match(float_dict, key, "weight")
|
||||
if match_key is not None:
|
||||
weight_dict[key] = {}
|
||||
weight_dict[key]["float"] = float_dict[match_key]
|
||||
weight_dict[key]["quantized"] = quantized_dict[key]
|
||||
continue
|
||||
|
||||
# For matching "fc.weight" and "fc._packed_params._packed_params"
|
||||
match_key = _find_match(float_dict, key, "_packed_params")
|
||||
if match_key is not None:
|
||||
weight_dict[key] = {}
|
||||
weight_dict[key]["float"] = float_dict[match_key]
|
||||
weight_dict[key]["quantized"] = quantized_dict[key][0]
|
||||
|
||||
# For LSTM
|
||||
split_str = key.split(".")
|
||||
if split_str[-1] == "param" and split_str[-3] == "_all_weight_values":
|
||||
layer = split_str[-2]
|
||||
module_name = ".".join(split_str[:-3])
|
||||
float_weight_ih_key = module_name + ".weight_ih_l" + layer
|
||||
float_weight_hh_key = module_name + ".weight_hh_l" + layer
|
||||
if float_weight_ih_key in float_dict and float_weight_hh_key in float_dict:
|
||||
weight_dict[key] = {}
|
||||
weight_dict[key]["float"] = float_dict[float_weight_ih_key]
|
||||
weight_dict[key]["quantized"] = (
|
||||
quantized_dict[key].__getstate__()[0][4][0].__getstate__()[0][0]
|
||||
)
|
||||
weight_dict[key]["float"] = float_dict[float_weight_hh_key]
|
||||
weight_dict[key]["quantized"] = (
|
||||
quantized_dict[key].__getstate__()[0][4][1].__getstate__()[0][0]
|
||||
)
|
||||
|
||||
return weight_dict
|
||||
|
||||
|
||||
def _get_logger_dict_helper(
|
||||
mod: nn.Module,
|
||||
target_dict: dict[str, Any],
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
r"""This is the helper function for get_logger_dict
|
||||
|
||||
Args:
|
||||
mod: module we want to save all logger stats
|
||||
prefix: prefix for the current module
|
||||
target_dict: the dictionary used to save all logger stats
|
||||
"""
|
||||
|
||||
def get_prefix(prefix):
|
||||
return prefix if prefix == "" else prefix + "."
|
||||
|
||||
for child in mod.children():
|
||||
if isinstance(child, Logger):
|
||||
target_dict[get_prefix(prefix) + "stats"] = child.stats
|
||||
break
|
||||
|
||||
for name, child in mod.named_children():
|
||||
module_prefix = get_prefix(prefix) + name if prefix else name
|
||||
_get_logger_dict_helper(child, target_dict, module_prefix)
|
||||
|
||||
|
||||
def get_logger_dict(mod: nn.Module, prefix: str = "") -> dict[str, dict]:
|
||||
r"""Traverse the modules and save all logger stats into target dict.
|
||||
This is mainly used for quantization accuracy debug.
|
||||
|
||||
Type of loggers supported:
|
||||
ShadowLogger: used to log the outputs of the quantized module and its matching float shadow module,
|
||||
OutputLogger: used to log the outputs of the modules
|
||||
|
||||
Args:
|
||||
mod: module we want to save all logger stats
|
||||
prefix: prefix for the current module
|
||||
|
||||
Return:
|
||||
target_dict: the dictionary used to save all logger stats
|
||||
|
||||
"""
|
||||
torch._C._log_api_usage_once("quantization_api._numeric_suite.get_logger_dict")
|
||||
|
||||
target_dict: dict[str, dict] = {}
|
||||
_get_logger_dict_helper(mod, target_dict, prefix)
|
||||
return target_dict
|
||||
|
||||
|
||||
class Logger(nn.Module):
|
||||
r"""Base class for stats logging"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.stats = {}
|
||||
# We only insert observer if the op is quantized with static quantization,
|
||||
# which is identified by activation_observer.dtype == quint8. This is needed
|
||||
# when attaching Logger as observer for FX mode
|
||||
self.dtype = torch.quint8
|
||||
|
||||
def forward(self, x):
|
||||
# fmt: off
|
||||
"""
|
||||
""" # blank docblock to make autodoc happy
|
||||
# fmt: on
|
||||
|
||||
|
||||
class ShadowLogger(Logger):
|
||||
r"""Class used in Shadow module to record the outputs of the original and
|
||||
shadow modules.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.stats["float"] = []
|
||||
self.stats["quantized"] = []
|
||||
|
||||
def forward(self, x, y): # type: ignore[override]
|
||||
# fmt: off
|
||||
"""
|
||||
""" # blank docblock to make autodoc happy
|
||||
# fmt: on
|
||||
if len(x) > 1:
|
||||
x = x[0]
|
||||
if len(y) > 1:
|
||||
y = y[0]
|
||||
self.stats["quantized"].append(x.detach())
|
||||
self.stats["float"].append(y.detach())
|
||||
|
||||
|
||||
class OutputLogger(Logger):
|
||||
r"""Class used to log the outputs of the module"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.stats["tensor_val"] = []
|
||||
|
||||
def forward(self, x):
|
||||
# fmt: off
|
||||
"""
|
||||
""" # blank docblock to make autodoc happy
|
||||
# fmt: on
|
||||
self.stats["tensor_val"].append(x)
|
||||
return x
|
||||
|
||||
|
||||
def _convert_tuple_to_list(t: Any) -> Any:
|
||||
return [_convert_tuple_to_list(x) for x in t] if type(t) is tuple else t
|
||||
|
||||
|
||||
def _dequantize_tensor_list(t: Any) -> Any:
|
||||
return (
|
||||
[_dequantize_tensor_list(x) for x in t]
|
||||
if type(t) is list
|
||||
else t.dequantize()
|
||||
if t.is_quantized
|
||||
else t
|
||||
)
|
||||
|
||||
|
||||
class Shadow(nn.Module):
|
||||
r"""Shadow module attaches the float module to its matching quantized module
|
||||
as the shadow. Then it uses Logger module to process the outputs of both
|
||||
modules.
|
||||
|
||||
Args:
|
||||
q_module: module quantized from float_module that we want to shadow
|
||||
float_module: float module used to shadow q_module
|
||||
logger_cls: type of logger used to process the outputs of q_module and
|
||||
float_module. ShadowLogger or custom loggers can be used.
|
||||
"""
|
||||
|
||||
def __init__(self, q_module, float_module, logger_cls):
|
||||
super().__init__()
|
||||
self.orig_module = q_module
|
||||
self.shadow_module = float_module
|
||||
self.dequant = nnq.DeQuantize()
|
||||
self.logger = logger_cls()
|
||||
|
||||
def forward(self, *x) -> torch.Tensor:
|
||||
# fmt: off
|
||||
"""
|
||||
""" # blank docblock to make autodoc happy
|
||||
# fmt: on
|
||||
xl = _convert_tuple_to_list(x)
|
||||
output = self.orig_module(*xl)
|
||||
xl_float = _dequantize_tensor_list(xl)
|
||||
shadow_output = self.shadow_module(*xl_float)
|
||||
self.logger(output, shadow_output)
|
||||
return output
|
||||
|
||||
def add(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
# fmt: off
|
||||
"""
|
||||
""" # blank docblock to make autodoc happy
|
||||
# fmt: on
|
||||
output = self.orig_module.add(x, y)
|
||||
x = x.dequantize()
|
||||
y = y.dequantize()
|
||||
shadow_output = self.shadow_module.add(x, y)
|
||||
self.logger(output, shadow_output)
|
||||
return output
|
||||
|
||||
def add_scalar(self, x: torch.Tensor, y: float) -> torch.Tensor:
|
||||
# fmt: off
|
||||
"""
|
||||
""" # blank docblock to make autodoc happy
|
||||
# fmt: on
|
||||
output = self.orig_module.add_scalar(x, y)
|
||||
x = x.dequantize()
|
||||
shadow_output = self.shadow_module.add_scalar(x, y)
|
||||
self.logger(output, shadow_output)
|
||||
return output
|
||||
|
||||
def mul(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
# fmt: off
|
||||
"""
|
||||
""" # blank docblock to make autodoc happy
|
||||
# fmt: on
|
||||
output = self.orig_module.mul(x, y)
|
||||
x = x.dequantize()
|
||||
y = y.dequantize()
|
||||
shadow_output = self.shadow_module.mul(x, y)
|
||||
self.logger(output, shadow_output)
|
||||
return output
|
||||
|
||||
def mul_scalar(self, x: torch.Tensor, y: float) -> torch.Tensor:
|
||||
# fmt: off
|
||||
"""
|
||||
""" # blank docblock to make autodoc happy
|
||||
# fmt: on
|
||||
output = self.orig_module.mul_scalar(x, y)
|
||||
x = x.dequantize()
|
||||
shadow_output = self.shadow_module.mul_scalar(x, y)
|
||||
self.logger(output, shadow_output)
|
||||
return output
|
||||
|
||||
def cat(self, x: list[torch.Tensor], dim: int = 0) -> torch.Tensor:
|
||||
# fmt: off
|
||||
"""
|
||||
""" # blank docblock to make autodoc happy
|
||||
# fmt: on
|
||||
output = self.orig_module.cat(x, dim)
|
||||
x = [y.dequantize() for y in x]
|
||||
shadow_output = self.shadow_module.cat(x, dim)
|
||||
self.logger(output, shadow_output)
|
||||
return output
|
||||
|
||||
def add_relu(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
# fmt: off
|
||||
"""
|
||||
""" # blank docblock to make autodoc happy
|
||||
# fmt: on
|
||||
output = self.orig_module.add_relu(x, y)
|
||||
x = x.dequantize()
|
||||
y = y.dequantize()
|
||||
shadow_output = self.shadow_module.add_relu(x, y)
|
||||
self.logger(output, shadow_output)
|
||||
return output
|
||||
|
||||
|
||||
def prepare_model_with_stubs(
|
||||
float_module: nn.Module,
|
||||
q_module: nn.Module,
|
||||
module_swap_list: set[type],
|
||||
logger_cls: Callable,
|
||||
) -> None:
|
||||
r"""Prepare the model by attaching the float module to its matching quantized
|
||||
module as the shadow if the float module type is in module_swap_list.
|
||||
|
||||
Example usage::
|
||||
|
||||
prepare_model_with_stubs(float_model, q_model, module_swap_list, Logger)
|
||||
q_model(data)
|
||||
ob_dict = get_logger_dict(q_model)
|
||||
|
||||
Args:
|
||||
float_module: float module used to generate the q_module
|
||||
q_module: module quantized from float_module
|
||||
module_swap_list: list of float module types to attach the shadow
|
||||
logger_cls: type of logger to be used in shadow module to process the outputs of
|
||||
quantized module and its float shadow module
|
||||
"""
|
||||
torch._C._log_api_usage_once(
|
||||
"quantization_api._numeric_suite.prepare_model_with_stubs"
|
||||
)
|
||||
|
||||
float_module_children = dict(float_module.named_children())
|
||||
|
||||
reassign = {}
|
||||
for name, mod in q_module.named_children():
|
||||
if name not in float_module_children:
|
||||
continue
|
||||
|
||||
float_mod = float_module_children[name]
|
||||
|
||||
if type(float_mod) not in module_swap_list:
|
||||
prepare_model_with_stubs(float_mod, mod, module_swap_list, logger_cls)
|
||||
|
||||
# Insert shadow module only if the module is not of the same type as
|
||||
# the floating point module
|
||||
if type(float_mod) in module_swap_list and not _is_identical_module_type(
|
||||
mod, float_mod
|
||||
):
|
||||
reassign[name] = Shadow(mod, float_mod, logger_cls)
|
||||
|
||||
for key, value in reassign.items():
|
||||
q_module._modules[key] = value
|
||||
|
||||
|
||||
def _is_identical_module_type(mod1, mod2):
|
||||
# Compare if two modules have the same dtype
|
||||
mod1_module_types = [type(mod) for mod in mod1.modules()]
|
||||
mod2_module_types = [type(mod) for mod in mod2.modules()]
|
||||
return mod1_module_types == mod2_module_types
|
||||
|
||||
|
||||
def compare_model_stub(
|
||||
float_model: nn.Module,
|
||||
q_model: nn.Module,
|
||||
module_swap_list: set[type],
|
||||
*data,
|
||||
logger_cls=ShadowLogger,
|
||||
) -> dict[str, dict]:
|
||||
r"""Compare quantized module in a model with its floating point counterpart,
|
||||
feeding both of them the same input. Return a dict with key corresponding to
|
||||
module names and each entry being a dictionary with two keys 'float' and
|
||||
'quantized', containing the output tensors of quantized and its matching
|
||||
float shadow module. This dict can be used to compare and compute the module
|
||||
level quantization error.
|
||||
|
||||
This function first call prepare_model_with_stubs() to swap the quantized
|
||||
module that we want to compare with the Shadow module, which takes quantized
|
||||
module, corresponding float module and logger as input, and creates a forward
|
||||
path inside to make the float module to shadow quantized module sharing the
|
||||
same input. The logger can be customizable, default logger is ShadowLogger
|
||||
and it will save the outputs of the quantized module and float module that
|
||||
can be used to compute the module level quantization error.
|
||||
|
||||
Example usage::
|
||||
|
||||
module_swap_list = [
|
||||
torchvision.models.quantization.resnet.QuantizableBasicBlock
|
||||
]
|
||||
ob_dict = compare_model_stub(float_model, qmodel, module_swap_list, data)
|
||||
for key in ob_dict:
|
||||
print(
|
||||
key,
|
||||
compute_error(
|
||||
ob_dict[key]["float"], ob_dict[key]["quantized"].dequantize()
|
||||
),
|
||||
)
|
||||
|
||||
Args:
|
||||
float_model: float model used to generate the q_model
|
||||
q_model: model quantized from float_model
|
||||
module_swap_list: list of float module types at which shadow modules will
|
||||
be attached.
|
||||
data: input data used to run the prepared q_model
|
||||
logger_cls: type of logger to be used in shadow module to process the outputs of
|
||||
quantized module and its float shadow module
|
||||
"""
|
||||
torch._C._log_api_usage_once("quantization_api._numeric_suite.compare_model_stub")
|
||||
prepare_model_with_stubs(float_model, q_model, module_swap_list, logger_cls)
|
||||
q_model(*data)
|
||||
ob_dict = get_logger_dict(q_model)
|
||||
return ob_dict
|
||||
|
||||
|
||||
def get_matching_activations(
|
||||
float_module: nn.Module,
|
||||
q_module: nn.Module,
|
||||
) -> dict[str, dict[str, torch.Tensor]]:
|
||||
r"""Find the matching activation between float and quantized modules.
|
||||
|
||||
Args:
|
||||
float_module: float module used to generate the q_module
|
||||
q_module: module quantized from float_module
|
||||
|
||||
Return:
|
||||
act_dict: dict with key corresponding to quantized module names and each
|
||||
entry being a dictionary with two keys 'float' and 'quantized', containing
|
||||
the matching float and quantized activations
|
||||
"""
|
||||
torch._C._log_api_usage_once(
|
||||
"quantization_api._numeric_suite.get_matching_activations"
|
||||
)
|
||||
float_dict = get_logger_dict(float_module)
|
||||
quantized_dict = get_logger_dict(q_module)
|
||||
act_dict: dict[str, dict] = {}
|
||||
for key in quantized_dict:
|
||||
if len(quantized_dict[key]["tensor_val"]) == 0:
|
||||
continue
|
||||
match_key = _find_match(sorted(float_dict, reverse=True), key, "stats")
|
||||
if match_key is not None:
|
||||
act_dict[key] = {}
|
||||
act_dict[key]["float"] = float_dict[match_key]["tensor_val"]
|
||||
act_dict[key]["quantized"] = quantized_dict[key]["tensor_val"]
|
||||
return act_dict
|
||||
|
||||
|
||||
def prepare_model_outputs(
|
||||
float_module: nn.Module,
|
||||
q_module: nn.Module,
|
||||
logger_cls=OutputLogger,
|
||||
allow_list=None,
|
||||
) -> None:
|
||||
r"""Prepare the model by attaching the logger to both float module
|
||||
and quantized module if they are in the allow_list.
|
||||
|
||||
Args:
|
||||
float_module: float module used to generate the q_module
|
||||
q_module: module quantized from float_module
|
||||
logger_cls: type of logger to be attached to float_module and q_module
|
||||
allow_list: list of module types to attach logger
|
||||
"""
|
||||
torch._C._log_api_usage_once(
|
||||
"quantization_api._numeric_suite.prepare_model_outputs"
|
||||
)
|
||||
if allow_list is None:
|
||||
allow_list = get_default_compare_output_module_list()
|
||||
|
||||
qconfig_debug = torch.ao.quantization.QConfig(activation=logger_cls, weight=None)
|
||||
float_module.qconfig = qconfig_debug # type: ignore[assignment]
|
||||
prepare(
|
||||
float_module, inplace=True, allow_list=allow_list, prepare_custom_config_dict={}
|
||||
)
|
||||
q_module.qconfig = qconfig_debug # type: ignore[assignment]
|
||||
prepare(
|
||||
q_module,
|
||||
inplace=True,
|
||||
allow_list=allow_list,
|
||||
observer_non_leaf_module_list=NON_LEAF_MODULE_TO_ADD_OBSERVER_ALLOW_LIST,
|
||||
prepare_custom_config_dict={},
|
||||
)
|
||||
|
||||
|
||||
def compare_model_outputs(
|
||||
float_model: nn.Module,
|
||||
q_model: nn.Module,
|
||||
*data,
|
||||
logger_cls=OutputLogger,
|
||||
allow_list=None,
|
||||
) -> dict[str, dict[str, torch.Tensor]]:
|
||||
r"""Compare output activations between float and quantized models at
|
||||
corresponding locations for the same input. Return a dict with key corresponding
|
||||
to quantized module names and each entry being a dictionary with two keys
|
||||
'float' and 'quantized', containing the activations of quantized model and
|
||||
float model at matching locations. This dict can be used to compare and
|
||||
compute the propagation quantization error.
|
||||
|
||||
Example usage::
|
||||
|
||||
act_compare_dict = compare_model_outputs(float_model, qmodel, data)
|
||||
for key in act_compare_dict:
|
||||
print(
|
||||
key,
|
||||
compute_error(
|
||||
act_compare_dict[key]["float"],
|
||||
act_compare_dict[key]["quantized"].dequantize(),
|
||||
),
|
||||
)
|
||||
|
||||
Args:
|
||||
float_model: float model used to generate the q_model
|
||||
q_model: model quantized from float_model
|
||||
data: input data used to run the prepared float_model and q_model
|
||||
logger_cls: type of logger to be attached to float_module and q_module
|
||||
allow_list: list of module types to attach logger
|
||||
|
||||
Return:
|
||||
act_compare_dict: dict with key corresponding to quantized module names
|
||||
and each entry being a dictionary with two keys 'float' and 'quantized',
|
||||
containing the matching float and quantized activations
|
||||
"""
|
||||
torch._C._log_api_usage_once(
|
||||
"quantization_api._numeric_suite.compare_model_outputs"
|
||||
)
|
||||
if allow_list is None:
|
||||
allow_list = get_default_compare_output_module_list()
|
||||
prepare_model_outputs(float_model, q_model, logger_cls, allow_list)
|
||||
float_model(*data)
|
||||
q_model(*data)
|
||||
act_compare_dict = get_matching_activations(float_model, q_model)
|
||||
return act_compare_dict
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,485 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import collections
|
||||
import enum
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch.ao.quantization import FakeQuantizeBase, ObserverBase
|
||||
from torch.ao.quantization.utils import getattr_from_fqn
|
||||
from torch.fx import GraphModule
|
||||
from torch.fx.graph import Graph, Node
|
||||
|
||||
from .mappings import get_base_name_to_sets_of_related_ops, get_unmatchable_types_map
|
||||
from .ns_types import NSNodeTargetType, NSSubgraph
|
||||
from .pattern_utils import (
|
||||
end_node_matches_reversed_fusion,
|
||||
get_reversed_fusions,
|
||||
get_type_a_related_to_b,
|
||||
)
|
||||
|
||||
|
||||
toq = torch.ops.quantized
|
||||
|
||||
|
||||
def _get_output_nodes(g: Graph) -> list[Node]:
|
||||
return [n for n in g.nodes if n.op == "output"]
|
||||
|
||||
|
||||
class _NSGraphMatchableSubgraphsIterator:
|
||||
"""
|
||||
Iterates through the graph of gm, starting with the output nodes
|
||||
and continuing backwards.
|
||||
1. Returns matchable subgraphs, in order. A subgraph is defined by
|
||||
(start_node, end_node).
|
||||
2. Skips over non-matchable subgraphs
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
gm: GraphModule,
|
||||
non_matchable_functions: set[NSNodeTargetType],
|
||||
non_matchable_modules: set[NSNodeTargetType],
|
||||
non_matchable_methods: set[NSNodeTargetType],
|
||||
):
|
||||
self.gm: GraphModule = gm
|
||||
self.non_matchable_functions: set[NSNodeTargetType] = non_matchable_functions
|
||||
self.non_matchable_modules: set[NSNodeTargetType] = non_matchable_modules
|
||||
self.non_matchable_methods: set[NSNodeTargetType] = non_matchable_methods
|
||||
self.seen_nodes: set[Node] = set()
|
||||
self.stack: list[Node] = []
|
||||
for start_node in _get_output_nodes(self.gm.graph):
|
||||
self.stack.append(start_node)
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self) -> NSSubgraph:
|
||||
"""
|
||||
Returns the next matchable subgraph.
|
||||
"""
|
||||
while len(self.stack) > 0:
|
||||
cur_end_node = self.stack.pop()
|
||||
if cur_end_node in self.seen_nodes:
|
||||
continue
|
||||
|
||||
# for subgraphs which are single nodes, start_node == end_node
|
||||
# for subgraphs with more than one node, start node != end_node
|
||||
cur_start_node = cur_end_node
|
||||
# Subgraphs like linear-relu have the base node as the start node.
|
||||
# Subgraphs like dequantize-linear-relu-to(torch.float16) have the
|
||||
# base node as the second node.
|
||||
# The cur_base_op_node var will move to the actual node during
|
||||
# the fusion matching later in this code block.
|
||||
cur_base_op_node = cur_end_node
|
||||
|
||||
# Check for potential fusions. For now, we are greedy
|
||||
# and always skip all non-base nodes of a fusion. For example,
|
||||
# if we match linear-relu backwards, we will always skip the
|
||||
# relu node and attempt to match the linear node. This can
|
||||
# be made configurable later if needed.
|
||||
for _reverse_fusion_ops, base_op_idx in get_reversed_fusions():
|
||||
is_match = end_node_matches_reversed_fusion(
|
||||
cur_end_node, _reverse_fusion_ops, self.gm, self.seen_nodes
|
||||
)
|
||||
if is_match:
|
||||
# navigate to the base node
|
||||
# pyrefly: ignore [bad-assignment, non-convergent-recursion]
|
||||
for rev_fusion_idx in range(len(_reverse_fusion_ops) - 1):
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
self.seen_nodes.add(cur_start_node)
|
||||
# for now, assume that there are no other nodes
|
||||
# which need to be added to the stack
|
||||
cur_start_node = cur_start_node.args[0] # type: ignore[assignment]
|
||||
# if the base op index matches the current node, set it
|
||||
rev_base_op_idx = len(_reverse_fusion_ops) - 2 - base_op_idx
|
||||
if rev_fusion_idx == rev_base_op_idx:
|
||||
cur_base_op_node = cur_start_node
|
||||
break
|
||||
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
self.seen_nodes.add(cur_start_node)
|
||||
# add args of previous nodes to stack
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
for arg in cur_start_node.all_input_nodes:
|
||||
self._recursively_add_node_arg_to_stack(arg)
|
||||
|
||||
# skip unmatchable nodes
|
||||
# note: this check is done on the start_node, i.e.
|
||||
# if we are matching linear-relu in reverse, this would do the matchable
|
||||
# check on the linear
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
if not self._is_matchable(cur_base_op_node):
|
||||
continue
|
||||
|
||||
# If an observer or a fake_quant was not matched as a part of
|
||||
# a pattern of multiple nodes, ignore it. One case where this is
|
||||
# relevant is an observer on a graph input, which was added because
|
||||
# it is necessary for the next node.
|
||||
if cur_end_node.op == "call_module" and cur_start_node is cur_end_node:
|
||||
maybe_obs = getattr_from_fqn(self.gm, cur_end_node.target) # type: ignore[arg-type]
|
||||
if isinstance(maybe_obs, (ObserverBase, FakeQuantizeBase)):
|
||||
continue
|
||||
|
||||
return NSSubgraph(
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
start_node=cur_start_node,
|
||||
end_node=cur_end_node,
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
base_op_node=cur_base_op_node,
|
||||
)
|
||||
|
||||
raise StopIteration
|
||||
|
||||
def _recursively_add_node_arg_to_stack(self, arg: Any) -> None:
|
||||
"""
|
||||
Adds all of the nodes in this arg to the stack, properly navigating
|
||||
through list, dicts and tuples.
|
||||
"""
|
||||
if isinstance(arg, Node):
|
||||
self.stack.append(arg)
|
||||
elif (
|
||||
isinstance(arg, torch.fx.immutable_collections.immutable_list)
|
||||
or type(arg) is tuple
|
||||
):
|
||||
for inner_arg in arg:
|
||||
self._recursively_add_node_arg_to_stack(inner_arg)
|
||||
elif isinstance(arg, torch.fx.immutable_collections.immutable_dict):
|
||||
for value in arg.values():
|
||||
self._recursively_add_node_arg_to_stack(value)
|
||||
|
||||
def _is_matchable(self, node: Node) -> bool:
|
||||
if node.op == "call_function":
|
||||
return node.target not in self.non_matchable_functions
|
||||
elif node.op == "call_module":
|
||||
if not isinstance(node.target, str):
|
||||
raise AssertionError(f"Expected str, got {type(node.target)}")
|
||||
target_mod = getattr_from_fqn(self.gm, node.target)
|
||||
return not any(
|
||||
isinstance(target_mod, t) # type: ignore[arg-type]
|
||||
for t in self.non_matchable_modules
|
||||
)
|
||||
elif node.op == "call_method":
|
||||
return node.target not in self.non_matchable_methods
|
||||
else:
|
||||
return False
|
||||
|
||||
|
||||
class GraphMatchingException(Exception):
|
||||
"""
|
||||
Exception raised when two graphs cannot be matched.
|
||||
"""
|
||||
|
||||
|
||||
class SubgraphTypeRelationship(enum.Enum):
|
||||
# same type, known
|
||||
# example: F.linear and F.linear, or nn.Conv2d and nn.Conv2d
|
||||
EQUAL = enum.auto()
|
||||
# same type, but the type is not known to Numerical Suite
|
||||
# (user defined type, etc).
|
||||
EQUAL_BUT_UKNOWN = enum.auto()
|
||||
# known, same subgraph_relationship set, but not the same type
|
||||
# example: F.linear and toq.linear
|
||||
RELATED_BUT_NOT_EQUAL = enum.auto()
|
||||
# not related
|
||||
NOT_RELATED = enum.auto()
|
||||
|
||||
|
||||
def _get_subgraph_relationship_type(
|
||||
subgraph_a: NSSubgraph,
|
||||
subgraph_b: NSSubgraph,
|
||||
gm_a: GraphModule,
|
||||
gm_b: GraphModule,
|
||||
type_a_related_to_b: set[tuple[NSNodeTargetType, NSNodeTargetType]],
|
||||
) -> SubgraphTypeRelationship:
|
||||
node_a = subgraph_a.base_op_node
|
||||
node_b = subgraph_b.base_op_node
|
||||
|
||||
# TODO(next): make this code handle matching by what is before the base op
|
||||
if node_a.op != node_b.op:
|
||||
if not (
|
||||
node_a.op in ("call_function", "call_method")
|
||||
and node_b.op in ("call_function", "call_method")
|
||||
):
|
||||
return SubgraphTypeRelationship.NOT_RELATED
|
||||
|
||||
if node_a.op in ("call_function", "call_method"):
|
||||
key = (node_a.target, node_b.target)
|
||||
|
||||
if key not in type_a_related_to_b:
|
||||
if node_a.target == node_b.target:
|
||||
return SubgraphTypeRelationship.EQUAL_BUT_UKNOWN
|
||||
else:
|
||||
return SubgraphTypeRelationship.NOT_RELATED
|
||||
# after this point, we are dealing with known types
|
||||
|
||||
if node_a.target == node_b.target:
|
||||
node_a_has_prev = subgraph_a.base_op_node == subgraph_a.start_node
|
||||
node_b_has_prev = subgraph_b.base_op_node == subgraph_b.start_node
|
||||
if node_a_has_prev and (not node_b_has_prev):
|
||||
return SubgraphTypeRelationship.RELATED_BUT_NOT_EQUAL
|
||||
elif (not node_a_has_prev) and node_b_has_prev:
|
||||
return SubgraphTypeRelationship.RELATED_BUT_NOT_EQUAL
|
||||
elif (not node_a_has_prev) and (not node_b_has_prev):
|
||||
return SubgraphTypeRelationship.EQUAL
|
||||
else:
|
||||
# TODO(future PR): check for matches start_op_node and base_op_node
|
||||
return SubgraphTypeRelationship.EQUAL
|
||||
|
||||
if key in type_a_related_to_b:
|
||||
return SubgraphTypeRelationship.RELATED_BUT_NOT_EQUAL
|
||||
else:
|
||||
return SubgraphTypeRelationship.NOT_RELATED
|
||||
elif node_a.op == "call_module":
|
||||
if (
|
||||
subgraph_a.base_op_node != subgraph_a.start_node
|
||||
or subgraph_b.base_op_node != subgraph_b.start_node
|
||||
):
|
||||
raise AssertionError(
|
||||
"Matching call_module patterns where base_op_node != start_node is not supported yet"
|
||||
)
|
||||
# for call_module, we need to look up the modules to do the type check
|
||||
if not isinstance(node_a.target, str):
|
||||
raise AssertionError(f"Expected str, got {type(node_a.target)}")
|
||||
mod_a = getattr_from_fqn(gm_a, node_a.target)
|
||||
if not isinstance(node_b.target, str):
|
||||
raise AssertionError(f"Expected str, got {type(node_b.target)}")
|
||||
mod_b = getattr_from_fqn(gm_b, node_b.target)
|
||||
|
||||
key = (type(mod_a), type(mod_b))
|
||||
|
||||
if key not in type_a_related_to_b:
|
||||
if type(mod_a) is type(mod_b):
|
||||
return SubgraphTypeRelationship.EQUAL_BUT_UKNOWN
|
||||
else:
|
||||
return SubgraphTypeRelationship.NOT_RELATED
|
||||
elif type(mod_a) is type(mod_b):
|
||||
return SubgraphTypeRelationship.EQUAL
|
||||
else:
|
||||
return SubgraphTypeRelationship.RELATED_BUT_NOT_EQUAL
|
||||
|
||||
return SubgraphTypeRelationship.NOT_RELATED
|
||||
|
||||
|
||||
def _get_name_for_subgraph(
|
||||
subgraph_a: NSSubgraph,
|
||||
gm_a: GraphModule,
|
||||
base_name_to_sets_of_related_ops: dict[str, set[NSNodeTargetType]],
|
||||
existing_names: set[str],
|
||||
) -> str:
|
||||
"""
|
||||
Returns a unique name for a subgraph. This name is based on two things:
|
||||
1. the name of the set containing the underlying type of the base op in the
|
||||
subgraph (i.e. 'torch.nn.functional.linear' if this is related to a linear op)
|
||||
2. the number of previous subgraphs with related underlying type of the base op
|
||||
|
||||
For example, in the graph
|
||||
|
||||
linear0 -> relu0 -> linear1 -> relu1
|
||||
|
||||
The subgraphs are (linear0, relu0) and (linear1, relu1). If we iterate
|
||||
from the output node backwards, the name given to (linear1, relu1) will be
|
||||
`base_op_torch.nn.functional.linear_0`, and the name given to (linear0, relu0)
|
||||
will be `base_op_torch.nn.functional.linear_1`.
|
||||
|
||||
Why are we not just using the node name? Answer: because of two requirements:
|
||||
A. fusions must be supported
|
||||
B. some Numeric Suite APIs can be called without having all of the models in memory
|
||||
|
||||
For example, let's say we need to match nodes of
|
||||
|
||||
(1) ... -> linear0 -> relu0 -> ...
|
||||
|
||||
And
|
||||
|
||||
(2) ... -> linear_relu0 -> ...
|
||||
|
||||
Without being able to inspect them together. With the current naming scheme, if
|
||||
we iterate through both of these graphs in the same order, and assuming the rest
|
||||
of the graphs match, both of these subgraphs will get the same name without
|
||||
(1) and (2) knowing anything about each other.
|
||||
"""
|
||||
target_type = _get_node_target_type(subgraph_a.base_op_node, gm_a)
|
||||
target_base_type = None
|
||||
for base_name, sets_of_related_ops in base_name_to_sets_of_related_ops.items():
|
||||
if target_type in sets_of_related_ops:
|
||||
target_base_type = base_name
|
||||
target_base_name = "base_op_" + str(target_base_type)
|
||||
counter = 0
|
||||
proposed_name = target_base_name + "_" + str(counter)
|
||||
while proposed_name in existing_names:
|
||||
counter += 1
|
||||
proposed_name = target_base_name + "_" + str(counter)
|
||||
existing_names.add(proposed_name)
|
||||
return proposed_name
|
||||
|
||||
|
||||
def _get_node_target_type(node: Node, gm: GraphModule) -> NSNodeTargetType | None:
|
||||
if node.op in ("call_function", "call_method"):
|
||||
return node.target
|
||||
elif node.op == "call_module":
|
||||
if not isinstance(node.target, str):
|
||||
raise AssertionError(f"Expected str, got {type(node.target)}")
|
||||
mod = getattr_from_fqn(gm, node.target)
|
||||
return type(mod)
|
||||
return None
|
||||
|
||||
|
||||
def get_matching_subgraph_pairs(
|
||||
gm_a: GraphModule,
|
||||
gm_b: GraphModule,
|
||||
base_name_to_sets_of_related_ops: dict[str, set[NSNodeTargetType]] | None = None,
|
||||
unmatchable_types_map: dict[str, set[NSNodeTargetType]] | None = None,
|
||||
) -> dict[str, tuple[NSSubgraph, NSSubgraph]]:
|
||||
"""
|
||||
Matches matchable subgraphs of graph_a to graph_b.
|
||||
|
||||
For a node, "matchable" is defined as a node which is not an observer,
|
||||
fake_quants, quant or dequant.
|
||||
|
||||
A subgraph can contain one or more nodes. A subgraph is matchable if
|
||||
at least one node inside of it is matchable. Currently, all nodes in
|
||||
a subgraph must be matchable (because we assume no observers will be
|
||||
inserted in the middle of a fusion).
|
||||
|
||||
A subgraph is defined by (start_node, end_node). We assume that only
|
||||
start_node and end_node are linked with the surrounding graph, all other
|
||||
nodes in a subgraph are self-contained.
|
||||
|
||||
A pair of nodes is "related" if both nodes represent the same mathematical
|
||||
operation across different quantization flavors. For example,
|
||||
`F.linear` and `torch.ops.quantized.linear` are related, and
|
||||
`F.linear` and `torch.nn.Conv` are not related.
|
||||
|
||||
For each matchable pair of nodes node_a and node_b, they will match
|
||||
if node_a and node_b are related.
|
||||
|
||||
For graphs A and B, they will match iff:
|
||||
1. the number of matchable subgraphs in A and B is equivalent
|
||||
2. when iterating through the matchable subgraphs of A and B in the same order, each
|
||||
corresponding pair of base nodes is related.
|
||||
|
||||
This enables us to find the corresponding subgraphs between
|
||||
graphs of related models. For example, if we had two graphs such as:
|
||||
|
||||
graph_a: x0 -> conv_0 (type: nn.Conv2d) -> obs_0 -> x1
|
||||
w -/
|
||||
b -/
|
||||
|
||||
graph_b: x0 -> quant_0 -> qconv_0 (type: nnq.Conv2d) -> dequant_0 -> x1
|
||||
packed_params_0 -/
|
||||
|
||||
This function will return the following result:
|
||||
{
|
||||
'conv_0': ( # the name of the node in graph_b
|
||||
(conv_0, conv_0), # (start_node_a, end_node_a)
|
||||
(qconv_0, qconv_0), # (start_node_b, end_node_b)
|
||||
),
|
||||
}
|
||||
|
||||
Or, if we have a fusion pattern,
|
||||
|
||||
graph_a: x0 -> linear_0 -> relu_0 -> obs_0 -> x1
|
||||
w -/
|
||||
b -/
|
||||
|
||||
graph_b: x0 -> quant_0 -> linear_relu_0 -> dequant_0 -> x1
|
||||
packed_params_0 -/
|
||||
|
||||
This function will return the following result:
|
||||
{
|
||||
'linear_relu_0': ( # the name of the node in graph_b
|
||||
(linear_0, relu_0), # (start_node_a, end_node_a)
|
||||
(linear_relu_0, linear_relu_0), # (start_node_b, end_node_b)
|
||||
),
|
||||
}
|
||||
"""
|
||||
if unmatchable_types_map is None:
|
||||
unmatchable_types_map = get_unmatchable_types_map()
|
||||
non_matchable_functions = unmatchable_types_map["funs_unmatchable"]
|
||||
non_matchable_modules = unmatchable_types_map["mods_unmatchable"]
|
||||
non_matchable_methods = unmatchable_types_map["meths_unmatchable"]
|
||||
|
||||
graph_a_iterator = _NSGraphMatchableSubgraphsIterator(
|
||||
gm_a, non_matchable_functions, non_matchable_modules, non_matchable_methods
|
||||
)
|
||||
graph_b_iterator = _NSGraphMatchableSubgraphsIterator(
|
||||
gm_b, non_matchable_functions, non_matchable_modules, non_matchable_methods
|
||||
)
|
||||
results = collections.OrderedDict()
|
||||
if base_name_to_sets_of_related_ops is None:
|
||||
base_name_to_sets_of_related_ops = get_base_name_to_sets_of_related_ops()
|
||||
type_a_related_to_b = get_type_a_related_to_b(base_name_to_sets_of_related_ops)
|
||||
|
||||
existing_names_a: set[str] = set()
|
||||
existing_names_b: set[str] = set()
|
||||
|
||||
while True:
|
||||
# fetch the next subgraphs from a and b
|
||||
cur_subgraph_a, cur_subgraph_b = None, None
|
||||
try:
|
||||
cur_subgraph_a = next(graph_a_iterator)
|
||||
except StopIteration:
|
||||
pass
|
||||
try:
|
||||
cur_subgraph_b = next(graph_b_iterator)
|
||||
except StopIteration:
|
||||
pass
|
||||
|
||||
# look up types of a and b for useful error messages
|
||||
type_start_a, type_start_b = None, None
|
||||
if cur_subgraph_a is not None:
|
||||
type_start_a = _get_node_target_type(cur_subgraph_a.start_node, gm_a)
|
||||
if cur_subgraph_b is not None:
|
||||
type_start_b = _get_node_target_type(cur_subgraph_b.start_node, gm_b)
|
||||
|
||||
# check for results and determine what to do next
|
||||
if cur_subgraph_a is not None and cur_subgraph_b is not None:
|
||||
# both nodes were fetched, check for subgraph_relationship
|
||||
# note: subgraph_relationship is checked on the start node, i.e.
|
||||
# if a linear-relu pattern is checked, we would check for subgraph_relationship
|
||||
# of the linear
|
||||
subgraph_relationship = _get_subgraph_relationship_type(
|
||||
cur_subgraph_a, cur_subgraph_b, gm_a, gm_b, type_a_related_to_b
|
||||
)
|
||||
if subgraph_relationship == SubgraphTypeRelationship.NOT_RELATED:
|
||||
msg = f"""
|
||||
The subgraphs
|
||||
({cur_subgraph_a}, {type_start_a}) and
|
||||
({cur_subgraph_b}, {type_start_b})
|
||||
are not related. Please ensure that the two models you pass in have the same number
|
||||
of subgraphs, and each pair of subgraphs is related to each other."""
|
||||
raise GraphMatchingException(msg)
|
||||
elif subgraph_relationship == SubgraphTypeRelationship.EQUAL_BUT_UKNOWN:
|
||||
# skip matching but unknown types
|
||||
continue
|
||||
key_name_a = _get_name_for_subgraph(
|
||||
cur_subgraph_a, gm_a, base_name_to_sets_of_related_ops, existing_names_a
|
||||
)
|
||||
key_name_b = _get_name_for_subgraph(
|
||||
cur_subgraph_b, gm_b, base_name_to_sets_of_related_ops, existing_names_b
|
||||
)
|
||||
if key_name_a != key_name_b:
|
||||
raise AssertionError(
|
||||
f"Subgraph names {key_name_a} and {key_name_b} do not match"
|
||||
)
|
||||
results[key_name_a] = (cur_subgraph_a, cur_subgraph_b)
|
||||
continue
|
||||
elif cur_subgraph_a is None and cur_subgraph_b is None:
|
||||
# we reached the end of both graphs
|
||||
break
|
||||
else:
|
||||
# only one node was fetched, no match possible, throw error
|
||||
msg = f"""
|
||||
Attempting to match
|
||||
({cur_subgraph_a}, {type_start_a}) and
|
||||
({cur_subgraph_b}, {type_start_b}),
|
||||
one of which is empty. Please ensure that the two models you pass in have the same number
|
||||
of subgraphs."""
|
||||
raise GraphMatchingException(msg)
|
||||
|
||||
# The subgraph pairs are originally created by traversing the two graphs
|
||||
# from the outputs to the inputs. Reverse the results to return the
|
||||
# subgraphs in their order of execution.
|
||||
results = collections.OrderedDict(reversed(results.items()))
|
||||
|
||||
return results
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,763 @@
|
||||
import operator
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
import torch.ao.nn.intrinsic as nni
|
||||
import torch.ao.nn.intrinsic.qat as nniqat
|
||||
import torch.ao.nn.intrinsic.quantized as nniq
|
||||
import torch.ao.nn.intrinsic.quantized.dynamic as nniqd
|
||||
import torch.ao.nn.qat as nnqat
|
||||
import torch.ao.nn.qat.dynamic as nnqatd
|
||||
import torch.ao.nn.quantized as nnq
|
||||
import torch.ao.nn.quantized.dynamic as nnqd
|
||||
import torch.ao.quantization.fx._lower_to_native_backend as _lower_to_native_backend
|
||||
import torch.ao.quantization.quantization_mappings as quantization_mappings
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.ao.quantization.backend_config import get_native_backend_config
|
||||
|
||||
from .ns_types import NSNodeTargetType
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
|
||||
toq = torch.ops.quantized
|
||||
|
||||
|
||||
def get_base_name_to_sets_of_related_ops() -> dict[str, set[NSNodeTargetType]]:
|
||||
# note: this set is modified below by items from backend_config
|
||||
sets_of_related_ops: list[set[NSNodeTargetType]] = [
|
||||
# conv modules
|
||||
{
|
||||
nn.Conv1d,
|
||||
},
|
||||
{
|
||||
nn.Conv2d,
|
||||
},
|
||||
{
|
||||
nn.Conv3d,
|
||||
},
|
||||
# conv functionals
|
||||
{
|
||||
F.conv1d,
|
||||
},
|
||||
{
|
||||
F.conv2d,
|
||||
},
|
||||
{
|
||||
F.conv3d,
|
||||
},
|
||||
# linear modules
|
||||
{
|
||||
nn.Linear,
|
||||
},
|
||||
# linear functionals
|
||||
{
|
||||
F.linear,
|
||||
},
|
||||
# average pool
|
||||
{
|
||||
nn.AvgPool1d,
|
||||
torch.avg_pool1d,
|
||||
},
|
||||
{
|
||||
nn.AvgPool2d,
|
||||
torch._C._nn.avg_pool2d,
|
||||
},
|
||||
{
|
||||
nn.AvgPool3d,
|
||||
torch._C._nn.avg_pool3d,
|
||||
},
|
||||
# adaptive average pool
|
||||
{
|
||||
nn.AdaptiveAvgPool1d,
|
||||
F.adaptive_avg_pool1d,
|
||||
},
|
||||
{
|
||||
nn.AdaptiveAvgPool2d,
|
||||
F.adaptive_avg_pool2d,
|
||||
},
|
||||
{
|
||||
nn.AdaptiveAvgPool3d,
|
||||
F.adaptive_avg_pool3d,
|
||||
},
|
||||
# LSTM
|
||||
{
|
||||
nn.LSTM,
|
||||
},
|
||||
# add
|
||||
{
|
||||
torch.add,
|
||||
operator.add, # x + y
|
||||
},
|
||||
# cat
|
||||
{
|
||||
torch.cat,
|
||||
},
|
||||
# mul
|
||||
{
|
||||
torch.mul,
|
||||
operator.mul,
|
||||
},
|
||||
# relu
|
||||
{
|
||||
F.relu,
|
||||
nn.ReLU,
|
||||
"relu",
|
||||
"relu_",
|
||||
torch.relu,
|
||||
},
|
||||
# maxpool
|
||||
{
|
||||
nn.MaxPool1d,
|
||||
F.max_pool1d,
|
||||
},
|
||||
{
|
||||
nn.MaxPool2d,
|
||||
F.max_pool2d,
|
||||
},
|
||||
{
|
||||
nn.MaxPool3d,
|
||||
F.max_pool3d,
|
||||
},
|
||||
# sigmoid
|
||||
{
|
||||
torch.sigmoid,
|
||||
"sigmoid",
|
||||
"sigmoid_",
|
||||
nn.Sigmoid,
|
||||
F.sigmoid,
|
||||
},
|
||||
# BatchNorm
|
||||
{
|
||||
nn.BatchNorm2d,
|
||||
},
|
||||
{
|
||||
nn.BatchNorm3d,
|
||||
},
|
||||
# ConvTranspose
|
||||
{
|
||||
nn.ConvTranspose1d,
|
||||
},
|
||||
{
|
||||
nn.ConvTranspose2d,
|
||||
},
|
||||
{
|
||||
nn.ConvTranspose3d,
|
||||
},
|
||||
# functional transposed conv
|
||||
{
|
||||
F.conv_transpose1d,
|
||||
},
|
||||
{
|
||||
F.conv_transpose2d,
|
||||
},
|
||||
{
|
||||
F.conv_transpose3d,
|
||||
},
|
||||
# ELU
|
||||
{
|
||||
nn.ELU,
|
||||
},
|
||||
# Embedding
|
||||
{
|
||||
nn.Embedding,
|
||||
},
|
||||
# EmbeddingBag
|
||||
{
|
||||
nn.EmbeddingBag,
|
||||
},
|
||||
# GroupNorm
|
||||
{
|
||||
nn.GroupNorm,
|
||||
},
|
||||
# Hardswish
|
||||
{
|
||||
nn.Hardswish,
|
||||
},
|
||||
# InstanceNorm
|
||||
{
|
||||
nn.InstanceNorm1d,
|
||||
},
|
||||
{
|
||||
nn.InstanceNorm2d,
|
||||
},
|
||||
{
|
||||
nn.InstanceNorm3d,
|
||||
},
|
||||
# LayerNorm
|
||||
{
|
||||
nn.LayerNorm,
|
||||
},
|
||||
# LeakyReLU
|
||||
{
|
||||
nn.LeakyReLU,
|
||||
},
|
||||
# ReLU6
|
||||
{
|
||||
nn.ReLU6,
|
||||
F.relu6,
|
||||
},
|
||||
# F.elu
|
||||
{
|
||||
F.elu,
|
||||
},
|
||||
# F.hardswish
|
||||
{
|
||||
F.hardswish,
|
||||
},
|
||||
# F.group_norm
|
||||
{
|
||||
F.group_norm,
|
||||
},
|
||||
# F.instance_norm
|
||||
{
|
||||
F.instance_norm,
|
||||
},
|
||||
# F.layer_norm
|
||||
{
|
||||
F.layer_norm,
|
||||
},
|
||||
# F.leaky_relu
|
||||
{
|
||||
F.leaky_relu,
|
||||
},
|
||||
# F.silu
|
||||
{
|
||||
nn.SiLU,
|
||||
F.silu,
|
||||
},
|
||||
# F.mish
|
||||
{
|
||||
nn.Mish,
|
||||
F.mish,
|
||||
},
|
||||
# F.tanh
|
||||
{
|
||||
nn.Tanh,
|
||||
F.tanh,
|
||||
torch.tanh,
|
||||
"tanh_",
|
||||
"tanh",
|
||||
},
|
||||
# F.hardsigmoid
|
||||
{
|
||||
"hardsigmoid_",
|
||||
"hardsigmoid",
|
||||
F.hardsigmoid,
|
||||
nn.Hardsigmoid,
|
||||
},
|
||||
# F.hardtanh
|
||||
{
|
||||
nn.Hardtanh,
|
||||
F.hardtanh,
|
||||
F.hardtanh_,
|
||||
},
|
||||
# floordiv
|
||||
{
|
||||
operator.floordiv,
|
||||
},
|
||||
# unsqueeze
|
||||
{
|
||||
torch.unsqueeze,
|
||||
},
|
||||
# stack
|
||||
{
|
||||
torch.stack,
|
||||
},
|
||||
# squeeze
|
||||
{
|
||||
torch.squeeze,
|
||||
},
|
||||
# sort
|
||||
{
|
||||
torch.sort,
|
||||
},
|
||||
# repeat_interleave
|
||||
{
|
||||
torch.repeat_interleave,
|
||||
},
|
||||
# min
|
||||
{
|
||||
torch.min,
|
||||
},
|
||||
# mean
|
||||
{
|
||||
torch.mean,
|
||||
},
|
||||
# max
|
||||
{
|
||||
torch.max,
|
||||
},
|
||||
# transpose
|
||||
{
|
||||
torch.transpose,
|
||||
},
|
||||
# flatten
|
||||
{
|
||||
torch.flatten,
|
||||
},
|
||||
# clamp
|
||||
{
|
||||
torch.clamp,
|
||||
},
|
||||
# chunk
|
||||
{
|
||||
torch.chunk,
|
||||
},
|
||||
# interpolate
|
||||
{
|
||||
torch.nn.functional.interpolate,
|
||||
},
|
||||
# dropout
|
||||
{
|
||||
nn.Dropout,
|
||||
},
|
||||
# F.dropout
|
||||
{
|
||||
F.dropout,
|
||||
},
|
||||
# matmul
|
||||
{
|
||||
torch.matmul,
|
||||
},
|
||||
# Softmax
|
||||
{
|
||||
nn.Softmax,
|
||||
},
|
||||
# PReLU
|
||||
{
|
||||
nn.PReLU,
|
||||
nnq.PReLU,
|
||||
},
|
||||
# F.prelu
|
||||
{
|
||||
F.prelu,
|
||||
toq.prelu,
|
||||
},
|
||||
# pixel shuffle
|
||||
{
|
||||
nn.PixelShuffle,
|
||||
},
|
||||
{
|
||||
F.pixel_shuffle,
|
||||
},
|
||||
# pixel unshuffle
|
||||
{
|
||||
nn.PixelUnshuffle,
|
||||
},
|
||||
{
|
||||
F.pixel_unshuffle,
|
||||
},
|
||||
# narrow
|
||||
{
|
||||
torch.narrow,
|
||||
},
|
||||
]
|
||||
|
||||
# for each floating point op, add versions of the op added by
|
||||
# backend_config
|
||||
backend_config = get_native_backend_config()
|
||||
|
||||
new_connections: list[tuple[Callable, Callable]] = [
|
||||
# technical debt edge case
|
||||
(nn.Linear, nn.modules.linear.NonDynamicallyQuantizableLinear),
|
||||
]
|
||||
|
||||
for pattern, config in backend_config._pattern_complex_format_to_config.items():
|
||||
# pattern format: (c, (b, a))
|
||||
first_element = pattern
|
||||
# look from the end, because pattern is in reverse order
|
||||
while isinstance(first_element, (list, tuple)):
|
||||
first_element = first_element[-1]
|
||||
|
||||
if config.fused_module is not None:
|
||||
# case 1: pattern fuses a pattern of ops into an op
|
||||
# example: nn.Conv1d, nn.ReLU fused into nni.ConvReLU1d
|
||||
new_connections.append((first_element, config.fused_module))
|
||||
|
||||
if config.qat_module is not None:
|
||||
# case 2: pattern swaps a module into a QAT module
|
||||
# example: nni.ConvReLU1d swapped into nniqat.ConvReLU1d
|
||||
new_connections.append((first_element, config.qat_module))
|
||||
|
||||
if config.reference_quantized_module is not None:
|
||||
# case 3: reference version of floating point module, such as
|
||||
# nn.Conv2d and nnqr.Conv2d
|
||||
new_connections.append((first_element, config.reference_quantized_module))
|
||||
|
||||
#
|
||||
# Add reference module swaps from default lowering path
|
||||
#
|
||||
|
||||
for source_to_target in (
|
||||
_lower_to_native_backend.STATIC_LOWER_MODULE_MAP,
|
||||
_lower_to_native_backend.DYNAMIC_LOWER_MODULE_MAP,
|
||||
_lower_to_native_backend.WEIGHT_ONLY_LOWER_MODULE_MAP,
|
||||
_lower_to_native_backend.SPECIAL_PATTERN_LOWER_MODULE_MAP,
|
||||
):
|
||||
for source, target in source_to_target.items(): # type: ignore[attr-defined]
|
||||
new_connections.append((source, target))
|
||||
|
||||
for source_to_double_target in (
|
||||
_lower_to_native_backend.STATIC_LOWER_FUSED_MODULE_MAP,
|
||||
_lower_to_native_backend.STATIC_LOWER_FUSED_MODULE_TWO_INPUTS_MAP,
|
||||
_lower_to_native_backend.DYNAMIC_LOWER_FUSED_MODULE_MAP,
|
||||
):
|
||||
for source, (target1, target2) in source_to_double_target.items(): # type: ignore[attr-defined]
|
||||
new_connections.append((source, target1))
|
||||
new_connections.append((source, target2))
|
||||
|
||||
#
|
||||
# Add function swaps from default lowering path
|
||||
#
|
||||
|
||||
for source, ( # type:ignore[assignment]
|
||||
target1,
|
||||
target2,
|
||||
) in _lower_to_native_backend.STATIC_LOWER_FUNCTIONAL_MAP.items():
|
||||
new_connections.append((source, target1))
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
new_connections.append((source, target2))
|
||||
|
||||
for source_to_target in (
|
||||
_lower_to_native_backend.QBIN_OP_MAPPING,
|
||||
_lower_to_native_backend.QBIN_RELU_OP_MAPPING,
|
||||
quantization_mappings.DEFAULT_FLOAT_TO_QUANTIZED_OPERATOR_MAPPINGS,
|
||||
):
|
||||
for source, target in source_to_target.items(): # type:ignore[assignment]
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
new_connections.append((source, target))
|
||||
|
||||
#
|
||||
# Add other swaps, ideally in the future this could be removed
|
||||
# after the lowering code stops using these.
|
||||
#
|
||||
for source_to_target in (
|
||||
quantization_mappings.DEFAULT_DYNAMIC_QUANT_MODULE_MAPPINGS,
|
||||
):
|
||||
for source, target in source_to_target.items(): # type:ignore[assignment]
|
||||
new_connections.append((source, target))
|
||||
|
||||
# add the new connections from backend_config
|
||||
for item1, item2 in new_connections:
|
||||
for set_of_related_ops in sets_of_related_ops:
|
||||
if item1 in set_of_related_ops or item2 in set_of_related_ops:
|
||||
set_of_related_ops.add(item1)
|
||||
set_of_related_ops.add(item2)
|
||||
break
|
||||
|
||||
base_name_to_sets_of_related_ops: dict[str, set[NSNodeTargetType]] = {}
|
||||
|
||||
for counter, set_of_related_ops in enumerate(sets_of_related_ops):
|
||||
base_name = str(counter)
|
||||
base_name_to_sets_of_related_ops[base_name] = set_of_related_ops
|
||||
|
||||
return base_name_to_sets_of_related_ops
|
||||
|
||||
|
||||
def get_base_name_for_op(
|
||||
base_name_to_sets_of_related_ops: dict[str, set[NSNodeTargetType]],
|
||||
op: NSNodeTargetType,
|
||||
) -> str | None:
|
||||
for base_name, set_of_related_ops in base_name_to_sets_of_related_ops.items():
|
||||
if op in set_of_related_ops:
|
||||
return base_name
|
||||
return None
|
||||
|
||||
|
||||
def add_op_to_sets_of_related_ops(
|
||||
base_name_to_sets_of_related_ops: dict[str, set[NSNodeTargetType]],
|
||||
op: NSNodeTargetType,
|
||||
related_op: NSNodeTargetType | None,
|
||||
) -> None:
|
||||
if related_op is not None:
|
||||
for set_of_related_ops in base_name_to_sets_of_related_ops.values():
|
||||
if related_op in set_of_related_ops:
|
||||
set_of_related_ops.add(op)
|
||||
return
|
||||
# if we got here, related_op was not found
|
||||
raise AssertionError(f"{related_op} was not found")
|
||||
else:
|
||||
counter = 0
|
||||
while str(counter) in base_name_to_sets_of_related_ops:
|
||||
counter += 1
|
||||
base_name_to_sets_of_related_ops[str(counter)] = {op}
|
||||
|
||||
|
||||
# TODO(future PR): clean this up
|
||||
def get_node_type_to_io_type_map() -> dict[str, set[NSNodeTargetType]]:
|
||||
FUNS_IO_TYPE_FP32: set[NSNodeTargetType] = {
|
||||
F.linear,
|
||||
F.conv1d,
|
||||
F.conv2d,
|
||||
F.conv3d,
|
||||
torch.cat,
|
||||
F.elu,
|
||||
F.hardswish,
|
||||
F.instance_norm,
|
||||
F.layer_norm,
|
||||
F.leaky_relu,
|
||||
F.dropout,
|
||||
F.silu,
|
||||
F.mish,
|
||||
operator.add,
|
||||
torch.add,
|
||||
operator.mul,
|
||||
torch.mul,
|
||||
torch.sum,
|
||||
F.prelu,
|
||||
}
|
||||
|
||||
FUNS_IO_TYPE_FP16: set[NSNodeTargetType] = set()
|
||||
|
||||
FUNS_IO_TYPE_INT8: set[NSNodeTargetType] = {
|
||||
toq.linear,
|
||||
toq.linear_relu,
|
||||
toq.conv1d,
|
||||
toq.conv1d_relu,
|
||||
toq.conv2d,
|
||||
toq.conv2d_relu,
|
||||
toq.conv3d,
|
||||
toq.conv3d_relu,
|
||||
toq.cat,
|
||||
toq.elu,
|
||||
toq.hardswish,
|
||||
toq.instance_norm,
|
||||
toq.layer_norm,
|
||||
toq.leaky_relu,
|
||||
toq.dropout,
|
||||
toq.prelu,
|
||||
# TODO(future PR): implement shadowing for binary ops and
|
||||
# uncomment below
|
||||
# toq.add,
|
||||
# toq.mul,
|
||||
}
|
||||
|
||||
FUNS_IO_TYPE_FP32_OR_INT8: set[NSNodeTargetType] = {
|
||||
F.relu,
|
||||
F.tanh,
|
||||
torch.tanh,
|
||||
F.sigmoid,
|
||||
torch.sigmoid,
|
||||
F.hardsigmoid,
|
||||
operator.floordiv,
|
||||
torch.adaptive_avg_pool1d,
|
||||
F.adaptive_avg_pool2d,
|
||||
F.adaptive_avg_pool3d,
|
||||
F.dropout,
|
||||
F.hardtanh,
|
||||
F.hardtanh_,
|
||||
F.interpolate,
|
||||
F.max_pool1d,
|
||||
F.max_pool2d,
|
||||
F.max_pool3d,
|
||||
F.relu6,
|
||||
F.pixel_shuffle,
|
||||
F.pixel_unshuffle,
|
||||
torch.avg_pool1d,
|
||||
torch._C._nn.avg_pool2d,
|
||||
torch._C._nn.avg_pool3d,
|
||||
torch.cat,
|
||||
torch.chunk,
|
||||
torch.clamp,
|
||||
torch.flatten,
|
||||
torch.transpose,
|
||||
torch.max,
|
||||
torch.mean,
|
||||
torch.min,
|
||||
torch.narrow,
|
||||
torch.repeat_interleave,
|
||||
torch.sort,
|
||||
torch.squeeze,
|
||||
torch.stack,
|
||||
torch.unsqueeze,
|
||||
operator.add,
|
||||
}
|
||||
|
||||
MODS_IO_TYPE_FP32: set[NSNodeTargetType] = {
|
||||
nn.Linear,
|
||||
nnqat.Linear,
|
||||
nnqatd.Linear,
|
||||
nnqd.Linear,
|
||||
torch.nn.modules.linear.NonDynamicallyQuantizableLinear,
|
||||
nn.Conv1d,
|
||||
nn.Conv2d,
|
||||
nn.Conv3d,
|
||||
nnqat.Conv1d,
|
||||
nnqat.Conv2d,
|
||||
nnqat.Conv3d,
|
||||
nnqat.Embedding,
|
||||
nnqat.EmbeddingBag,
|
||||
nn.LSTM,
|
||||
# note: nnqd.Linear is an instance of nnq.Linear, so this
|
||||
# check has to happen before the int8 module check
|
||||
nnqd.LSTM,
|
||||
nn.BatchNorm2d,
|
||||
nn.BatchNorm3d,
|
||||
nn.Dropout,
|
||||
nn.ConvTranspose1d,
|
||||
nn.ConvTranspose2d,
|
||||
nn.ConvTranspose3d,
|
||||
nn.ELU,
|
||||
nn.GroupNorm,
|
||||
nn.InstanceNorm1d,
|
||||
nn.InstanceNorm2d,
|
||||
nn.InstanceNorm3d,
|
||||
nn.LayerNorm,
|
||||
nn.Hardswish,
|
||||
nn.LeakyReLU,
|
||||
nn.ReLU6,
|
||||
nn.SiLU,
|
||||
nn.Mish,
|
||||
nn.Softmax,
|
||||
nn.PReLU,
|
||||
nni.BNReLU2d,
|
||||
nni.BNReLU3d,
|
||||
nni.ConvReLU1d,
|
||||
nni.ConvReLU2d,
|
||||
nni.ConvReLU3d,
|
||||
nni.LinearReLU,
|
||||
nni.LinearBn1d,
|
||||
nni.ConvBn1d,
|
||||
nni.ConvBn2d,
|
||||
nni.ConvBn3d,
|
||||
nniqat.ConvBn1d,
|
||||
nniqat.ConvBn2d,
|
||||
nniqat.ConvBn3d,
|
||||
nniqat.ConvBnReLU1d,
|
||||
nniqat.ConvBnReLU2d,
|
||||
nniqat.ConvBnReLU3d,
|
||||
nniqat.ConvReLU1d,
|
||||
nniqat.ConvReLU2d,
|
||||
nniqat.ConvReLU3d,
|
||||
nniqat.LinearReLU,
|
||||
nniqat.LinearBn1d,
|
||||
nniqd.LinearReLU,
|
||||
nni.LinearLeakyReLU,
|
||||
nni.LinearTanh,
|
||||
nni.ConvAdd2d,
|
||||
nni.ConvAddReLU2d,
|
||||
}
|
||||
|
||||
MODS_IO_TYPE_INT8: set[NSNodeTargetType] = {
|
||||
nnq.Linear,
|
||||
nnq.Conv1d,
|
||||
nnq.Conv2d,
|
||||
nnq.Conv3d,
|
||||
nnq.BatchNorm2d,
|
||||
nnq.BatchNorm3d,
|
||||
nnq.Dropout,
|
||||
nnq.ConvTranspose1d,
|
||||
nnq.ConvTranspose2d,
|
||||
nnq.ELU,
|
||||
nnq.InstanceNorm1d,
|
||||
nnq.InstanceNorm2d,
|
||||
nnq.InstanceNorm3d,
|
||||
nnq.LayerNorm,
|
||||
nnq.Hardswish,
|
||||
nnq.LeakyReLU,
|
||||
nnq.Embedding,
|
||||
nnq.EmbeddingBag,
|
||||
nnq.Dropout,
|
||||
nnq.Softmax,
|
||||
nnq.PReLU,
|
||||
nniq.BNReLU2d,
|
||||
nniq.BNReLU3d,
|
||||
nniq.ConvReLU1d,
|
||||
nniq.ConvReLU2d,
|
||||
nniq.ConvReLU3d,
|
||||
nniq.LinearReLU,
|
||||
nniq.LinearLeakyReLU,
|
||||
nniq.LinearTanh,
|
||||
nniq.ConvAdd2d,
|
||||
nniq.ConvAddReLU2d,
|
||||
}
|
||||
|
||||
MODS_IO_TYPE_FP32_OR_INT8: set[NSNodeTargetType] = {
|
||||
nn.ReLU,
|
||||
nn.Tanh,
|
||||
nn.Sigmoid,
|
||||
nn.Hardsigmoid,
|
||||
nn.AdaptiveAvgPool1d,
|
||||
nn.AdaptiveAvgPool2d,
|
||||
nn.AdaptiveAvgPool3d,
|
||||
nn.AvgPool1d,
|
||||
nn.AvgPool2d,
|
||||
nn.AvgPool3d,
|
||||
nn.Dropout,
|
||||
nn.Hardtanh,
|
||||
nn.Identity,
|
||||
nn.MaxPool1d,
|
||||
nn.MaxPool2d,
|
||||
nn.MaxPool3d,
|
||||
nn.PixelShuffle,
|
||||
nn.PixelUnshuffle,
|
||||
nn.ReLU6,
|
||||
}
|
||||
|
||||
METHS_IO_TYPE_FP32_OR_INT8: set[NSNodeTargetType] = {
|
||||
"sigmoid_",
|
||||
"sigmoid",
|
||||
"tanh_",
|
||||
"tanh",
|
||||
"hardsigmoid_",
|
||||
"hardsigmoid",
|
||||
"relu_",
|
||||
"relu",
|
||||
}
|
||||
|
||||
return {
|
||||
"funs_io_type_fp32": FUNS_IO_TYPE_FP32,
|
||||
"funs_io_type_fp16": FUNS_IO_TYPE_FP16,
|
||||
"funs_io_type_int8": FUNS_IO_TYPE_INT8,
|
||||
"funs_io_type_fp32_or_int8": FUNS_IO_TYPE_FP32_OR_INT8,
|
||||
"mods_io_type_fp32": MODS_IO_TYPE_FP32,
|
||||
"mods_io_type_int8": MODS_IO_TYPE_INT8,
|
||||
"mods_io_type_fp32_or_int8": MODS_IO_TYPE_FP32_OR_INT8,
|
||||
"meths_io_type_fp32_or_int8": METHS_IO_TYPE_FP32_OR_INT8,
|
||||
}
|
||||
|
||||
|
||||
def get_unmatchable_types_map() -> dict[str, set[NSNodeTargetType]]:
|
||||
FUNS_UNMATCHABLE: set[NSNodeTargetType] = {
|
||||
torch.quantize_per_tensor,
|
||||
operator.getitem,
|
||||
}
|
||||
|
||||
MODS_UNMATCHABLE: set[NSNodeTargetType] = {
|
||||
nn.Identity,
|
||||
}
|
||||
|
||||
METHS_UNMATCHABLE: set[NSNodeTargetType] = {
|
||||
"to",
|
||||
"dequantize",
|
||||
"reshape",
|
||||
"view",
|
||||
"unsqueeze_",
|
||||
"unsqueeze",
|
||||
"transpose",
|
||||
"squeeze_",
|
||||
"squeeze",
|
||||
"size",
|
||||
"shape",
|
||||
"resize_",
|
||||
"repeat_interleave",
|
||||
"repeat",
|
||||
"permute",
|
||||
"numel",
|
||||
"mean",
|
||||
"detach_",
|
||||
"detach",
|
||||
"contiguous",
|
||||
"clamp",
|
||||
"chunk",
|
||||
}
|
||||
|
||||
return {
|
||||
"funs_unmatchable": FUNS_UNMATCHABLE,
|
||||
"mods_unmatchable": MODS_UNMATCHABLE,
|
||||
"meths_unmatchable": METHS_UNMATCHABLE,
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,66 @@
|
||||
import enum
|
||||
from collections.abc import Callable
|
||||
from typing import Any, NamedTuple
|
||||
|
||||
from torch.fx.graph import Node
|
||||
|
||||
|
||||
class NSSingleResultValuesType(str, enum.Enum):
|
||||
WEIGHT = "weight"
|
||||
NODE_OUTPUT = "node_output"
|
||||
NODE_INPUT = "node_input"
|
||||
|
||||
|
||||
class NSSubgraph(NamedTuple):
|
||||
start_node: Node
|
||||
end_node: Node
|
||||
base_op_node: Node
|
||||
|
||||
|
||||
# TODO(future PR): see if we can use typing_extensions's TypedDict instead
|
||||
# to properly type the various keys
|
||||
# {
|
||||
# # one of NSSingleResultValuesType
|
||||
# 'type': 'weight',
|
||||
# # the values of type specified above
|
||||
# 'values': [torch.tensor(...), ...],
|
||||
# # name of the node directly before the logger
|
||||
# 'prev_node_name': 'linear1',
|
||||
# # type of the underlying function or module
|
||||
# 'prev_node_target_type': torch.nn.functional.linear # or torch.nn.Linear, etc
|
||||
# # name of the node responsible for adding this logger
|
||||
# # Note: this may differ from prev_node_name if we are logging inputs
|
||||
# 'ref_node_name': 'linear1',
|
||||
# # index of this node within the arg of the input/output node
|
||||
# # for example, in cat([x1, x2, x3], dim=0), x2 would have index_within_arg == 1
|
||||
# 'index_within_arg': 0,
|
||||
# # index of this node within the args of the input/output node
|
||||
# # for example, in add(x1, x2), x2 would have index_of_arg == 1
|
||||
# 'index_of_arg': 0,
|
||||
# # precomputed comparisons of logger values to reference values
|
||||
# 'comparisons': [torch.tensor(...), ...]
|
||||
# # name of function used for precomputed comparisons
|
||||
# 'comparison_fn_name': 'sqnr',
|
||||
# # string representation of qconfig responsible for creating this logger
|
||||
# 'qconfig_str': 'QConfig(...)',
|
||||
# }
|
||||
NSSingleResultType = dict[str, Any]
|
||||
|
||||
# {
|
||||
# 'layer_name_1': { # subgraph name
|
||||
# 'node_output': { # results type (node_output, node_input, weight)
|
||||
# 'model_name_a': # model name
|
||||
# [NSSingleResultType, ...], # results, ordered by index_within_arg
|
||||
# 'model_name_b':
|
||||
# [NSSingleResultType, ...],
|
||||
# },
|
||||
# },
|
||||
# }
|
||||
#
|
||||
NSResultsType = dict[str, dict[str, dict[str, list[NSSingleResultType]]]]
|
||||
|
||||
# Defines the underlying target type of a node, for example:
|
||||
# `F.conv1d` for a `call_function` conv node
|
||||
# `nn.Conv1d` for a `call_module` node calling the forward of a `nn.Conv1d` module
|
||||
# `'sigmoid'` for a `call_method` node calling `x.sigmoid()`
|
||||
NSNodeTargetType = Callable | str
|
||||
@@ -0,0 +1,215 @@
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.ao.quantization import FakeQuantizeBase, ObserverBase
|
||||
from torch.ao.quantization.backend_config import get_native_backend_config
|
||||
from torch.ao.quantization.fx.quantize_handler import _get_pattern_to_quantize_handlers
|
||||
from torch.ao.quantization.utils import getattr_from_fqn
|
||||
from torch.fx import GraphModule
|
||||
from torch.fx.graph import Node
|
||||
|
||||
from .ns_types import NSNodeTargetType
|
||||
|
||||
|
||||
toq = torch.ops.quantized
|
||||
|
||||
|
||||
def get_type_a_related_to_b(
|
||||
base_name_to_sets_of_related_ops: dict[str, set[NSNodeTargetType]],
|
||||
) -> set[tuple[NSNodeTargetType, NSNodeTargetType]]:
|
||||
# TODO(future PR): allow customizations
|
||||
# TODO(future PR): reuse existing quantization mappings
|
||||
# TODO(future PR): add the rest of modules and ops here
|
||||
type_a_related_to_b: set[tuple[NSNodeTargetType, NSNodeTargetType]] = set()
|
||||
|
||||
for s in base_name_to_sets_of_related_ops.values():
|
||||
s_list = list(s)
|
||||
# add every bidirectional pair
|
||||
for idx_0 in range(len(s_list)):
|
||||
for idx_1 in range(idx_0, len(s_list)):
|
||||
type_a_related_to_b.add((s_list[idx_0], s_list[idx_1]))
|
||||
type_a_related_to_b.add((s_list[idx_1], s_list[idx_0]))
|
||||
|
||||
return type_a_related_to_b
|
||||
|
||||
|
||||
NSFusionElType = (
|
||||
Callable # call_function or call_module type, example: F.linear or nn.Conv2d
|
||||
| str # call_method name, example: "dequantize"
|
||||
| tuple[
|
||||
str, Any
|
||||
] # call_method name and first argument, example: ("to", torch.float16)
|
||||
)
|
||||
NSFusionType = (
|
||||
tuple[NSFusionElType, NSFusionElType]
|
||||
| tuple[NSFusionElType, NSFusionElType, NSFusionElType, NSFusionElType]
|
||||
)
|
||||
|
||||
|
||||
def get_reversed_fusions() -> list[tuple[NSFusionType, int]]:
|
||||
"""
|
||||
Set of potential fusions, in reverse order. The order is reversed
|
||||
to match how fusion patterns are defined in quantization code.
|
||||
|
||||
Fusion format:
|
||||
((fusion_op_0, fusion_op_1), base_op_idx)
|
||||
|
||||
Where base_op_idx is the idx of the op we should use to match other related
|
||||
ops. Note: base_op_idx is specified in non-reverse order, i.e. a base_op_idx
|
||||
of 0 represents the first op in regular (non-reverse) order, 1 represents the
|
||||
second op, etc.
|
||||
"""
|
||||
results: list[tuple[NSFusionType, int]] = []
|
||||
|
||||
# Possible syntaxes:
|
||||
# * single op: torch.nn.Conv2d
|
||||
# * multiple ops: (torch.nn.ReLU, torch.nn.Conv2d)
|
||||
# For fusions, we only care about patterns composed of multiple ops.
|
||||
# TODO(future PR): allow customizations from default patterns.
|
||||
all_quant_patterns = _get_pattern_to_quantize_handlers(get_native_backend_config())
|
||||
|
||||
default_base_op_idx = 0
|
||||
for quant_pattern in all_quant_patterns:
|
||||
# TODO: this is a temporary hack to flatten the patterns from quantization so
|
||||
# that it works with the ns matcher function, maybe we should use `_is_match`
|
||||
# in torch.ao.quantization.fx.match_utils to match the patterns
|
||||
if (
|
||||
isinstance(quant_pattern, tuple)
|
||||
and len(quant_pattern) == 2
|
||||
and isinstance(quant_pattern[1], tuple)
|
||||
and len(quant_pattern[1]) == 2
|
||||
):
|
||||
# flatten the pattern with form (nn.ReLU, (nn.BatchNorm2d, nn.Conv2d))
|
||||
quant_pattern = (quant_pattern[0], quant_pattern[1][0], quant_pattern[1][1])
|
||||
|
||||
# Only patterns of multiple ops are fusions, ignore
|
||||
# patterns which contain a single ops (they get matched
|
||||
# without caring about fusions).
|
||||
if isinstance(quant_pattern, tuple):
|
||||
results.append((quant_pattern, default_base_op_idx)) # type: ignore[arg-type]
|
||||
|
||||
# For each pattern, add additional patterns with observers and
|
||||
# fake quants at the end.
|
||||
# TODO(future PR): if needed, implement matching for a node
|
||||
# having multiple output observers.
|
||||
for cls in (ObserverBase, FakeQuantizeBase):
|
||||
if isinstance(quant_pattern, tuple):
|
||||
# pyrefly: ignore [not-iterable]
|
||||
new_pattern = (cls, *quant_pattern)
|
||||
else:
|
||||
new_pattern = (cls, quant_pattern)
|
||||
results.append((new_pattern, default_base_op_idx)) # type: ignore[arg-type]
|
||||
|
||||
# After this point, results contains values such as
|
||||
# [..., ((torch.nn.Relu, torch.nn.Conv2d), 0), ...]
|
||||
|
||||
# Patterns for matching fp16 emulation are not specified in the quantization
|
||||
# fusion mappings. For now, define them here.
|
||||
fp16_em_base_op_idx = 1
|
||||
patterns_to_add = [
|
||||
# linear-relu fp16 emulation:
|
||||
# fp16_to_fp32 -> linear -> relu -> fp32_to_fp16
|
||||
(
|
||||
(("to", torch.float16), F.relu, F.linear, "dequantize"),
|
||||
fp16_em_base_op_idx,
|
||||
),
|
||||
# Conv-BN fusion (this happens outside of quantization patterns,
|
||||
# which is why it is defined separately here).
|
||||
((nn.BatchNorm1d, nn.Conv1d), default_base_op_idx),
|
||||
((nn.BatchNorm2d, nn.Conv2d), default_base_op_idx),
|
||||
((nn.BatchNorm3d, nn.Conv3d), default_base_op_idx),
|
||||
((nn.ReLU, nn.BatchNorm1d, nn.Conv1d), default_base_op_idx),
|
||||
((nn.ReLU, nn.BatchNorm2d, nn.Conv2d), default_base_op_idx),
|
||||
((nn.ReLU, nn.BatchNorm3d, nn.Conv3d), default_base_op_idx),
|
||||
]
|
||||
for p in patterns_to_add:
|
||||
results.append(p) # type: ignore[arg-type]
|
||||
results.append(((ObserverBase, *p[0]), p[1])) # type: ignore[arg-type]
|
||||
results.append(((FakeQuantizeBase, *p[0]), p[1])) # type: ignore[arg-type]
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def end_node_matches_reversed_fusion(
|
||||
end_node: Node,
|
||||
reversed_fusion: NSFusionType,
|
||||
gm: GraphModule,
|
||||
seen_nodes: set[Node],
|
||||
) -> bool:
|
||||
"""
|
||||
Returns true if a pattern ending with `end_node` matches
|
||||
the fusion pattern.
|
||||
"""
|
||||
cur_node = end_node
|
||||
for fusion_idx in range(len(reversed_fusion)):
|
||||
# each node can only belong to one matched pattern
|
||||
if cur_node in seen_nodes:
|
||||
return False
|
||||
|
||||
cur_fusion_el = reversed_fusion[fusion_idx]
|
||||
|
||||
if cur_node.op == "call_function":
|
||||
fusion_el_is_fun = (not isinstance(cur_fusion_el, str)) and (
|
||||
not isinstance(cur_fusion_el, type)
|
||||
)
|
||||
if fusion_el_is_fun:
|
||||
if cur_node.target != cur_fusion_el:
|
||||
return False
|
||||
if len(cur_node.args) > 0 and isinstance(cur_node.args[0], Node):
|
||||
cur_node = cur_node.args[0]
|
||||
else:
|
||||
return False
|
||||
else:
|
||||
return False
|
||||
|
||||
elif cur_node.op == "call_module":
|
||||
fusion_el_is_mod = isinstance(cur_fusion_el, type)
|
||||
if fusion_el_is_mod:
|
||||
if not isinstance(cur_node.target, str):
|
||||
raise AssertionError(f"Expected str, got {type(cur_node.target)}")
|
||||
target_mod = getattr_from_fqn(gm, cur_node.target)
|
||||
if not isinstance(cur_fusion_el, type):
|
||||
return False
|
||||
if not isinstance(target_mod, cur_fusion_el):
|
||||
return False
|
||||
if len(cur_node.args) > 0 and isinstance(cur_node.args[0], Node):
|
||||
cur_node = cur_node.args[0]
|
||||
else:
|
||||
return False
|
||||
else:
|
||||
return False
|
||||
|
||||
elif cur_node.op == "call_method":
|
||||
fusion_el_is_meth_with_second_arg = (
|
||||
isinstance(cur_fusion_el, tuple) and len(cur_fusion_el) == 2
|
||||
)
|
||||
fusion_el_is_meth_without_args = isinstance(cur_fusion_el, str)
|
||||
if fusion_el_is_meth_without_args or fusion_el_is_meth_with_second_arg:
|
||||
if fusion_el_is_meth_without_args:
|
||||
if cur_node.target != cur_fusion_el:
|
||||
return False
|
||||
else:
|
||||
if not isinstance(cur_fusion_el, tuple):
|
||||
raise AssertionError(
|
||||
f"Expected tuple, got {type(cur_fusion_el)}"
|
||||
)
|
||||
if cur_node.target != cur_fusion_el[0]:
|
||||
return False
|
||||
elif len(cur_node.args) < 2:
|
||||
return False
|
||||
elif cur_node.args[1] != cur_fusion_el[1]:
|
||||
return False
|
||||
|
||||
if len(cur_node.args) > 0 and isinstance(cur_node.args[0], Node):
|
||||
cur_node = cur_node.args[0]
|
||||
else:
|
||||
return False
|
||||
else:
|
||||
return False
|
||||
else:
|
||||
return False
|
||||
|
||||
return True
|
||||
@@ -0,0 +1,251 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
from torch.ao.quantization import QConfigMapping
|
||||
from torch.ao.quantization.qconfig_mapping import _QCONFIG_STYLE_ORDER
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
from torch.ao.quantization.qconfig import QConfigAny
|
||||
|
||||
__all__ = ["QConfigMultiMapping"]
|
||||
|
||||
_QCONFIG_STYLE_TO_METHOD: dict[str, str] = {
|
||||
"global_qconfig": "set_global",
|
||||
"object_type_qconfigs": "set_object_type",
|
||||
"module_name_regex_qconfigs": "set_module_name_regex",
|
||||
"module_name_qconfigs": "set_module_name",
|
||||
"module_name_object_type_order_qconfigs": "set_module_name_object_type_order",
|
||||
}
|
||||
|
||||
|
||||
def _remove_duplicates_and_none(qconfig_list: list[QConfigAny]) -> None:
|
||||
to_remove = []
|
||||
for index, cur_qconfig in enumerate(qconfig_list):
|
||||
if cur_qconfig is None:
|
||||
to_remove.append(index)
|
||||
break
|
||||
for checked_qconfig in qconfig_list[:index]:
|
||||
if torch.ao.quantization.qconfig_equals(cur_qconfig, checked_qconfig):
|
||||
to_remove.append(index)
|
||||
break
|
||||
for index in to_remove[::-1]:
|
||||
qconfig_list.pop(index)
|
||||
|
||||
|
||||
class QConfigMultiMapping:
|
||||
"""
|
||||
This class, used with the prepare_n_shadows_model API, stores a list of :class:`torch.ao.quantization.QConfigMapping`s
|
||||
so that multiple QConfigs can be specified for each QConfig matching style.
|
||||
|
||||
The user can specify QConfigs using the following methods (in increasing match priority):
|
||||
|
||||
``set_global`` : sets the global (default) QConfigs
|
||||
|
||||
``set_object_type`` : sets the QConfigs for a given module type, function, or method name
|
||||
|
||||
``set_module_name_regex`` : sets the QConfigs for modules matching the given regex string
|
||||
|
||||
``set_module_name`` : sets the QConfigs for modules matching the given module name
|
||||
|
||||
``set_module_name_object_type_order`` : sets the QConfigs for modules matching a combination
|
||||
of the given module name, object type, and the index at which the module appears
|
||||
|
||||
Note: Usage of set methods is the same as in QConfigMapping except with a passed in list of QConfigs rather than a
|
||||
single QConfig.
|
||||
|
||||
Example usage::
|
||||
|
||||
qconfig_mapping = QConfigMultiMapping()
|
||||
.set_global([qconfig1, qconfig2])
|
||||
.set_object_type(torch.nn.Linear, [qconfig2, qconfig3])
|
||||
.set_object_type(torch.nn.ReLU, [qconfig1])
|
||||
.set_module_name_regex("foo.*bar.*conv[0-9]+", [qconfig2])
|
||||
.set_module_name_regex("foo.*", [qconfig1, qconfig2, qconfig3])
|
||||
.set_module_name("module1", [None])
|
||||
.set_module_name("module2", [qconfig2])
|
||||
.set_module_name_object_type_order("foo.bar", torch.nn.functional.linear, 0, [qconfig3])
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
# initialize this with 1 QConfigMapping to avoid corner cases
|
||||
self.qconfig_mappings_list: list[QConfigMapping] = [QConfigMapping()]
|
||||
|
||||
def _handle_list_size_mismatch(
|
||||
self, qconfig_list: list[QConfigAny], style: str
|
||||
) -> None:
|
||||
# this method handles cases where the size of qconfig_list does not match
|
||||
# the size of qconfig_mappings_list.
|
||||
# Issue: Consider a user inserting global_qconfig A and B first, then inserting
|
||||
# qconfig C as an object_type_qconfig for conv ops. If we internally store
|
||||
# 1 QConfigMapping with A and C and another with just B, then the
|
||||
# second QConfigMapping will match B to conv ops (which is not wanted), since B is global.
|
||||
|
||||
# we avoid this by maintaining the invariant that if any QConfigMapping
|
||||
# has a qconfig style+key with a qconfig in it, all QConfigMappings must
|
||||
# have either a qconfig or None for that same style+key. In the above
|
||||
# example, a None qconfig would prevent the unwanted match in the
|
||||
# second QConfigMapping
|
||||
|
||||
if len(qconfig_list) > len(self.qconfig_mappings_list):
|
||||
# Case: we have more qconfigs (in qconfig_list) than QConfigMappings
|
||||
|
||||
# Add new QConfigMappings (initialized so we maintain the `invariant`)
|
||||
|
||||
new_qconfig_mapping = QConfigMapping()
|
||||
# searches other QConfigMappings for qconfig style+keys
|
||||
# that need to be inserted as `None` into the new QConfigMapping
|
||||
for qconfig_mapping in self.qconfig_mappings_list:
|
||||
# global_qconfig has None by default
|
||||
for check_style in _QCONFIG_STYLE_ORDER[1:]:
|
||||
qconfigs_dict = getattr(qconfig_mapping, check_style)
|
||||
target_qconfigs_dict = getattr(new_qconfig_mapping, check_style)
|
||||
for key in qconfigs_dict:
|
||||
target_qconfigs_dict[key] = None
|
||||
break
|
||||
|
||||
# insert copies of this new QConfigMapping until all entries
|
||||
# in qconfig_list can fit among the QConfigMappings
|
||||
while len(qconfig_list) > len(self.qconfig_mappings_list):
|
||||
self.qconfig_mappings_list.append(copy.deepcopy(new_qconfig_mapping))
|
||||
else:
|
||||
# Case: we have fewer qconfigs in qconfig_list than QConfigMappings
|
||||
|
||||
# pad qconfig_list with `None` until length is same
|
||||
while len(qconfig_list) < len(self.qconfig_mappings_list):
|
||||
qconfig_list.append(None)
|
||||
|
||||
# this function applies the insertion method across each QConfigMapping
|
||||
def _insert_qconfig_list(
|
||||
self,
|
||||
style: str,
|
||||
args: list[str | int | Callable],
|
||||
qconfig_list: list[QConfigAny],
|
||||
) -> None:
|
||||
# we remove duplicates and None to make the ordering of qconfigs
|
||||
# deterministic upon insertion.
|
||||
_remove_duplicates_and_none(qconfig_list)
|
||||
|
||||
self._handle_list_size_mismatch(qconfig_list, style)
|
||||
method_name = _QCONFIG_STYLE_TO_METHOD[style]
|
||||
for qconfig_mapping, qconfig in zip(self.qconfig_mappings_list, qconfig_list):
|
||||
# uses QConfigMapping set method to insert qconfig
|
||||
set_method = getattr(qconfig_mapping, method_name)
|
||||
set_method(*args, qconfig)
|
||||
|
||||
def set_global(self, global_qconfig_list: list[QConfigAny]) -> QConfigMultiMapping:
|
||||
"""
|
||||
Set global QConfigs
|
||||
see :func:`~torch.ao.quantization.QConfigMapping.set_global()` for more info
|
||||
"""
|
||||
self._insert_qconfig_list("global_qconfig", [], global_qconfig_list)
|
||||
return self
|
||||
|
||||
def set_object_type(
|
||||
self, object_type: Callable | str, qconfig_list: list[QConfigAny]
|
||||
) -> QConfigMultiMapping:
|
||||
"""
|
||||
Set object type QConfigs
|
||||
see :func:`~torch.ao.quantization.QConfigMapping.set_object_type()` for more info
|
||||
"""
|
||||
self._insert_qconfig_list("object_type_qconfigs", [object_type], qconfig_list)
|
||||
return self
|
||||
|
||||
def set_module_name_regex(
|
||||
self, module_name_regex: str, qconfig_list: list[QConfigAny]
|
||||
) -> QConfigMultiMapping:
|
||||
"""
|
||||
Set module_name_regex QConfigs
|
||||
see :func:`~torch.ao.quantization.QConfigMapping.set_module_name_regex()` for more info
|
||||
"""
|
||||
self._insert_qconfig_list(
|
||||
"module_name_regex_qconfigs", [module_name_regex], qconfig_list
|
||||
)
|
||||
return self
|
||||
|
||||
def set_module_name(
|
||||
self, module_name: str, qconfig_list: list[QConfigAny]
|
||||
) -> QConfigMultiMapping:
|
||||
"""
|
||||
Set module_name QConfigs
|
||||
see :func:`~torch.ao.quantization.QConfigMapping.set_module_name()` for more info
|
||||
"""
|
||||
self._insert_qconfig_list("module_name_qconfigs", [module_name], qconfig_list)
|
||||
return self
|
||||
|
||||
def set_module_name_object_type_order(
|
||||
self,
|
||||
module_name: str,
|
||||
object_type: Callable,
|
||||
index: int,
|
||||
qconfig_list: list[QConfigAny],
|
||||
) -> QConfigMultiMapping:
|
||||
"""
|
||||
Set module_name QConfigs
|
||||
see :func:`~torch.ao.quantization.QConfigMapping.set_module_name_object_type_order()` for more info
|
||||
"""
|
||||
self._insert_qconfig_list(
|
||||
"module_name_object_type_order_qconfigs",
|
||||
[module_name, object_type, index],
|
||||
qconfig_list,
|
||||
)
|
||||
return self
|
||||
|
||||
def __repr__(self):
|
||||
return (
|
||||
self.__class__.__name__
|
||||
+ " ["
|
||||
+ "".join(
|
||||
f"\n{qconfig_mapping.__repr__()},"
|
||||
for qconfig_mapping in self.qconfig_mappings_list
|
||||
)
|
||||
+ "\n]"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_list_qconfig_mapping(
|
||||
cls, qconfig_mapping_list: list[QConfigMapping]
|
||||
) -> QConfigMultiMapping:
|
||||
"""
|
||||
Creates a QConfigMultiMapping from a list of QConfigMappings
|
||||
"""
|
||||
new_qconfig_multi_mapping = cls()
|
||||
|
||||
new_qconfig_multi_mapping.qconfig_mappings_list = copy.deepcopy(
|
||||
qconfig_mapping_list
|
||||
)
|
||||
|
||||
# we need to avoid the issue described in _handle_list_size_mismatch,
|
||||
# so we reinsert all the qconfigs using the QConfigMultiMapping
|
||||
# set methods
|
||||
|
||||
# go through all qconfig styles
|
||||
# note: global can be ignored since it is None by default
|
||||
for style in _QCONFIG_STYLE_ORDER[1:]:
|
||||
# gather all key+qconfigs for current style
|
||||
# into qconfig_dict_list
|
||||
qconfig_dict_list: dict[Any, list[QConfigAny]] = {}
|
||||
for qconfig_mapping in qconfig_mapping_list:
|
||||
qconfig_dict = getattr(qconfig_mapping, style)
|
||||
for key, qconfig in qconfig_dict.items():
|
||||
if key not in qconfig_dict_list:
|
||||
qconfig_dict_list[key] = []
|
||||
qconfig_dict_list[key].append(qconfig)
|
||||
|
||||
# reinsert all gathered key+qconfigs
|
||||
set_method_name = _QCONFIG_STYLE_TO_METHOD[style]
|
||||
set_method = getattr(new_qconfig_multi_mapping, set_method_name)
|
||||
for key, qconfig_list in qconfig_dict_list.items():
|
||||
if isinstance(key, tuple):
|
||||
set_method(*key, qconfig_list)
|
||||
else:
|
||||
set_method(key, qconfig_list)
|
||||
|
||||
return new_qconfig_multi_mapping
|
||||
@@ -0,0 +1,579 @@
|
||||
# mypy: allow-untyped-decorators
|
||||
# mypy: allow-untyped-defs
|
||||
import enum
|
||||
import operator
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
import torch.ao.nn.intrinsic.quantized as nniq
|
||||
import torch.ao.nn.quantized as nnq
|
||||
import torch.nn as nn
|
||||
from torch.ao.quantization import FakeQuantizeBase, ObserverBase
|
||||
from torch.ao.quantization.observer import _is_activation_post_process
|
||||
from torch.ao.quantization.utils import getattr_from_fqn
|
||||
from torch.fx import GraphModule
|
||||
from torch.fx.graph import Node
|
||||
|
||||
from .ns_types import NSNodeTargetType, NSResultsType
|
||||
|
||||
|
||||
toq = torch.ops.quantized
|
||||
|
||||
|
||||
# TODO(future PR): consider deleting this enum and using the torch types
|
||||
# directly. This might be tricky because it is not a one to one mapping.
|
||||
class NodeInputOrOutputType(enum.Enum):
|
||||
FP32 = enum.auto() # torch.float
|
||||
INT8 = enum.auto() # torch.qint8 or torch.quint8
|
||||
FP16 = enum.auto() # torch.float16
|
||||
UNKNOWN = enum.auto() # we cannot determine input/output dtype
|
||||
# TODO(future PR): while these functions can support multiple dtypes,
|
||||
# for the purposes of numerical debugging we want to get the actual
|
||||
# dtype used in the model. We will likely need some kind of dtype
|
||||
# propagation to estimate this.
|
||||
FP32_OR_INT8 = enum.auto() # either torch.float or torch.quint8 or torch.qint8
|
||||
# TODO(future PRs): dynamic quant, fake quant, etc
|
||||
|
||||
|
||||
def get_node_first_input_and_output_type(
|
||||
node: Node,
|
||||
gm: GraphModule,
|
||||
logger_cls: Callable,
|
||||
node_type_to_io_type_map: dict[str, set[NSNodeTargetType]],
|
||||
) -> tuple[NodeInputOrOutputType, NodeInputOrOutputType]:
|
||||
# TODO(future PR): clean this up
|
||||
FUNS_IO_TYPE_FP32 = node_type_to_io_type_map["funs_io_type_fp32"]
|
||||
FUNS_IO_TYPE_FP16 = node_type_to_io_type_map["funs_io_type_fp16"]
|
||||
FUNS_IO_TYPE_INT8 = node_type_to_io_type_map["funs_io_type_int8"]
|
||||
FUNS_IO_TYPE_FP32_OR_INT8 = node_type_to_io_type_map["funs_io_type_fp32_or_int8"]
|
||||
MODS_IO_TYPE_FP32 = node_type_to_io_type_map["mods_io_type_fp32"]
|
||||
MODS_IO_TYPE_INT8 = node_type_to_io_type_map["mods_io_type_int8"]
|
||||
MODS_IO_TYPE_FP32_OR_INT8 = node_type_to_io_type_map["mods_io_type_fp32_or_int8"]
|
||||
METHS_IO_TYPE_FP32_OR_INT8 = node_type_to_io_type_map["meths_io_type_fp32_or_int8"]
|
||||
|
||||
if node.op == "call_function":
|
||||
if node.target in FUNS_IO_TYPE_FP32:
|
||||
return (NodeInputOrOutputType.FP32, NodeInputOrOutputType.FP32)
|
||||
if node.target in FUNS_IO_TYPE_FP16:
|
||||
return (NodeInputOrOutputType.FP16, NodeInputOrOutputType.FP16)
|
||||
elif node.target in FUNS_IO_TYPE_INT8:
|
||||
return (NodeInputOrOutputType.INT8, NodeInputOrOutputType.INT8)
|
||||
elif node.target in FUNS_IO_TYPE_FP32_OR_INT8:
|
||||
first_arg = get_normalized_nth_input(node, gm, 0)
|
||||
if not isinstance(first_arg, Node):
|
||||
raise AssertionError(f"Expected Node, got {type(first_arg)}")
|
||||
(
|
||||
_prev_node_input_type,
|
||||
prev_node_output_type,
|
||||
) = get_node_first_input_and_output_type(
|
||||
first_arg, gm, logger_cls, node_type_to_io_type_map
|
||||
)
|
||||
return (prev_node_output_type, prev_node_output_type)
|
||||
else:
|
||||
return (NodeInputOrOutputType.UNKNOWN, NodeInputOrOutputType.UNKNOWN)
|
||||
|
||||
elif node.op == "call_module":
|
||||
if node.op != "call_module":
|
||||
raise AssertionError(f"Expected call_module, got '{node.op}'")
|
||||
if not isinstance(node.target, str):
|
||||
raise AssertionError(f"Expected str, but got {type(node.target)}")
|
||||
|
||||
mod = getattr_from_fqn(gm, node.target)
|
||||
is_known_fp32_or_int8_input_module = any(
|
||||
isinstance(mod, target_type) # type: ignore[arg-type]
|
||||
for target_type in MODS_IO_TYPE_FP32_OR_INT8
|
||||
)
|
||||
if (
|
||||
isinstance(mod, (logger_cls, ObserverBase, FakeQuantizeBase)) # type: ignore[arg-type]
|
||||
or is_known_fp32_or_int8_input_module
|
||||
):
|
||||
# A logger or observer's input and output type is the output
|
||||
# type of the preceding node.
|
||||
first_arg = get_normalized_nth_input(node, gm, 0)
|
||||
if not isinstance(first_arg, Node):
|
||||
raise AssertionError(f"Expected Node, got {type(first_arg)}")
|
||||
(
|
||||
_prev_node_input_type,
|
||||
prev_node_output_type,
|
||||
) = get_node_first_input_and_output_type(
|
||||
first_arg, gm, logger_cls, node_type_to_io_type_map
|
||||
)
|
||||
return (prev_node_output_type, prev_node_output_type)
|
||||
is_known_fp32_input_module = any(
|
||||
isinstance(mod, target_type) # type: ignore[arg-type]
|
||||
for target_type in MODS_IO_TYPE_FP32
|
||||
)
|
||||
is_known_int8_input_module = any(
|
||||
isinstance(mod, target_type) # type: ignore[arg-type]
|
||||
for target_type in MODS_IO_TYPE_INT8
|
||||
)
|
||||
if is_known_fp32_input_module:
|
||||
return (NodeInputOrOutputType.FP32, NodeInputOrOutputType.FP32)
|
||||
elif is_known_int8_input_module:
|
||||
return (NodeInputOrOutputType.INT8, NodeInputOrOutputType.INT8)
|
||||
else:
|
||||
return (NodeInputOrOutputType.UNKNOWN, NodeInputOrOutputType.UNKNOWN)
|
||||
|
||||
elif node.op == "call_method":
|
||||
if node.target == "dequantize":
|
||||
# Dequantize is a special node because it allows multiple input types.
|
||||
# So, we look up the output type of the previous node and return that
|
||||
# as the input type of this node instance.
|
||||
prev_node = get_normalized_nth_input(node, gm, 0)
|
||||
if not isinstance(prev_node, Node):
|
||||
raise AssertionError(f"Expected Node, got {type(prev_node)}")
|
||||
(
|
||||
_prev_node_input_type,
|
||||
prev_node_output_type,
|
||||
) = get_node_first_input_and_output_type(
|
||||
prev_node, gm, logger_cls, node_type_to_io_type_map
|
||||
)
|
||||
return (prev_node_output_type, NodeInputOrOutputType.FP32)
|
||||
|
||||
elif node.target == "to":
|
||||
# to is a special node because it allows multiple input types.
|
||||
# So, we look up the output type of the previous node and return that
|
||||
# as the input type of this node instance. We also look up the target
|
||||
# of to and return the correct output type.
|
||||
prev_node = get_normalized_nth_input(node, gm, 0)
|
||||
if not isinstance(prev_node, Node):
|
||||
raise AssertionError(f"Expected Node, got {type(prev_node)}")
|
||||
(
|
||||
_prev_node_input_type,
|
||||
prev_node_output_type,
|
||||
) = get_node_first_input_and_output_type(
|
||||
prev_node, gm, logger_cls, node_type_to_io_type_map
|
||||
)
|
||||
|
||||
cur_node_dtype_target = get_normalized_nth_input(node, gm, 1)
|
||||
if cur_node_dtype_target is not torch.float16:
|
||||
raise AssertionError(
|
||||
f"{cur_node_dtype_target} handling needs to be added"
|
||||
)
|
||||
|
||||
return (prev_node_output_type, NodeInputOrOutputType.FP16)
|
||||
|
||||
elif node.target in METHS_IO_TYPE_FP32_OR_INT8:
|
||||
first_arg = get_normalized_nth_input(node, gm, 0)
|
||||
if not isinstance(first_arg, Node):
|
||||
raise AssertionError(f"Expected Node, got {type(first_arg)}")
|
||||
(
|
||||
_prev_node_input_type,
|
||||
prev_node_output_type,
|
||||
) = get_node_first_input_and_output_type(
|
||||
first_arg, gm, logger_cls, node_type_to_io_type_map
|
||||
)
|
||||
return (prev_node_output_type, prev_node_output_type)
|
||||
|
||||
return (NodeInputOrOutputType.UNKNOWN, NodeInputOrOutputType.UNKNOWN)
|
||||
else:
|
||||
return (NodeInputOrOutputType.UNKNOWN, NodeInputOrOutputType.UNKNOWN)
|
||||
|
||||
|
||||
def get_node_input_qparams(
|
||||
node: Node,
|
||||
gm: GraphModule,
|
||||
node_type_to_io_type_map: dict[str, set[NSNodeTargetType]],
|
||||
) -> tuple[torch.Tensor | float, torch.Tensor | int] | None:
|
||||
"""
|
||||
Returns the qparams (scale, zero_point) of the first input to `node`,
|
||||
if they can be inferred from the graph.
|
||||
"""
|
||||
prev_node = get_normalized_nth_input(node, gm, 0)
|
||||
|
||||
if not isinstance(prev_node, Node):
|
||||
return None
|
||||
|
||||
MODS_IO_TYPE_FP32_OR_INT8 = node_type_to_io_type_map["mods_io_type_fp32_or_int8"]
|
||||
|
||||
def _get_scale_zp_from_function_args(node, gm, scale_arg_idx, zp_arg_idx):
|
||||
scale_node = get_normalized_nth_input(node, gm, scale_arg_idx)
|
||||
zp_node = get_normalized_nth_input(node, gm, zp_arg_idx)
|
||||
if not isinstance(scale_node, Node):
|
||||
raise AssertionError(f"Expected Node, got {type(scale_node)}")
|
||||
if not isinstance(scale_node.target, str):
|
||||
raise AssertionError(f"Expected str, got {type(scale_node.target)}")
|
||||
if not isinstance(zp_node, Node):
|
||||
raise AssertionError(f"Expected Node, got {type(zp_node)}")
|
||||
if not isinstance(zp_node.target, str):
|
||||
raise AssertionError(f"Expected str, got {type(zp_node.target)}")
|
||||
scale_obj = getattr_from_fqn(gm, scale_node.target)
|
||||
zp_obj = getattr_from_fqn(gm, zp_node.target)
|
||||
return (scale_obj, zp_obj)
|
||||
|
||||
if prev_node.op == "call_function":
|
||||
# quantize - read the args directly
|
||||
if prev_node.target is torch.quantize_per_tensor:
|
||||
return _get_scale_zp_from_function_args(prev_node, gm, 1, 2)
|
||||
elif prev_node.target in (toq.add, toq.add_relu, toq.mul, toq.mul_relu):
|
||||
return _get_scale_zp_from_function_args(prev_node, gm, 2, 3)
|
||||
|
||||
return None
|
||||
# TODO(future PR): handle more functionals
|
||||
# TODO(future PR): handle functional ops which inherit qparams from input
|
||||
|
||||
elif prev_node.op == "call_module":
|
||||
# get type of the module
|
||||
if not isinstance(prev_node.target, str):
|
||||
raise AssertionError(f"Expected str, got {type(prev_node.target)}")
|
||||
module_obj = getattr_from_fqn(gm, prev_node.target)
|
||||
if isinstance(
|
||||
module_obj,
|
||||
(
|
||||
nnq.Linear,
|
||||
nnq.Conv1d,
|
||||
nnq.Conv2d,
|
||||
nniq.ConvReLU2d,
|
||||
nnq.Conv3d,
|
||||
nnq.BatchNorm2d,
|
||||
nnq.BatchNorm3d,
|
||||
nnq.ConvTranspose1d,
|
||||
nnq.ConvTranspose2d,
|
||||
nnq.ELU,
|
||||
nnq.GroupNorm,
|
||||
nnq.InstanceNorm1d,
|
||||
nnq.InstanceNorm2d,
|
||||
nnq.InstanceNorm3d,
|
||||
nnq.LayerNorm,
|
||||
nnq.Hardswish,
|
||||
nnq.LeakyReLU,
|
||||
nnq.ReLU6,
|
||||
nniq.BNReLU2d,
|
||||
nniq.BNReLU3d,
|
||||
nniq.ConvReLU1d,
|
||||
nniq.ConvReLU2d,
|
||||
nniq.ConvReLU3d,
|
||||
nniq.LinearReLU,
|
||||
),
|
||||
):
|
||||
return (module_obj.scale, module_obj.zero_point) # type: ignore[return-value]
|
||||
|
||||
is_known_fp32_or_int8_input_module = any(
|
||||
isinstance(module_obj, target_type) # type: ignore[arg-type]
|
||||
for target_type in MODS_IO_TYPE_FP32_OR_INT8
|
||||
)
|
||||
if is_known_fp32_or_int8_input_module:
|
||||
return get_node_input_qparams(prev_node, gm, node_type_to_io_type_map)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def return_first_non_observer_node(
|
||||
node: Node,
|
||||
gm: GraphModule,
|
||||
) -> Node:
|
||||
"""
|
||||
If node is not an observer, returns it. If node is an observer,
|
||||
navigates up the graph and returns the first parent which is not an
|
||||
observer. For example,
|
||||
|
||||
graph: (node_non_obs), node = node_non_obs : returns node_non_obs
|
||||
graph: (node_non_obs -> obs0), node = obs0 : returns node_non_obs
|
||||
graph: (node_non_obs -> obs0 -> fq0), node = fq0 : returns node_non_obs
|
||||
"""
|
||||
if node.op == "call_module":
|
||||
node_obj = getattr_from_fqn(gm, node.target) # type: ignore[arg-type]
|
||||
if _is_activation_post_process(node_obj):
|
||||
if len(node.args) != 1:
|
||||
raise AssertionError(
|
||||
f"Expected node.args to have length 1, got {len(node.args)}"
|
||||
)
|
||||
if not isinstance(node.args[0], Node):
|
||||
raise AssertionError(f"Expected Node, got {type(node.args[0])}")
|
||||
node = node.args[0]
|
||||
# code duplication intended, not worth refactoring
|
||||
if not isinstance(node.target, str):
|
||||
raise AssertionError(f"Expected str, got {type(node.target)}")
|
||||
node_obj = getattr_from_fqn(gm, node.target)
|
||||
if _is_activation_post_process(node_obj):
|
||||
if len(node.args) != 1:
|
||||
raise AssertionError(
|
||||
f"Expected node.args to have length 1, got {len(node.args)}"
|
||||
)
|
||||
if not isinstance(node.args[0], Node):
|
||||
raise AssertionError(f"Expected Node, got {type(node.args[0])}")
|
||||
node = node.args[0]
|
||||
return node
|
||||
|
||||
|
||||
def get_number_of_non_param_args(
|
||||
node: Node,
|
||||
gm: GraphModule,
|
||||
) -> int:
|
||||
"""
|
||||
Assumes that all non-param args occur first. Returns the number of
|
||||
non-param args expected for a node. For example, for
|
||||
|
||||
F.linear(x, weight, bias)
|
||||
|
||||
Returns 1, because x is a non-param arg and weight and bias are params.
|
||||
For
|
||||
|
||||
lstm_mod(x, hid)
|
||||
|
||||
Returns 2, because both x and hid are non-param args.
|
||||
"""
|
||||
if node.op == "call_module":
|
||||
node_obj = getattr_from_fqn(gm, node.target) # type: ignore[arg-type]
|
||||
if isinstance(node_obj, nn.LSTM):
|
||||
return 2
|
||||
|
||||
# default is 1
|
||||
return 1
|
||||
|
||||
|
||||
def get_arg_indices_of_inputs_to_log(node: Node) -> list[int]:
|
||||
"""
|
||||
Returns the indices of args of the node which we should attach
|
||||
loggers to, if input logging is enabled.
|
||||
|
||||
For example,
|
||||
* for (x + y), returns [0, 1]
|
||||
* for (1 + y), returns [1]
|
||||
* for (x + 1), returns [0]
|
||||
* for (linear(x, w, b)) returns [0]
|
||||
* by default, returns [0]
|
||||
"""
|
||||
if len(node.args) == 0:
|
||||
return []
|
||||
if node.op == "call_function" and (
|
||||
# TODO(future PR): use relationship map instead of hardcoding
|
||||
node.target in (torch.add, torch.ops.quantized.add, operator.add)
|
||||
or node.target in (torch.mul, torch.ops.quantized.mul, operator.mul)
|
||||
):
|
||||
result = [i for i in range(2) if type(node.args[i]) is Node]
|
||||
return result
|
||||
return [0]
|
||||
|
||||
|
||||
def get_target_type_str(node: Node, gm: GraphModule) -> str:
|
||||
"""
|
||||
Returns a string representation of the type of the function or module
|
||||
pointed to by this node, or '' for other node types.
|
||||
"""
|
||||
target_type = ""
|
||||
if node.op in ("call_function", "call_method"):
|
||||
target_type = torch.typename(node.target)
|
||||
elif node.op == "call_module":
|
||||
if not isinstance(node.target, str):
|
||||
raise AssertionError(f"Expected str, got {type(node.target)}")
|
||||
target_mod = getattr_from_fqn(gm, node.target)
|
||||
target_type = torch.typename(target_mod)
|
||||
return target_type
|
||||
|
||||
|
||||
def rekey_logger_info_on_node_name_of_model(
|
||||
results: NSResultsType,
|
||||
model_name: str,
|
||||
) -> NSResultsType:
|
||||
"""
|
||||
Rekeys the layer name of a results dictionary to use node names
|
||||
from `model_name`.
|
||||
|
||||
For example, transforms
|
||||
|
||||
{'base_op_1_0': {'node_output': {'model_a':
|
||||
[{'ref_node_name': 'linear1', ...}]}}}
|
||||
|
||||
into
|
||||
|
||||
{'linear1': {'node_output': {'model_a':
|
||||
[{'ref_node_name': 'linear1', ...}]}}}
|
||||
|
||||
Note: we cannot use these node names directly because they are not
|
||||
guaranteed to be consistent across models. This is why we extract
|
||||
the results first and rekey afterwards.
|
||||
"""
|
||||
new_results = {}
|
||||
for old_layer_name, result_type_to_results in results.items():
|
||||
new_layer_name = None
|
||||
for model_name_to_results in result_type_to_results.values():
|
||||
for cur_model_name, list_of_results in model_name_to_results.items():
|
||||
if cur_model_name == model_name:
|
||||
if len(list_of_results) == 0:
|
||||
raise AssertionError("Expected list_of_results to be not empty")
|
||||
new_layer_name = list_of_results[0]["ref_node_name"]
|
||||
else:
|
||||
continue
|
||||
if new_layer_name is not None:
|
||||
new_results[new_layer_name] = result_type_to_results
|
||||
else:
|
||||
new_results[old_layer_name] = result_type_to_results
|
||||
return new_results
|
||||
|
||||
|
||||
def maybe_add_missing_fqns(results: NSResultsType) -> None:
|
||||
"""
|
||||
If `fqn` entries are filled in for one of the models in `results`, copies
|
||||
them over to any models which do not have them filled out.
|
||||
|
||||
A common use case benefitting from this is comparing a model prepared by
|
||||
quantization to a quantized model. In this case, the model prepared by
|
||||
quantization would have `fqn` entries, and the quantized model would not.
|
||||
"""
|
||||
|
||||
# Check in the first result to find any model with fqn entries defined.
|
||||
model_name_with_fqns = None
|
||||
for result_type_to_results in results.values():
|
||||
for model_name_to_results in result_type_to_results.values():
|
||||
for model_name, model_results in model_name_to_results.items():
|
||||
if len(model_results) > 0:
|
||||
if model_results[0]["fqn"] is not None:
|
||||
model_name_with_fqns = model_name
|
||||
break
|
||||
break
|
||||
break
|
||||
|
||||
if model_name_with_fqns:
|
||||
for result_type_to_results in results.values():
|
||||
for model_name_to_results in result_type_to_results.values():
|
||||
ref_model_results = model_name_to_results[model_name_with_fqns]
|
||||
for model_name, model_results in model_name_to_results.items():
|
||||
if model_name == model_name_with_fqns:
|
||||
continue
|
||||
|
||||
for i in range(len(model_results)):
|
||||
fqn = ref_model_results[i]["fqn"]
|
||||
model_results[i]["fqn"] = fqn
|
||||
|
||||
|
||||
def maybe_dequantize_first_two_tensor_args_and_handle_tuples(f):
|
||||
def inner(*args, **kwargs):
|
||||
a0, a1, *a_other = args
|
||||
|
||||
if (isinstance(a0, tuple) and isinstance(a1, tuple)) or (
|
||||
isinstance(a0, list) and isinstance(a1, list)
|
||||
):
|
||||
results = []
|
||||
for el0, el1 in zip(a0, a1):
|
||||
new_args = (el0, el1, *a_other)
|
||||
results.append(inner(*new_args, **kwargs))
|
||||
return results
|
||||
|
||||
elif isinstance(a0, torch.Tensor) and isinstance(a1, torch.Tensor):
|
||||
if a0.is_quantized:
|
||||
a0 = a0.dequantize()
|
||||
if a1.is_quantized:
|
||||
a1 = a1.dequantize()
|
||||
|
||||
# for the purposes of this util, only handle floats
|
||||
if a0.dtype != torch.float or a1.dtype != torch.float:
|
||||
return None
|
||||
|
||||
new_args = (a0, a1, *a_other)
|
||||
return f(*new_args, **kwargs)
|
||||
|
||||
return inner
|
||||
|
||||
|
||||
@maybe_dequantize_first_two_tensor_args_and_handle_tuples
|
||||
def compute_sqnr(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Computes the SQNR between `x` and `y`.
|
||||
|
||||
Args:
|
||||
x: Tensor or tuple of tensors
|
||||
y: Tensor or tuple of tensors
|
||||
|
||||
Return:
|
||||
float or tuple of floats
|
||||
"""
|
||||
Ps = torch.norm(x)
|
||||
Pn = torch.norm(x - y)
|
||||
return 20 * torch.log10(Ps / Pn)
|
||||
|
||||
|
||||
@maybe_dequantize_first_two_tensor_args_and_handle_tuples
|
||||
def compute_normalized_l2_error(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Computes the normalized L2 error between `x` and `y`.
|
||||
|
||||
Args:
|
||||
x: Tensor or tuple of tensors
|
||||
y: Tensor or tuple of tensors
|
||||
|
||||
Return:
|
||||
float or tuple of floats
|
||||
"""
|
||||
|
||||
return torch.sqrt(((x - y) ** 2).sum() / (x**2).sum())
|
||||
|
||||
|
||||
@maybe_dequantize_first_two_tensor_args_and_handle_tuples
|
||||
def compute_cosine_similarity(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Computes the cosine similarity between `x` and `y`.
|
||||
|
||||
Args:
|
||||
x: Tensor or tuple of tensors
|
||||
y: Tensor or tuple of tensors
|
||||
|
||||
Return:
|
||||
float or tuple of floats
|
||||
"""
|
||||
# For convolutions, the shape of the quantized weight has one additional
|
||||
# dimension compared to the shape of the fp32 weight. Match the shapes
|
||||
# to enable cosine similarity comparison.
|
||||
x = x.reshape(1, -1)
|
||||
y = y.reshape(1, -1)
|
||||
return torch.nn.functional.cosine_similarity(x, y)
|
||||
|
||||
|
||||
def op_type_supports_shadowing(node: Node) -> bool:
|
||||
if node.op == "call_function":
|
||||
if node.target in (
|
||||
torch.add,
|
||||
torch.mul,
|
||||
operator.add,
|
||||
operator.mul,
|
||||
torch.cat,
|
||||
torch.stack,
|
||||
):
|
||||
# shadowing for ops with multiple tensor inputs is not implemented yet
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def get_normalized_nth_input(node: Node, gm: GraphModule, idx: int) -> Node:
|
||||
"""
|
||||
Given a node, gets the n'th input to that node, normalizing
|
||||
args and kwargs to the best of its ability.
|
||||
"""
|
||||
try:
|
||||
norm_args_and_kwargs = node.normalized_arguments(
|
||||
gm, normalize_to_only_use_kwargs=True
|
||||
)
|
||||
if norm_args_and_kwargs is not None:
|
||||
norm_args, norm_kwargs = norm_args_and_kwargs
|
||||
if len(norm_args) + len(norm_kwargs) <= idx:
|
||||
raise AssertionError(
|
||||
f"Index {idx} out of range: total = {len(norm_args) + len(norm_kwargs)}"
|
||||
)
|
||||
if idx < len(norm_args):
|
||||
return norm_args[idx]
|
||||
else:
|
||||
# note: in Python 3.7+ dicts are ordered
|
||||
return list(norm_kwargs.values())[idx]
|
||||
else:
|
||||
if len(node.args) + len(node.kwargs) <= idx:
|
||||
raise AssertionError(
|
||||
f"Index {idx} out of range: total = {len(node.args) + len(node.kwargs)}"
|
||||
)
|
||||
if idx < len(node.args):
|
||||
return node.args[idx] # type: ignore[return-value]
|
||||
else:
|
||||
kwargs_idx = idx + len(node.args)
|
||||
return list(node.kwargs.values())[kwargs_idx] # type: ignore[return-value]
|
||||
except RuntimeError:
|
||||
# this RuntimeError happens when node argument normalization
|
||||
# requires typehints to proceed, such as for torch.add where
|
||||
# either the first, second or both arguments could be tensors
|
||||
if len(node.args) + len(node.kwargs) <= idx:
|
||||
raise AssertionError(
|
||||
f"Index {idx} out of range: total = {len(node.args) + len(node.kwargs)}"
|
||||
) from None
|
||||
if idx < len(node.args):
|
||||
return node.args[idx] # type: ignore[return-value]
|
||||
else:
|
||||
kwargs_idx = idx + len(node.args)
|
||||
return list(node.kwargs.values())[kwargs_idx] # type: ignore[return-value]
|
||||
@@ -0,0 +1,302 @@
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
import torch.ao.nn.intrinsic as nni
|
||||
import torch.ao.nn.intrinsic.qat as nniqat
|
||||
import torch.ao.nn.intrinsic.quantized as nniq
|
||||
import torch.ao.nn.qat as nnqat
|
||||
import torch.ao.nn.quantized as nnq
|
||||
import torch.ao.nn.quantized.dynamic as nnqd
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.fx import GraphModule
|
||||
from torch.fx.graph import Node
|
||||
|
||||
from .ns_types import NSSingleResultType, NSSingleResultValuesType
|
||||
from .utils import get_target_type_str, getattr_from_fqn, return_first_non_observer_node
|
||||
|
||||
|
||||
toq = torch.ops.quantized
|
||||
|
||||
|
||||
def mod_weight_detach(mod: nn.Module) -> torch.Tensor:
|
||||
return mod.weight.detach() # type: ignore[operator]
|
||||
|
||||
|
||||
def mod_0_weight_detach(mod: nn.Module) -> torch.Tensor:
|
||||
return mod[0].weight.detach() # type: ignore[index]
|
||||
|
||||
|
||||
def mod_weight_bias_0(mod: nn.Module) -> torch.Tensor:
|
||||
return mod._weight_bias()[0] # type: ignore[operator]
|
||||
|
||||
|
||||
def get_lstm_weight(mod: nn.Module) -> list[torch.Tensor]:
|
||||
res = []
|
||||
for idx, param_name in enumerate(mod._flat_weights_names): # type: ignore[arg-type]
|
||||
if "weight_ih_l" in param_name or "weight_hh_l" in param_name:
|
||||
param_value = mod._flat_weights[idx].detach() # type: ignore[index,union-attr]
|
||||
res.append(param_value)
|
||||
return res
|
||||
|
||||
|
||||
def get_qlstm_weight(mod: nn.Module) -> list[torch.Tensor]:
|
||||
res = []
|
||||
for weight_value in mod._all_weight_values: # type: ignore[union-attr]
|
||||
res.append(weight_value.param.__getstate__()[0][4][0].__getstate__()[0][0])
|
||||
res.append(weight_value.param.__getstate__()[0][4][1].__getstate__()[0][0])
|
||||
return res
|
||||
|
||||
|
||||
def get_conv_mod_weight(mod: nn.Module) -> torch.Tensor:
|
||||
if isinstance(mod, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
|
||||
return mod.weight.detach()
|
||||
elif isinstance(mod, (nni.ConvReLU1d, nni.ConvReLU2d, nni.ConvReLU3d)):
|
||||
return mod[0].weight.detach() # type: ignore[operator]
|
||||
else:
|
||||
return mod._weight_bias()[0] # type: ignore[operator]
|
||||
|
||||
|
||||
def get_linear_mod_weight(mod: nn.Module) -> torch.Tensor:
|
||||
if isinstance(mod, nn.Linear):
|
||||
return mod.weight.detach()
|
||||
elif isinstance(mod, nni.LinearReLU):
|
||||
return mod[0].weight.detach() # type: ignore[operator]
|
||||
else:
|
||||
return mod._weight_bias()[0] # type: ignore[operator]
|
||||
|
||||
|
||||
def get_lstm_mod_weights(mod: nn.Module) -> list[torch.Tensor]:
|
||||
# TODO(future PR): make more generic, handle everything
|
||||
if isinstance(mod, nn.LSTM):
|
||||
res = []
|
||||
for idx, param_name in enumerate(mod._flat_weights_names):
|
||||
if "weight_ih_l" in param_name or "weight_hh_l" in param_name:
|
||||
param_value = mod._flat_weights[idx].detach() # type: ignore[index,union-attr]
|
||||
res.append(param_value)
|
||||
return res
|
||||
else:
|
||||
if not isinstance(mod, nnqd.LSTM):
|
||||
raise AssertionError(f"type {type(mod)} not handled yet")
|
||||
res = []
|
||||
for weight_value in mod._all_weight_values:
|
||||
res.append(
|
||||
weight_value.param.__getstate__()[0][4][0].__getstate__()[0][0] # type: ignore[index]
|
||||
)
|
||||
res.append(
|
||||
weight_value.param.__getstate__()[0][4][1].__getstate__()[0][0] # type: ignore[index]
|
||||
)
|
||||
return res
|
||||
|
||||
|
||||
def get_conv_fun_weight(node: Node, gm: GraphModule) -> torch.Tensor:
|
||||
# traverse backwards from the weight arg, accounting for any observers
|
||||
weight_arg_node = node.args[1]
|
||||
if not isinstance(weight_arg_node, Node):
|
||||
raise AssertionError(f"Expected Node, got {type(weight_arg_node)}")
|
||||
weight_node = return_first_non_observer_node(weight_arg_node, gm)
|
||||
if not isinstance(weight_node, Node):
|
||||
raise AssertionError(f"Expected Node, got {type(weight_node)}")
|
||||
if weight_node.op != "get_attr":
|
||||
raise AssertionError(f"Expected get_attr, got {weight_node.op}")
|
||||
weight = getattr_from_fqn(gm, weight_node.target) # type: ignore[arg-type]
|
||||
return weight.detach()
|
||||
|
||||
|
||||
def get_qconv_fun_weight(node: Node, gm: GraphModule) -> torch.Tensor:
|
||||
# qconv state is arg 1
|
||||
qconv_state_node = node.args[1]
|
||||
if not isinstance(qconv_state_node, Node):
|
||||
raise AssertionError(f"Expected Node, got {type(qconv_state_node)}")
|
||||
if qconv_state_node.op != "get_attr":
|
||||
raise AssertionError(f"Expected get_attr, got {qconv_state_node.op}")
|
||||
qconv_state_obj = getattr_from_fqn(gm, qconv_state_node.target) # type: ignore[arg-type]
|
||||
return qconv_state_obj.weight()
|
||||
|
||||
|
||||
def get_linear_fun_weight(node: Node, gm: GraphModule) -> torch.Tensor:
|
||||
# traverse backwards from the weight arg, accounting for any observers
|
||||
# supported patterns:
|
||||
# weight -> obs -> linear
|
||||
# weight -> to(torch.float16) -> dequantize -> linear
|
||||
linear_second_arg = node.args[1]
|
||||
if not isinstance(linear_second_arg, Node):
|
||||
raise AssertionError(f"Expected Node, got {type(linear_second_arg)}")
|
||||
|
||||
if linear_second_arg.op == "call_module":
|
||||
# weight -> obs -> linear
|
||||
weight_arg_node = node.args[1]
|
||||
if not isinstance(weight_arg_node, Node):
|
||||
raise AssertionError(f"Expected Node, got {type(weight_arg_node)}")
|
||||
weight_node = weight_arg_node.args[0]
|
||||
if not isinstance(weight_node, Node):
|
||||
raise AssertionError(f"Expected Node, got {type(weight_node)}")
|
||||
if weight_node.op != "get_attr":
|
||||
raise AssertionError(f"Expected get_attr, got {weight_node.op}")
|
||||
weight = getattr_from_fqn(gm, weight_node.target) # type: ignore[arg-type]
|
||||
return weight.detach()
|
||||
elif linear_second_arg.op == "call_method":
|
||||
# weight -> to(torch.float16) -> dequantize -> linear
|
||||
if linear_second_arg.op != "call_method":
|
||||
raise AssertionError(f"Expected call_method, got {linear_second_arg.op}")
|
||||
dequant_node = node.args[1]
|
||||
if not isinstance(dequant_node, Node):
|
||||
raise AssertionError(f"Expected Node, got {type(dequant_node)}")
|
||||
to_fp16_node = dequant_node.args[0]
|
||||
if not isinstance(to_fp16_node, Node):
|
||||
raise AssertionError(f"Expected Node, got {type(to_fp16_node)}")
|
||||
# extract the dtype, so we can cast to it before returning
|
||||
target_dtype = to_fp16_node.args[1]
|
||||
weight_node = to_fp16_node.args[0]
|
||||
if not isinstance(weight_node, Node):
|
||||
raise AssertionError(f"Expected Node, got {type(weight_node)}")
|
||||
if weight_node.op != "get_attr":
|
||||
raise AssertionError(f"Expected get_attr, got {weight_node.op}")
|
||||
weight = getattr_from_fqn(gm, weight_node.target) # type: ignore[arg-type]
|
||||
# return the weight with fp16 cast
|
||||
return weight.detach().to(target_dtype)
|
||||
else:
|
||||
if linear_second_arg.op != "get_attr":
|
||||
raise AssertionError(f"Expected get_attr, got {linear_second_arg.op}")
|
||||
weight = getattr_from_fqn(gm, linear_second_arg.target) # type: ignore[arg-type]
|
||||
return weight.detach()
|
||||
|
||||
|
||||
def get_qlinear_fun_weight(node: Node, gm: GraphModule) -> torch.Tensor:
|
||||
# packed weight is arg 1
|
||||
packed_weight_node = node.args[1]
|
||||
if not isinstance(packed_weight_node, Node):
|
||||
raise AssertionError(f"Expected Node, got {type(packed_weight_node)}")
|
||||
if packed_weight_node.op != "get_attr":
|
||||
raise AssertionError(f"Expected get_attr, got {packed_weight_node.op}")
|
||||
packed_weight = getattr_from_fqn(gm, packed_weight_node.target) # type: ignore[arg-type]
|
||||
# TODO(future PR): why does packed_weight.unpack() not work?
|
||||
(weight, _bias), _name = packed_weight.__getstate__()
|
||||
return weight
|
||||
|
||||
|
||||
def get_op_to_type_to_weight_extraction_fn() -> dict[str, dict[Callable, Callable]]:
|
||||
op_to_type_to_weight_extraction_fn: dict[str, dict[Callable, Callable]] = {
|
||||
"call_module": {
|
||||
# Conv1d
|
||||
nn.Conv1d: mod_weight_detach,
|
||||
nni.ConvReLU1d: mod_0_weight_detach,
|
||||
nnq.Conv1d: mod_weight_bias_0,
|
||||
nnqat.Conv1d: mod_weight_detach,
|
||||
nniqat.ConvBn1d: mod_weight_detach,
|
||||
nniqat.ConvBnReLU1d: mod_weight_detach,
|
||||
nniqat.ConvReLU1d: mod_weight_detach,
|
||||
nniq.ConvReLU1d: mod_weight_bias_0,
|
||||
# Conv2d
|
||||
nn.Conv2d: mod_weight_detach,
|
||||
nni.ConvReLU2d: mod_0_weight_detach,
|
||||
nnq.Conv2d: mod_weight_bias_0,
|
||||
nnqat.Conv2d: mod_weight_detach,
|
||||
nniqat.ConvBn2d: mod_weight_detach,
|
||||
nniqat.ConvBnReLU2d: mod_weight_detach,
|
||||
nniqat.ConvReLU2d: mod_weight_detach,
|
||||
nniq.ConvReLU2d: mod_weight_bias_0,
|
||||
# Conv3d
|
||||
nn.Conv3d: mod_weight_detach,
|
||||
nni.ConvReLU3d: mod_0_weight_detach,
|
||||
nnq.Conv3d: mod_weight_bias_0,
|
||||
nnqat.Conv3d: mod_weight_detach,
|
||||
nniqat.ConvBn3d: mod_weight_detach,
|
||||
nniqat.ConvBnReLU3d: mod_weight_detach,
|
||||
nniqat.ConvReLU3d: mod_weight_detach,
|
||||
nniq.ConvReLU3d: mod_weight_bias_0,
|
||||
# Linear
|
||||
nn.Linear: mod_weight_detach,
|
||||
nnq.Linear: mod_weight_bias_0,
|
||||
nni.LinearReLU: mod_0_weight_detach,
|
||||
nniq.LinearReLU: mod_weight_bias_0,
|
||||
nnqat.Linear: mod_weight_detach,
|
||||
nnqd.Linear: mod_weight_bias_0,
|
||||
nniqat.LinearReLU: mod_weight_detach,
|
||||
nniqat.LinearBn1d: mod_weight_detach,
|
||||
nn.modules.linear.NonDynamicallyQuantizableLinear: mod_weight_detach,
|
||||
# LSTM
|
||||
nn.LSTM: get_lstm_weight,
|
||||
nnqd.LSTM: get_qlstm_weight,
|
||||
},
|
||||
"call_function": {
|
||||
# Conv
|
||||
F.conv1d: get_conv_fun_weight,
|
||||
F.conv2d: get_conv_fun_weight,
|
||||
F.conv3d: get_conv_fun_weight,
|
||||
toq.conv1d: get_qconv_fun_weight,
|
||||
toq.conv2d: get_qconv_fun_weight,
|
||||
toq.conv3d: get_qconv_fun_weight,
|
||||
toq.conv1d_relu: get_qconv_fun_weight,
|
||||
toq.conv2d_relu: get_qconv_fun_weight,
|
||||
toq.conv3d_relu: get_qconv_fun_weight,
|
||||
# Linear
|
||||
F.linear: get_linear_fun_weight,
|
||||
toq.linear: get_qlinear_fun_weight,
|
||||
toq.linear_relu: get_qlinear_fun_weight,
|
||||
},
|
||||
}
|
||||
|
||||
return op_to_type_to_weight_extraction_fn
|
||||
|
||||
|
||||
def extract_weight_from_node(
|
||||
node: Node,
|
||||
gm: GraphModule,
|
||||
op_to_type_to_weight_extraction_fn: dict[str, dict[Callable, Callable]]
|
||||
| None = None,
|
||||
) -> NSSingleResultType | None:
|
||||
res_type = NSSingleResultValuesType.WEIGHT.value
|
||||
|
||||
# Not all graphmodules have _node_name_to_scope, so only fill it
|
||||
# out if it exists.
|
||||
fqn = None
|
||||
if hasattr(gm, "_node_name_to_scope"):
|
||||
fqn = gm._node_name_to_scope[node.name][0] # type: ignore[index]
|
||||
|
||||
if op_to_type_to_weight_extraction_fn is None:
|
||||
op_to_type_to_weight_extraction_fn = get_op_to_type_to_weight_extraction_fn()
|
||||
|
||||
ref_node_type = get_target_type_str(node, gm)
|
||||
# for extracting weights, these are always the same
|
||||
prev_node_type = ref_node_type
|
||||
|
||||
if node.op == "call_function":
|
||||
function_mapping = op_to_type_to_weight_extraction_fn["call_function"]
|
||||
for target_fn_type, weight_extraction_fn in function_mapping.items():
|
||||
if node.target == target_fn_type:
|
||||
weight = weight_extraction_fn(node, gm)
|
||||
return {
|
||||
"type": res_type,
|
||||
"values": [weight],
|
||||
"prev_node_name": node.name,
|
||||
"prev_node_target_type": prev_node_type,
|
||||
"ref_node_name": node.name,
|
||||
"ref_node_target_type": ref_node_type,
|
||||
"index_within_arg": 0,
|
||||
"index_of_arg": 0,
|
||||
"fqn": fqn,
|
||||
}
|
||||
|
||||
elif node.op == "call_module":
|
||||
# for call_module, we need to look up the modules to do the type check
|
||||
if not isinstance(node.target, str):
|
||||
raise AssertionError(f"Expected str, got {type(node.target)}")
|
||||
mod = getattr_from_fqn(gm, node.target)
|
||||
module_mapping = op_to_type_to_weight_extraction_fn["call_module"]
|
||||
for target_mod_type, weight_extraction_fn in module_mapping.items():
|
||||
if type(mod) is target_mod_type:
|
||||
weight = weight_extraction_fn(mod)
|
||||
return {
|
||||
"type": res_type,
|
||||
"values": [weight],
|
||||
"prev_node_name": node.name,
|
||||
"prev_node_target_type": prev_node_type,
|
||||
"ref_node_name": node.name,
|
||||
"ref_node_target_type": ref_node_type,
|
||||
"index_within_arg": 0,
|
||||
"index_of_arg": 0,
|
||||
"fqn": fqn,
|
||||
}
|
||||
|
||||
return None
|
||||
@@ -0,0 +1,23 @@
|
||||
# Variables
|
||||
from ._mappings import (
|
||||
get_dynamic_sparse_quantized_mapping,
|
||||
get_static_sparse_quantized_mapping,
|
||||
)
|
||||
|
||||
# Scheduler
|
||||
from .scheduler.base_scheduler import BaseScheduler
|
||||
from .scheduler.cubic_scheduler import CubicSL
|
||||
from .scheduler.lambda_scheduler import LambdaSL
|
||||
|
||||
# Sparsifier
|
||||
from .sparsifier.base_sparsifier import BaseSparsifier
|
||||
from .sparsifier.nearly_diagonal_sparsifier import NearlyDiagonalSparsifier
|
||||
|
||||
# Parametrizations
|
||||
from .sparsifier.utils import (
|
||||
FakeSparsity,
|
||||
fqn_to_module,
|
||||
get_arg_info_from_tensor_fqn,
|
||||
module_to_fqn,
|
||||
)
|
||||
from .sparsifier.weight_norm_sparsifier import WeightNormSparsifier
|
||||
+482
@@ -0,0 +1,482 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import copy
|
||||
import warnings
|
||||
from collections import defaultdict
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.ao.pruning.sparsifier.utils import fqn_to_module, module_to_fqn
|
||||
|
||||
|
||||
__all__ = ["ActivationSparsifier"]
|
||||
|
||||
|
||||
class ActivationSparsifier:
|
||||
r"""
|
||||
The Activation sparsifier class aims to sparsify/prune activations in a neural
|
||||
network. The idea is to attach the sparsifier to a layer (or layers) and it
|
||||
zeroes out the activations based on the mask_fn (or sparsification function)
|
||||
input by the user.
|
||||
The mask_fn is applied once all the inputs are aggregated and reduced i.e.
|
||||
mask = mask_fn(reduce_fn(aggregate_fn(activations)))
|
||||
|
||||
Note::
|
||||
The sparsification mask is computed on the input **before it goes through the attached layer**.
|
||||
|
||||
Args:
|
||||
model (nn.Module):
|
||||
The model whose layers will be sparsified. The layers that needs to be
|
||||
sparsified should be added separately using the register_layer() function
|
||||
aggregate_fn (Optional, Callable):
|
||||
default aggregate_fn that is used if not specified while registering the layer.
|
||||
specifies how inputs should be aggregated over time.
|
||||
The aggregate_fn should usually take 2 torch tensors and return the aggregated tensor.
|
||||
Example
|
||||
def add_agg_fn(tensor1, tensor2): return tensor1 + tensor2
|
||||
reduce_fn (Optional, Callable):
|
||||
default reduce_fn that is used if not specified while registering the layer.
|
||||
reduce_fn will be called on the aggregated tensor i.e. the tensor obtained after
|
||||
calling agg_fn() on all inputs.
|
||||
Example
|
||||
def mean_reduce_fn(agg_tensor): return agg_tensor.mean(dim=0)
|
||||
mask_fn (Optional, Callable):
|
||||
default mask_fn that is used to create the sparsification mask using the tensor obtained after
|
||||
calling the reduce_fn(). This is used by default if a custom one is passed in the
|
||||
register_layer().
|
||||
Note that the mask_fn() definition should contain the sparse arguments that is passed in sparse_config
|
||||
arguments.
|
||||
features (Optional, list):
|
||||
default selected features to sparsify.
|
||||
If this is non-empty, then the mask_fn will be applied for each feature of the input.
|
||||
For example,
|
||||
mask = [mask_fn(reduce_fn(aggregated_fn(input[feature])) for feature in features]
|
||||
feature_dim (Optional, int):
|
||||
default dimension of input features. Again, features along this dim will be chosen
|
||||
for sparsification.
|
||||
sparse_config (Dict):
|
||||
Default configuration for the mask_fn. This config will be passed
|
||||
with the mask_fn()
|
||||
|
||||
Example:
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> model = SomeModel()
|
||||
>>> act_sparsifier = ActivationSparsifier(...) # init activation sparsifier
|
||||
>>> # Initialize aggregate_fn
|
||||
>>> def agg_fn(x, y):
|
||||
>>> return x + y
|
||||
>>>
|
||||
>>> # Initialize reduce_fn
|
||||
>>> def reduce_fn(x):
|
||||
>>> return torch.mean(x, dim=0)
|
||||
>>>
|
||||
>>> # Initialize mask_fn
|
||||
>>> def mask_fn(data):
|
||||
>>> return torch.eye(data.shape).to(data.device)
|
||||
>>>
|
||||
>>>
|
||||
>>> act_sparsifier.register_layer(
|
||||
... model.some_layer,
|
||||
... aggregate_fn=agg_fn,
|
||||
... reduce_fn=reduce_fn,
|
||||
... mask_fn=mask_fn,
|
||||
... )
|
||||
>>>
|
||||
>>> # start training process
|
||||
>>> for _ in [...]:
|
||||
>>> # epoch starts
|
||||
>>> # model.forward(), compute_loss() and model.backwards()
|
||||
>>> # epoch ends
|
||||
>>> act_sparsifier.step()
|
||||
>>> # end training process
|
||||
>>> sparsifier.squash_mask()
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: nn.Module,
|
||||
aggregate_fn=None,
|
||||
reduce_fn=None,
|
||||
mask_fn=None,
|
||||
features=None,
|
||||
feature_dim=None,
|
||||
**sparse_config,
|
||||
):
|
||||
self.model = model
|
||||
self.defaults: dict[str, Any] = defaultdict()
|
||||
self.defaults["sparse_config"] = sparse_config
|
||||
|
||||
# functions
|
||||
self.defaults["aggregate_fn"] = aggregate_fn
|
||||
self.defaults["reduce_fn"] = reduce_fn
|
||||
self.defaults["mask_fn"] = mask_fn
|
||||
|
||||
# default feature and feature_dim
|
||||
self.defaults["features"] = features
|
||||
self.defaults["feature_dim"] = feature_dim
|
||||
|
||||
self.data_groups: dict[str, dict] = defaultdict(
|
||||
dict
|
||||
) # contains all relevant info w.r.t each registered layer
|
||||
|
||||
self.state: dict[str, Any] = defaultdict(dict) # layer name -> mask
|
||||
|
||||
@staticmethod
|
||||
def _safe_rail_checks(args):
|
||||
"""Makes sure that some of the functions and attributes are not passed incorrectly"""
|
||||
|
||||
# if features are not None, then feature_dim must not be None
|
||||
features, feature_dim = args["features"], args["feature_dim"]
|
||||
if features is not None:
|
||||
if feature_dim is None:
|
||||
raise AssertionError("need feature dim to select features")
|
||||
|
||||
# all the *_fns should be callable
|
||||
fn_keys = ["aggregate_fn", "reduce_fn", "mask_fn"]
|
||||
for key in fn_keys:
|
||||
fn = args[key]
|
||||
if not callable(fn):
|
||||
raise AssertionError(f"{fn} must be callable")
|
||||
|
||||
def _aggregate_hook(self, name):
|
||||
"""Returns hook that computes aggregate of activations passing through."""
|
||||
|
||||
# gather some data
|
||||
feature_dim = self.data_groups[name]["feature_dim"]
|
||||
features = self.data_groups[name]["features"]
|
||||
agg_fn = self.data_groups[name]["aggregate_fn"]
|
||||
|
||||
def hook(module, input) -> None:
|
||||
input_data = input[0]
|
||||
|
||||
data = self.data_groups[name].get("data") # aggregated data
|
||||
if features is None:
|
||||
# no features associated, data should not be a list
|
||||
if data is None:
|
||||
data = torch.zeros_like(input_data)
|
||||
self.state[name]["mask"] = torch.ones_like(input_data)
|
||||
out_data = agg_fn(data, input_data)
|
||||
else:
|
||||
# data should be a list [aggregated over each feature only]
|
||||
if data is None:
|
||||
out_data = [
|
||||
0 for _ in range(len(features))
|
||||
] # create one in case of 1st forward
|
||||
self.state[name]["mask"] = [0 for _ in range(len(features))]
|
||||
else:
|
||||
out_data = data # a list
|
||||
|
||||
# compute aggregate over each feature
|
||||
for feature_idx in range(len(features)):
|
||||
# each feature is either a list or scalar, convert it to torch tensor
|
||||
feature_tensor = (
|
||||
torch.Tensor([features[feature_idx]])
|
||||
.long()
|
||||
.to(input_data.device)
|
||||
)
|
||||
data_feature = torch.index_select(
|
||||
input_data, feature_dim, feature_tensor
|
||||
)
|
||||
if data is None:
|
||||
curr_data = torch.zeros_like(data_feature)
|
||||
self.state[name]["mask"][feature_idx] = torch.ones_like(
|
||||
data_feature
|
||||
)
|
||||
else:
|
||||
curr_data = data[feature_idx]
|
||||
out_data[feature_idx] = agg_fn(curr_data, data_feature)
|
||||
self.data_groups[name]["data"] = out_data
|
||||
|
||||
return hook
|
||||
|
||||
def register_layer(
|
||||
self,
|
||||
layer: nn.Module,
|
||||
aggregate_fn=None,
|
||||
reduce_fn=None,
|
||||
mask_fn=None,
|
||||
features=None,
|
||||
feature_dim=None,
|
||||
**sparse_config,
|
||||
):
|
||||
r"""
|
||||
Registers a layer for sparsification. The layer should be part of self.model.
|
||||
Specifically, registers a pre-forward hook to the layer. The hook will apply the aggregate_fn
|
||||
and store the aggregated activations that is input over each step.
|
||||
|
||||
Note::
|
||||
- There is no need to pass in the name of the layer as it is automatically computed as per
|
||||
the fqn convention.
|
||||
|
||||
- All the functions (fn) passed as argument will be called at a dim, feature level.
|
||||
"""
|
||||
name = module_to_fqn(self.model, layer)
|
||||
if name is None:
|
||||
raise AssertionError("layer not found in the model")
|
||||
|
||||
if name in self.data_groups: # unregister layer if already present
|
||||
warnings.warn(
|
||||
"layer already attached to the sparsifier, deregistering the layer and registering with new config",
|
||||
stacklevel=2,
|
||||
)
|
||||
self.unregister_layer(name=name)
|
||||
|
||||
local_args = copy.deepcopy(self.defaults)
|
||||
update_dict = {
|
||||
"aggregate_fn": aggregate_fn,
|
||||
"reduce_fn": reduce_fn,
|
||||
"mask_fn": mask_fn,
|
||||
"features": features,
|
||||
"feature_dim": feature_dim,
|
||||
"layer": layer,
|
||||
}
|
||||
local_args.update(
|
||||
(arg, val) for arg, val in update_dict.items() if val is not None
|
||||
)
|
||||
local_args["sparse_config"].update(sparse_config)
|
||||
|
||||
self._safe_rail_checks(local_args)
|
||||
|
||||
self.data_groups[name] = local_args
|
||||
agg_hook = layer.register_forward_pre_hook(self._aggregate_hook(name=name))
|
||||
|
||||
self.state[name]["mask"] = (
|
||||
None # mask will be created when model forward is called.
|
||||
)
|
||||
|
||||
# attach agg hook
|
||||
self.data_groups[name]["hook"] = agg_hook
|
||||
|
||||
# for serialization purposes, we know whether aggregate_hook is attached
|
||||
# or sparsify_hook()
|
||||
self.data_groups[name]["hook_state"] = "aggregate" # aggregate hook is attached
|
||||
|
||||
def get_mask(self, name: str | None = None, layer: nn.Module | None = None):
|
||||
"""
|
||||
Returns mask associated to the layer.
|
||||
|
||||
The mask is
|
||||
- a torch tensor is features for that layer is None.
|
||||
- a list of torch tensors for each feature, otherwise
|
||||
|
||||
Note::
|
||||
The shape of the mask is unknown until model.forward() is applied.
|
||||
Hence, if get_mask() is called before model.forward(), an
|
||||
error will be raised.
|
||||
"""
|
||||
if name is None and layer is None:
|
||||
raise AssertionError("Need at least name or layer obj to retrieve mask")
|
||||
|
||||
if name is None:
|
||||
if layer is None:
|
||||
raise AssertionError("layer must be provided when name is None")
|
||||
name = module_to_fqn(self.model, layer)
|
||||
if name is None:
|
||||
raise AssertionError("layer not found in the specified model")
|
||||
|
||||
if name not in self.state:
|
||||
raise ValueError("Error: layer with the given name not found")
|
||||
|
||||
mask = self.state[name].get("mask", None)
|
||||
|
||||
if mask is None:
|
||||
raise ValueError(
|
||||
"Error: shape unknown, call layer() routine at least once to infer mask"
|
||||
)
|
||||
return mask
|
||||
|
||||
def unregister_layer(self, name):
|
||||
"""Detaches the sparsifier from the layer"""
|
||||
|
||||
# detach any hooks attached
|
||||
self.data_groups[name]["hook"].remove()
|
||||
|
||||
# pop from the state dict
|
||||
self.state.pop(name)
|
||||
|
||||
# pop from the data groups
|
||||
self.data_groups.pop(name)
|
||||
|
||||
def step(self):
|
||||
"""Internally calls the update_mask() function for each layer"""
|
||||
with torch.no_grad():
|
||||
for name, configs in self.data_groups.items():
|
||||
data = configs["data"]
|
||||
self.update_mask(name, data, configs)
|
||||
|
||||
self.data_groups[name].pop("data") # reset the accumulated data
|
||||
|
||||
def update_mask(self, name, data, configs):
|
||||
"""
|
||||
Called for each registered layer and does the following-
|
||||
1. apply reduce_fn on the aggregated activations
|
||||
2. use mask_fn to compute the sparsification mask
|
||||
|
||||
Note:
|
||||
the reduce_fn and mask_fn is called for each feature, dim over the data
|
||||
"""
|
||||
mask = self.get_mask(name)
|
||||
sparse_config = configs["sparse_config"]
|
||||
features = configs["features"]
|
||||
reduce_fn = configs["reduce_fn"]
|
||||
mask_fn = configs["mask_fn"]
|
||||
if features is None:
|
||||
data = reduce_fn(data)
|
||||
mask.data = mask_fn(data, **sparse_config)
|
||||
else:
|
||||
for feature_idx in range(len(features)):
|
||||
data_feature = reduce_fn(data[feature_idx])
|
||||
mask[feature_idx].data = mask_fn(data_feature, **sparse_config)
|
||||
|
||||
def _sparsify_hook(self, name):
|
||||
"""Returns hook that applies sparsification mask to input entering the attached layer"""
|
||||
mask = self.get_mask(name)
|
||||
features = self.data_groups[name]["features"]
|
||||
feature_dim = self.data_groups[name]["feature_dim"]
|
||||
|
||||
def hook(module, input):
|
||||
input_data = input[0]
|
||||
if features is None:
|
||||
# apply to all the features
|
||||
return input_data * mask
|
||||
else:
|
||||
# apply per feature, feature_dim
|
||||
for feature_idx in range(len(features)):
|
||||
feature = (
|
||||
torch.Tensor([features[feature_idx]])
|
||||
.long()
|
||||
.to(input_data.device)
|
||||
)
|
||||
sparsified = (
|
||||
torch.index_select(input_data, feature_dim, feature)
|
||||
* mask[feature_idx]
|
||||
)
|
||||
input_data.index_copy_(feature_dim, feature, sparsified)
|
||||
return input_data
|
||||
|
||||
return hook
|
||||
|
||||
def squash_mask(self, attach_sparsify_hook=True, **kwargs):
|
||||
"""
|
||||
Unregisters aggregate hook that was applied earlier and registers sparsification hooks if
|
||||
attach_sparsify_hook = True.
|
||||
"""
|
||||
for name, configs in self.data_groups.items():
|
||||
# unhook agg hook
|
||||
configs["hook"].remove()
|
||||
configs.pop("hook")
|
||||
self.data_groups[name]["hook_state"] = "None"
|
||||
if attach_sparsify_hook:
|
||||
configs["hook"] = configs["layer"].register_forward_pre_hook(
|
||||
self._sparsify_hook(name)
|
||||
)
|
||||
configs["hook_state"] = (
|
||||
"sparsify" # signals that sparsify hook is now attached
|
||||
)
|
||||
|
||||
def _get_serializable_data_groups(self):
|
||||
"""Exclude hook and layer from the config keys before serializing
|
||||
|
||||
TODO: Might have to treat functions (reduce_fn, mask_fn etc) in a different manner while serializing.
|
||||
For time-being, functions are treated the same way as other attributes
|
||||
"""
|
||||
data_groups: dict[str, Any] = defaultdict()
|
||||
for name, config in self.data_groups.items():
|
||||
new_config = {
|
||||
key: value
|
||||
for key, value in config.items()
|
||||
if key not in ["hook", "layer"]
|
||||
}
|
||||
data_groups[name] = new_config
|
||||
return data_groups
|
||||
|
||||
def _convert_mask(self, states_dict, sparse_coo=True):
|
||||
r"""Converts the mask to sparse coo or dense depending on the `sparse_coo` argument.
|
||||
If `sparse_coo=True`, then the mask is stored as sparse coo else dense tensor
|
||||
"""
|
||||
states = copy.deepcopy(states_dict)
|
||||
for state in states.values():
|
||||
if state["mask"] is not None:
|
||||
if isinstance(state["mask"], list):
|
||||
for idx in range(len(state["mask"])):
|
||||
if sparse_coo:
|
||||
state["mask"][idx] = state["mask"][idx].to_sparse_coo()
|
||||
else:
|
||||
state["mask"][idx] = state["mask"][idx].to_dense()
|
||||
else:
|
||||
if sparse_coo:
|
||||
state["mask"] = state["mask"].to_sparse_coo()
|
||||
else:
|
||||
state["mask"] = state["mask"].to_dense()
|
||||
return states
|
||||
|
||||
def state_dict(self) -> dict[str, Any]:
|
||||
r"""Returns the state of the sparsifier as a :class:`dict`.
|
||||
|
||||
It contains:
|
||||
* state - contains name -> mask mapping.
|
||||
* data_groups - a dictionary containing all config information for each
|
||||
layer
|
||||
* defaults - the default config while creating the constructor
|
||||
"""
|
||||
data_groups = self._get_serializable_data_groups()
|
||||
state = self._convert_mask(self.state)
|
||||
return {"state": state, "data_groups": data_groups, "defaults": self.defaults}
|
||||
|
||||
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
|
||||
r"""The load_state_dict() restores the state of the sparsifier based on the state_dict
|
||||
|
||||
Args:
|
||||
* state_dict - the dictionary that to which the current sparsifier needs to be restored to
|
||||
"""
|
||||
state = state_dict["state"]
|
||||
data_groups, defaults = state_dict["data_groups"], state_dict["defaults"]
|
||||
|
||||
self.__set_state__(
|
||||
{"state": state, "data_groups": data_groups, "defaults": defaults}
|
||||
)
|
||||
|
||||
def __get_state__(self) -> dict[str, Any]:
|
||||
data_groups = self._get_serializable_data_groups()
|
||||
state = self._convert_mask(self.state)
|
||||
return {
|
||||
"defaults": self.defaults,
|
||||
"state": state,
|
||||
"data_groups": data_groups,
|
||||
}
|
||||
|
||||
def __set_state__(self, state: dict[str, Any]) -> None:
|
||||
state["state"] = self._convert_mask(
|
||||
state["state"], sparse_coo=False
|
||||
) # convert mask to dense tensor
|
||||
self.__dict__.update(state)
|
||||
|
||||
# need to attach layer and hook info into the data_groups
|
||||
for name, config in self.data_groups.items():
|
||||
# fetch layer
|
||||
layer = fqn_to_module(self.model, name)
|
||||
if layer is None:
|
||||
raise AssertionError(f"layer {name} not found in the model")
|
||||
|
||||
# if agg_mode is True, then layer in aggregate mode
|
||||
if "hook_state" in config and config["hook_state"] == "aggregate":
|
||||
hook = layer.register_forward_pre_hook(self._aggregate_hook(name))
|
||||
|
||||
elif "hook_state" in config and config["hook_state"] == "sparsify":
|
||||
hook = layer.register_forward_pre_hook(self._sparsify_hook(name))
|
||||
|
||||
config["layer"] = layer
|
||||
config["hook"] = hook # type: ignore[possibly-undefined]
|
||||
|
||||
def __repr__(self):
|
||||
format_string = self.__class__.__name__ + " ("
|
||||
for name, config in self.data_groups.items():
|
||||
format_string += "\n"
|
||||
format_string += "\tData Group\n"
|
||||
format_string += f"\t name: {name}\n"
|
||||
for key in sorted(config.keys()):
|
||||
if key in ["data", "hook", "reduce_fn", "mask_fn", "aggregate_fn"]:
|
||||
continue
|
||||
format_string += f"\t {key}: {config[key]}\n"
|
||||
format_string += ")"
|
||||
return format_string
|
||||
+6
@@ -0,0 +1,6 @@
|
||||
from .base_data_scheduler import BaseDataScheduler
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BaseDataScheduler",
|
||||
]
|
||||
+199
@@ -0,0 +1,199 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import abc
|
||||
import warnings
|
||||
import weakref
|
||||
from functools import wraps
|
||||
|
||||
from torch.ao.pruning._experimental.data_sparsifier import BaseDataSparsifier
|
||||
|
||||
|
||||
__all__ = ["BaseDataScheduler"]
|
||||
|
||||
|
||||
class BaseDataScheduler:
|
||||
r"""
|
||||
The BaseDataScheduler is the abstract scheduler class specifically for the
|
||||
BaseDataSparsifier class. This class controls a specific hyperparameter of
|
||||
the sparsifier class and varies it across the training process (or across time).
|
||||
|
||||
Args:
|
||||
data_sparsifier (instance of BaseDataSparsifier)
|
||||
Implemented class data sparsifier class wherein the update_mask is implemented
|
||||
schedule_param (str)
|
||||
A specific hyperparameter of the passed sparsifier that needs to be scheduled/varied
|
||||
last_epoch (int, default=-1)
|
||||
This is specifically is passed when training needs to be resumed from a particular
|
||||
point.
|
||||
verbose (bool, default=False)
|
||||
Verbosity of the BaseDataScheduler
|
||||
|
||||
The *get_hyperparam()* function needs to be implemented by the user.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, data_sparsifier, schedule_param: str, last_epoch=-1, verbose=False
|
||||
):
|
||||
# Attach sparsifier
|
||||
if not isinstance(data_sparsifier, BaseDataSparsifier):
|
||||
raise TypeError(
|
||||
f"{type(data_sparsifier).__name__} is not an instance of torch.ao.pruning.BaseDataSparsifier"
|
||||
)
|
||||
self.data_sparsifier = data_sparsifier
|
||||
self.schedule_param = schedule_param
|
||||
|
||||
# Initialize epoch and base hyper-params
|
||||
self.base_param = {
|
||||
name: config.get(schedule_param, None)
|
||||
for name, config in self.data_sparsifier.data_groups.items()
|
||||
}
|
||||
|
||||
self.last_epoch = last_epoch
|
||||
|
||||
# Following https://github.com/pytorch/pytorch/issues/20124
|
||||
# We would like to ensure that `scheduler.step()` is called after
|
||||
# `sparsifier.step()`
|
||||
def with_counter(method):
|
||||
if getattr(method, "_with_counter", False):
|
||||
# `sparsifier.step()` has already been replaced, return.
|
||||
return method
|
||||
|
||||
# Keep a weak reference to the sparsifier instance to prevent
|
||||
# cyclic references.
|
||||
instance_ref = weakref.ref(method.__self__)
|
||||
# Get the unbound method for the same purpose.
|
||||
func = method.__func__
|
||||
cls = instance_ref().__class__
|
||||
del method
|
||||
|
||||
@wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
instance = instance_ref()
|
||||
instance._step_count += 1 # type: ignore[union-attr]
|
||||
wrapped = func.__get__(instance, cls)
|
||||
return wrapped(*args, **kwargs)
|
||||
|
||||
# Note that the returned function here is no longer a bound method,
|
||||
# so attributes like `__func__` and `__self__` no longer exist.
|
||||
wrapper._with_counter = True # type: ignore[attr-defined]
|
||||
return wrapper
|
||||
|
||||
self.data_sparsifier.step = with_counter(self.data_sparsifier.step) # type: ignore[assignment]
|
||||
self.data_sparsifier._step_count = 0 # type: ignore[attr-defined]
|
||||
self._step_count: int = 0
|
||||
self.verbose = verbose
|
||||
|
||||
# Housekeeping
|
||||
self._get_sp_called_within_step: bool = False # sp -> schedule parameter
|
||||
self.step()
|
||||
|
||||
@abc.abstractmethod
|
||||
def get_schedule_param(self):
|
||||
r"""
|
||||
Abstract method that needs to be implemented by the child class.
|
||||
The expected return type should is a dictionary of name to schedule_param value
|
||||
The returned values will be updated in sparsifier when the scheduler step() function
|
||||
is called.
|
||||
|
||||
Example:
|
||||
>>> def get_schedule_param(self):
|
||||
... new_param = {}
|
||||
... for name in self.sparsifier.data_groups.keys():
|
||||
... new_param[name] = (
|
||||
... self.sparsifier.data_groups[name][self.schedule_param] * 0.5
|
||||
... )
|
||||
... return new_param
|
||||
|
||||
When the step() function is called, the value in self.sparsifier.data_groups[name][self.schedule_param]
|
||||
would be halved
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def __repr__(self):
|
||||
format_string = self.__class__.__name__ + " ("
|
||||
format_string += "\n"
|
||||
format_string += f"Data Sparsifier {self.data_sparsifier}\n"
|
||||
format_string += f" {self.schedule_param}: {self.base_param}\n"
|
||||
format_string += ")"
|
||||
return format_string
|
||||
|
||||
def state_dict(self):
|
||||
"""Returns the state of the scheduler as a :class:`dict`.
|
||||
|
||||
It contains an entry for every variable in self.__dict__ which
|
||||
is not the sparsifier.
|
||||
|
||||
Note:
|
||||
The scheduler class does not track the state of the data_sparsifier.
|
||||
Make sure to store the state of the sparsifier before storing the
|
||||
state of the scheduler
|
||||
"""
|
||||
return {
|
||||
key: value
|
||||
for key, value in self.__dict__.items()
|
||||
if key != "data_sparsifier"
|
||||
}
|
||||
|
||||
def load_state_dict(self, state_dict):
|
||||
"""Loads the schedulers state.
|
||||
|
||||
Note:
|
||||
Remember to restore the state of the data_sparsifier before the scheduler.
|
||||
|
||||
Args:
|
||||
state_dict (dict): scheduler state. Should be an object returned
|
||||
from a call to :meth:`state_dict`.
|
||||
"""
|
||||
self.__dict__.update(state_dict)
|
||||
|
||||
def get_last_param(self):
|
||||
return self._last_param
|
||||
|
||||
def step(self):
|
||||
# Raise warning if trying to call scheduler step before the sparsifier.
|
||||
# https://github.com/pytorch/pytorch/issues/20124
|
||||
if self._step_count == 1:
|
||||
if not hasattr(self.data_sparsifier.step, "_with_counter"):
|
||||
warnings.warn(
|
||||
"Seems like `data_sparsifier.step()` has been overridden after sparsity scheduler "
|
||||
"initialization. Please, make sure to call `data_sparsifier.step()` before "
|
||||
"`scheduler.step()`.",
|
||||
UserWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
# Just check if there were two first scheduler.step() calls before sparsifier.step()
|
||||
elif self.data_sparsifier._step_count < 1: # type: ignore[attr-defined]
|
||||
warnings.warn(
|
||||
"Detected call of `scheduler.step()` before `data_sparsifier.step()`. "
|
||||
"You have to make sure you run the data_sparsifier.step() BEFORE any "
|
||||
"calls to the scheduler.step().",
|
||||
UserWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
self._step_count += 1
|
||||
|
||||
class _enable_get_sp_call:
|
||||
def __init__(self, o):
|
||||
self.o = o
|
||||
|
||||
def __enter__(self):
|
||||
self.o._get_sp_called_within_step = True
|
||||
return self
|
||||
|
||||
def __exit__(self, type, value, traceback):
|
||||
self.o._get_sp_called_within_step = False
|
||||
|
||||
with _enable_get_sp_call(self):
|
||||
self.last_epoch += 1
|
||||
updated_scheduler_params = self.get_schedule_param()
|
||||
|
||||
for name, param in updated_scheduler_params.items():
|
||||
self.data_sparsifier.data_groups[name][self.schedule_param] = param
|
||||
if self.verbose:
|
||||
print(f"Adjusting {self.schedule_param} for group {name} to {param}")
|
||||
|
||||
self._last_param = {
|
||||
name: config.get(self.schedule_param, None)
|
||||
for name, config in self.data_sparsifier.data_groups.items()
|
||||
}
|
||||
self.data_sparsifier.enable_mask_update = True
|
||||
+8
@@ -0,0 +1,8 @@
|
||||
from .base_data_sparsifier import BaseDataSparsifier
|
||||
from .data_norm_sparsifier import DataNormSparsifier
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BaseDataSparsifier",
|
||||
"DataNormSparsifier",
|
||||
]
|
||||
+334
@@ -0,0 +1,334 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import abc
|
||||
import copy
|
||||
import sys
|
||||
import warnings
|
||||
from collections import defaultdict
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.ao.pruning.sparsifier import base_sparsifier, utils
|
||||
from torch.nn.utils import parametrize
|
||||
|
||||
|
||||
if not sys.warnoptions:
|
||||
# to suppress repeated warnings when being used in a training loop.
|
||||
warnings.simplefilter("once")
|
||||
|
||||
__all__ = ["BaseDataSparsifier"]
|
||||
|
||||
EMBEDDING_TYPES = {
|
||||
nn.Embedding,
|
||||
nn.EmbeddingBag,
|
||||
}
|
||||
|
||||
SUPPORTED_TYPES = {
|
||||
torch.Tensor,
|
||||
nn.Parameter,
|
||||
*EMBEDDING_TYPES,
|
||||
}
|
||||
|
||||
|
||||
class _Container(nn.Module):
|
||||
pass
|
||||
|
||||
|
||||
class BaseDataSparsifier(base_sparsifier.BaseSparsifier):
|
||||
r"""
|
||||
Base Data Sparsifier class for all Data sparsifiers.
|
||||
The abstract class accepts raw torch tensors / embedding / embedding bags (refer to SUPPORTED_TYPES above)
|
||||
to prepare for sparsification.
|
||||
In this case, mask (and parametrizations) is owned by the class and not by the user.
|
||||
Specifically, the container object inside the class maintains the mask and parametrizations of the input data
|
||||
|
||||
Args:
|
||||
data_list (list of tuples)
|
||||
list of (name, data) tuples to sparsify. Lookup SUPPORTED_TYPES
|
||||
for type of data. Internally, a container module handles the data sparsification.
|
||||
|
||||
defaults (dict)
|
||||
default configurations will be attached to the
|
||||
configuration. Only the keys that don't exist in the `config` will
|
||||
be updated.
|
||||
Example::
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> data_list = [('tensor_1', torch.randn(3,3)), ('tensor_2', torch.randn(4,4))]
|
||||
>>> defaults = {'sparsity_level': 0.7}
|
||||
>>> sparsifier = DerivedDataSparsifier(data_list = data_list, **defaults) # Some sparsifier that inherits BaseDataSparsifier
|
||||
>>> new_tensor_to_add = {'name': 'tensor_3', 'data': torch.randn(5,5), 'sparsity_level': 0.3}
|
||||
>>> sparsifier.add_data(**new_tensor_to_add)
|
||||
>>> # tensor_1 and tensor_2 will have sparsity_level of 0.7 but tensor_3 will have sparsity_level=0.3
|
||||
"""
|
||||
|
||||
def __init__(self, data_list: list[tuple[str, Any]] | None = None, **defaults):
|
||||
super().__init__(defaults=defaults)
|
||||
|
||||
self._container = _Container()
|
||||
|
||||
self.data_groups: dict[str, dict] = defaultdict(dict) # name -> {**config}
|
||||
if data_list is not None:
|
||||
# add data with default config here
|
||||
[self.add_data(name, data, **self.defaults) for name, data in data_list]
|
||||
|
||||
def prepare(self, model, config):
|
||||
raise NotImplementedError("this function is undefined for this class")
|
||||
|
||||
def _extract_weight(self, data):
|
||||
# extract the weight parameter instead of underlying data
|
||||
if type(data) in [torch.Tensor, nn.Parameter]:
|
||||
return data
|
||||
elif type(data) in EMBEDDING_TYPES:
|
||||
return data.weight
|
||||
|
||||
def add_data(self, name: str, data, reuse_mask=True, **config):
|
||||
r"""Configures and parametrizes the internal container model with name and data.
|
||||
|
||||
**Note**:
|
||||
1. If the data with name already exists, it replaces the data.
|
||||
2. While replacing, the old mask is reused when `reuse_mask=True`
|
||||
3. If `reuse_mask=True`, then the replacing data needs to have the same shape as that of old data.
|
||||
4. By default, the config of the replaced data is used as config for the replacing data, unless something
|
||||
is specified in the config dictionary.
|
||||
"""
|
||||
if type(data) not in SUPPORTED_TYPES:
|
||||
raise AssertionError(
|
||||
f"specified data type:{type(data)} not supported at the moment"
|
||||
)
|
||||
local_args = copy.deepcopy(self.defaults)
|
||||
local_args.update(config)
|
||||
weight = self._extract_weight(data)
|
||||
|
||||
# Bookkeeping in the container class
|
||||
mask = local_args.get("mask", torch.ones_like(weight))
|
||||
param_class = local_args.get("parametrization", utils.FakeSparsity)
|
||||
|
||||
if name in self.state:
|
||||
# If the named data already exists - replace
|
||||
warnings.warn(
|
||||
"Replacing existing data of the same name. - Did you mean a different name?",
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
# reuse old config
|
||||
old_args = self.data_groups[name]
|
||||
local_args = copy.deepcopy(old_args)
|
||||
local_args.update(config)
|
||||
|
||||
if reuse_mask:
|
||||
current_data = self.get_data(name=name)
|
||||
if weight.shape != current_data.shape:
|
||||
raise AssertionError(
|
||||
"to retain the old mask, the shape of the new data must be the same as the previous one"
|
||||
)
|
||||
mask = self.get_mask(
|
||||
name=name
|
||||
) # reuse mask instead of creating a new one
|
||||
|
||||
self._delete_data(name=name)
|
||||
|
||||
# parameter creates a deepcopy of the weight inside, so create a buffer
|
||||
self._container.register_buffer(name=name, tensor=weight)
|
||||
parametrize.register_parametrization(self._container, name, param_class(mask))
|
||||
self.state[name]["mask"] = mask
|
||||
self.data_groups[name] = local_args
|
||||
return getattr(self._container, name)
|
||||
|
||||
def get_data(self, name: str, return_original: bool = True):
|
||||
r"""Returns weight tensor (or data)
|
||||
Args:
|
||||
- name: name of the data to be returned
|
||||
- return_original returns weight tensor without applying parametrization if True
|
||||
else - returns the sparsified version (parametrized)
|
||||
"""
|
||||
if name not in self.data_groups:
|
||||
raise ValueError("data with specified name does not exist")
|
||||
|
||||
if return_original:
|
||||
if not parametrize.is_parametrized(self._container, name):
|
||||
raise ValueError("mask squashed - original mask value does not exist")
|
||||
data = getattr(self._container.parametrizations, name).original
|
||||
return data
|
||||
else:
|
||||
return getattr(self._container, name)
|
||||
|
||||
def _convert_mask(self, states, sparse_coo=True):
|
||||
r"""Converts the mask to sparse coo or dense tensors depending on the `sparse_coo` argument."""
|
||||
states = copy.deepcopy(states)
|
||||
for state in states.values():
|
||||
if sparse_coo:
|
||||
state["mask"] = state["mask"].to_sparse_coo()
|
||||
else:
|
||||
state["mask"] = state["mask"].to_dense()
|
||||
|
||||
return states
|
||||
|
||||
def state_dict(self):
|
||||
r"""Returns the state of the optimizer as a :class:`dict`.
|
||||
|
||||
It contains:
|
||||
* state - contains name -> mask mapping.
|
||||
* data_groups - a list containing all sparsity configuration groups
|
||||
with the key name specifying the name of the data
|
||||
* container_state_dict - the state dictionary of the internal
|
||||
container model used for sparsification
|
||||
"""
|
||||
state = self._convert_mask(self.state)
|
||||
return {
|
||||
"state": state,
|
||||
"data_groups": self.data_groups,
|
||||
"_container": self._container.state_dict(),
|
||||
}
|
||||
|
||||
def _load_container_from_state(self, states, data_groups, container_state_dict):
|
||||
r"""This restores the state of the container specifically based on the data present in state and data_groups
|
||||
If the data was parametrized, then the data would be added to the container and then parametrized,
|
||||
else it would just add the attribute the container.
|
||||
"""
|
||||
for name, state in states.items():
|
||||
config_name = data_groups.get(name, None)
|
||||
if config_name is None:
|
||||
raise RuntimeError(f"Error loading {name}")
|
||||
|
||||
# check if the data with such a name was parametrized, if so parametrize
|
||||
# otherwise just set the attribute and continue
|
||||
parametrized_name = f"parametrizations.{name}.original"
|
||||
parametrized = False
|
||||
data = container_state_dict.get(name, None)
|
||||
if name in container_state_dict:
|
||||
# the parametrization was probably removed for this
|
||||
data = container_state_dict.get(name)
|
||||
|
||||
elif parametrized_name in container_state_dict:
|
||||
# so the weight was parametrized
|
||||
data = container_state_dict.get(parametrized_name)
|
||||
parametrized = True
|
||||
|
||||
else:
|
||||
raise RuntimeError(f"Error loading {name}")
|
||||
|
||||
self._container.register_buffer(name=name, tensor=data)
|
||||
|
||||
if parametrized:
|
||||
# register parameter if parametrized
|
||||
mask = state.get("mask", torch.ones_like(data))
|
||||
param_class = data_groups.get(
|
||||
"parametrization", utils.FakeSparsity
|
||||
) # change once public_api for utils is fixed!
|
||||
parametrize.register_parametrization(
|
||||
self._container, name, param_class(mask)
|
||||
)
|
||||
|
||||
def load_state_dict(self, state_dict, strict=True):
|
||||
r"""The load_state_dict() restores the state of the sparsifier based on the state_dict
|
||||
|
||||
Args:
|
||||
* state_dict - the dictionary that to which the current sparsifier needs to be restored to
|
||||
* strict - If True - the sparsifier is reset and is restored exactly to the state in state_dict.
|
||||
If False - the current sparsifier is not reset before loading the state_dict i.e. data added
|
||||
before loading the state_dict is not erased.
|
||||
"""
|
||||
states = copy.deepcopy(state_dict["state"])
|
||||
data_groups = copy.deepcopy(state_dict["data_groups"])
|
||||
container_state_dict = copy.deepcopy(state_dict["_container"])
|
||||
|
||||
states = self._convert_mask(
|
||||
states, sparse_coo=False
|
||||
) # convert sparse coo mask to dense
|
||||
if strict:
|
||||
# if strict load -> then reset container
|
||||
self._container = _Container()
|
||||
|
||||
self._load_container_from_state(states, data_groups, container_state_dict)
|
||||
|
||||
if not strict:
|
||||
states.update(self.state)
|
||||
data_groups.update(self.data_groups)
|
||||
|
||||
self.__setstate__({"state": states, "data_groups": data_groups})
|
||||
|
||||
def __setstate__(self, state):
|
||||
if "_container" in state: # If container object is in state then load model
|
||||
container_dict = state.pop("_container")
|
||||
self._container = _Container()
|
||||
state["state"] = self._convert_mask(
|
||||
state["state"], sparse_coo=False
|
||||
) # convert sparse coo mask to dense
|
||||
self._load_container_from_state(
|
||||
state["state"], state["data_groups"], container_dict
|
||||
)
|
||||
|
||||
self.__dict__.update(state)
|
||||
|
||||
def __getstate__(self):
|
||||
state = self._convert_mask(self.state)
|
||||
return {
|
||||
"defaults": self.defaults,
|
||||
"state": state,
|
||||
"data_groups": self.data_groups,
|
||||
"_container": self._container.state_dict(),
|
||||
}
|
||||
|
||||
def __repr__(self): # type:ignore[override]
|
||||
format_string = self.__class__.__name__ + " ("
|
||||
for name, sparse_args in self.data_groups.items():
|
||||
format_string += "\n"
|
||||
format_string += "\tData Group\n"
|
||||
format_string += f"\t name: {name}\n"
|
||||
for key in sorted(sparse_args.keys()):
|
||||
if key == "data":
|
||||
continue
|
||||
format_string += f"\t {key}: {sparse_args[key]}\n"
|
||||
format_string += ")"
|
||||
return format_string
|
||||
|
||||
def get_mask(self, name: str):
|
||||
if name not in self.state:
|
||||
raise ValueError("data with specified name does not exist")
|
||||
return self.state[name]["mask"]
|
||||
|
||||
def squash_mask(self, *args, leave_parametrized=True, names=None, **kwargs):
|
||||
r"""Squashes the sparse masks into the appropriate tensors. Also, accepts list of strings
|
||||
to squash mask for. If none, squashes mask for all the keys
|
||||
kwargs:
|
||||
* names: list of strings to squash mask for
|
||||
* sparsified: if true - applies the mask before squashing
|
||||
if false - does not apply the mask before squashing
|
||||
"""
|
||||
if names is None:
|
||||
names = list(self.data_groups.keys())
|
||||
for name in names:
|
||||
parametrize.remove_parametrizations(
|
||||
self._container, name, leave_parametrized=leave_parametrized
|
||||
)
|
||||
|
||||
def step(self): # type:ignore[override]
|
||||
if not self.enable_mask_update:
|
||||
return
|
||||
with torch.no_grad():
|
||||
for name, config in self.data_groups.items():
|
||||
# get non-sparsified data
|
||||
data = self.get_data(name)
|
||||
# need name for the mask otherwise can directly pass mask?
|
||||
self.update_mask(name, data, **config)
|
||||
|
||||
@abc.abstractmethod
|
||||
def update_mask(self, name, data, **kwargs): # type: ignore[override]
|
||||
pass
|
||||
|
||||
def _delete_data(self, name):
|
||||
"""Detaches some data from the sparsifier.
|
||||
|
||||
Args:
|
||||
name (str)
|
||||
Name of the data to be removed from the sparsifier
|
||||
|
||||
Note:
|
||||
Currently private. Kind of used as a helper function when replacing data of the same name
|
||||
"""
|
||||
self.squash_mask(
|
||||
names=[name], leave_parametrized=False
|
||||
) # do not apply the mask while deleting
|
||||
delattr(self._container, name)
|
||||
self.state.pop(name)
|
||||
self.data_groups.pop(name)
|
||||
+204
@@ -0,0 +1,204 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import operator
|
||||
from functools import reduce
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch.nn import functional as F
|
||||
|
||||
from .base_data_sparsifier import BaseDataSparsifier
|
||||
|
||||
|
||||
__all__ = ["DataNormSparsifier"]
|
||||
|
||||
|
||||
class DataNormSparsifier(BaseDataSparsifier):
|
||||
r"""L1-Norm Sparsifier
|
||||
This sparsifier computes the *L1-norm* of every sparse block and "zeroes-out" the
|
||||
ones with the lowest norm. The level of sparsity defines how many of the
|
||||
blocks is removed.
|
||||
This sparsifier is controlled by three variables:
|
||||
1. `sparsity_level` defines the number of *sparse blocks* that are zeroed-out
|
||||
2. `sparse_block_shape` defines the shape of the sparse blocks. Note that
|
||||
the sparse blocks originate at the zero-index of the tensor.
|
||||
3. `zeros_per_block` is the number of zeros that we are expecting in each
|
||||
sparse block. By default we assume that all elements within a block are
|
||||
zeroed-out. However, setting this variable sets the target number of
|
||||
zeros per block. The zeros within each block are chosen as the *smallest
|
||||
absolute values*.
|
||||
Args:
|
||||
sparsity_level: The target level of sparsity
|
||||
sparse_block_shape: The shape of a sparse block
|
||||
zeros_per_block: Number of zeros in a sparse block
|
||||
Note::
|
||||
All arguments to the DataNormSparsifier constructor are "default"
|
||||
arguments and could be overridden by the configuration provided in the
|
||||
`add_data` step.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
data_list: list[tuple[str, Any]] | None = None,
|
||||
sparsity_level: float = 0.5,
|
||||
sparse_block_shape: tuple[int, int] = (1, 4),
|
||||
zeros_per_block: int | None = None,
|
||||
norm: str = "L1",
|
||||
):
|
||||
if zeros_per_block is None:
|
||||
zeros_per_block = reduce(operator.mul, sparse_block_shape)
|
||||
|
||||
if norm not in ["L1", "L2"]:
|
||||
raise AssertionError("only L1 and L2 norm supported at the moment")
|
||||
|
||||
defaults = {
|
||||
"sparsity_level": sparsity_level,
|
||||
"sparse_block_shape": sparse_block_shape,
|
||||
"zeros_per_block": zeros_per_block,
|
||||
}
|
||||
self.norm = norm
|
||||
super().__init__(data_list=data_list, **defaults)
|
||||
|
||||
def __get_scatter_folded_mask(
|
||||
self, data, dim, indices, output_size, sparse_block_shape
|
||||
):
|
||||
mask = torch.ones_like(data)
|
||||
mask.scatter_(dim=dim, index=indices, value=0) # zeroing out
|
||||
mask = F.fold(
|
||||
mask,
|
||||
output_size=output_size,
|
||||
kernel_size=sparse_block_shape,
|
||||
stride=sparse_block_shape,
|
||||
)
|
||||
mask = mask.to(torch.int8)
|
||||
return mask
|
||||
|
||||
def __get_block_level_mask(self, data, sparse_block_shape, zeros_per_block):
|
||||
# Assume data is a squeezed tensor
|
||||
height, width = data.shape[-2], data.shape[-1]
|
||||
block_height, block_width = sparse_block_shape
|
||||
values_per_block = block_height * block_width
|
||||
|
||||
# just return zeros if zeroing all elements in block
|
||||
if values_per_block == zeros_per_block:
|
||||
return torch.zeros_like(data, dtype=torch.int8)
|
||||
|
||||
# creating additional height and width to support padding
|
||||
dh = (block_height - height % block_height) % block_height
|
||||
dw = (block_width - width % block_width) % block_width
|
||||
|
||||
# create a new padded tensor like data (to match the block_shape)
|
||||
padded_data = torch.ones(
|
||||
height + dh, width + dw, dtype=data.dtype, device=data.device
|
||||
)
|
||||
padded_data = (
|
||||
padded_data * torch.nan
|
||||
) # can also be replaced with 0 to stop the removal of edge data
|
||||
padded_data[0:height, 0:width] = data
|
||||
unfolded_data = F.unfold(
|
||||
padded_data[None, None, :],
|
||||
kernel_size=sparse_block_shape,
|
||||
stride=sparse_block_shape,
|
||||
)
|
||||
|
||||
_, sorted_idx = torch.sort(unfolded_data, dim=1)
|
||||
sorted_idx = sorted_idx[
|
||||
:, :zeros_per_block, :
|
||||
] # zero out zeros_per_block number of elements
|
||||
|
||||
mask = self.__get_scatter_folded_mask(
|
||||
data=unfolded_data,
|
||||
dim=1,
|
||||
indices=sorted_idx,
|
||||
output_size=padded_data.shape,
|
||||
sparse_block_shape=sparse_block_shape,
|
||||
)
|
||||
|
||||
mask = (
|
||||
mask.squeeze(0).squeeze(0)[:height, :width].contiguous()
|
||||
) # remove padding and make contiguous
|
||||
return mask
|
||||
|
||||
def __get_data_level_mask(self, data, sparsity_level, sparse_block_shape):
|
||||
height, width = data.shape[-2], data.shape[-1]
|
||||
block_height, block_width = sparse_block_shape
|
||||
dh = (block_height - height % block_height) % block_height
|
||||
dw = (block_width - width % block_width) % block_width
|
||||
|
||||
data_norm = F.avg_pool2d(
|
||||
data[None, None, :],
|
||||
kernel_size=sparse_block_shape,
|
||||
stride=sparse_block_shape,
|
||||
ceil_mode=True,
|
||||
)
|
||||
|
||||
values_per_block = reduce(operator.mul, sparse_block_shape)
|
||||
|
||||
data_norm = data_norm.flatten()
|
||||
num_blocks = len(data_norm)
|
||||
|
||||
data_norm = data_norm.repeat(
|
||||
1, values_per_block, 1
|
||||
) # get similar shape after unfold
|
||||
_, sorted_idx = torch.sort(data_norm, dim=2)
|
||||
|
||||
threshold_idx = round(sparsity_level * num_blocks) # number of blocks to remove
|
||||
sorted_idx = sorted_idx[:, :, :threshold_idx]
|
||||
|
||||
mask = self.__get_scatter_folded_mask(
|
||||
data=data_norm,
|
||||
dim=2,
|
||||
indices=sorted_idx,
|
||||
output_size=(height + dh, width + dw),
|
||||
sparse_block_shape=sparse_block_shape,
|
||||
)
|
||||
|
||||
mask = mask.squeeze(0).squeeze(0)[
|
||||
:height, :width
|
||||
] # squeeze only the first 2 dimension
|
||||
return mask
|
||||
|
||||
def update_mask( # type: ignore[override]
|
||||
self, name, data, sparsity_level, sparse_block_shape, zeros_per_block, **kwargs
|
||||
):
|
||||
values_per_block = reduce(operator.mul, sparse_block_shape)
|
||||
if zeros_per_block > values_per_block:
|
||||
raise ValueError(
|
||||
"Number of zeros per block cannot be more than "
|
||||
"the total number of elements in that block."
|
||||
)
|
||||
if zeros_per_block < 0:
|
||||
raise ValueError("Number of zeros per block should be positive.")
|
||||
|
||||
if self.norm == "L1":
|
||||
data_norm = torch.abs(data).squeeze() # absolute value based (L1)
|
||||
else:
|
||||
data_norm = (data * data).squeeze() # square every element for L2
|
||||
|
||||
if len(data_norm.shape) > 2: # only supports 2 dimensional data at the moment
|
||||
raise ValueError("only supports 2-D at the moment")
|
||||
|
||||
elif len(data_norm.shape) == 1: # in case the data is bias (or 1D)
|
||||
data_norm = data_norm[None, :]
|
||||
|
||||
mask = self.get_mask(name)
|
||||
if sparsity_level <= 0 or zeros_per_block == 0:
|
||||
mask.data = torch.ones_like(mask)
|
||||
elif sparsity_level >= 1.0 and (zeros_per_block == values_per_block):
|
||||
mask.data = torch.zeros_like(mask)
|
||||
|
||||
# Fetch the high level mask that zeros out entire blocks
|
||||
data_lvl_mask = self.__get_data_level_mask(
|
||||
data=data_norm,
|
||||
sparsity_level=sparsity_level,
|
||||
sparse_block_shape=sparse_block_shape,
|
||||
)
|
||||
|
||||
# Fetch block level mask that zeros out 'zeros_per_block' number of elements in every block
|
||||
block_lvl_mask = self.__get_block_level_mask(
|
||||
data=data_norm,
|
||||
sparse_block_shape=sparse_block_shape,
|
||||
zeros_per_block=zeros_per_block,
|
||||
)
|
||||
|
||||
# zero out the entries inside those blocks whose block is sparsified
|
||||
mask.data = torch.where(data_lvl_mask == 1, data_lvl_mask, block_lvl_mask)
|
||||
+44
@@ -0,0 +1,44 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import logging
|
||||
|
||||
from torch.ao.pruning._experimental.data_sparsifier.base_data_sparsifier import (
|
||||
SUPPORTED_TYPES,
|
||||
)
|
||||
|
||||
|
||||
logger: logging.Logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _attach_model_to_data_sparsifier(module, data_sparsifier, config=None):
|
||||
"""Attaches a data sparsifier to all the layers of the module.
|
||||
Essentially, loop over all the weight parameters in the module and
|
||||
attach it to the data sparsifier.
|
||||
Note::
|
||||
The '.' in the layer names are replaced with '_' (refer to _get_valid_name() below)
|
||||
before attaching to the sparsifier. This is because, the data
|
||||
sparsifier uses a dummy model inside to store the weight parameters.
|
||||
"""
|
||||
if config is None:
|
||||
config = {}
|
||||
for name, parameter in module.named_parameters():
|
||||
if type(parameter) in SUPPORTED_TYPES:
|
||||
valid_name = _get_valid_name(name)
|
||||
# will be defaulted to default configs
|
||||
data_sparsifier.add_data(
|
||||
name=valid_name, data=parameter, **config.get(valid_name, {})
|
||||
)
|
||||
|
||||
|
||||
def _get_valid_name(name):
|
||||
return name.replace(".", "_") # . is not allowed as a name
|
||||
|
||||
|
||||
def _log_sparsified_level(model, data_sparsifier) -> None:
|
||||
# Show the level of sparsity AFTER step:
|
||||
for name, parameter in model.named_parameters():
|
||||
if type(parameter) not in SUPPORTED_TYPES:
|
||||
continue
|
||||
valid_name = _get_valid_name(name)
|
||||
mask = data_sparsifier.get_mask(name=valid_name)
|
||||
sparsity_level = 1.0 - mask.float().mean()
|
||||
logger.info("Sparsity in layer %s = % .2%", name, sparsity_level)
|
||||
+181
@@ -0,0 +1,181 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from collections import defaultdict
|
||||
from copy import deepcopy
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
import pytorch_lightning as pl # type: ignore[import]
|
||||
|
||||
from ._data_sparstity_utils import (
|
||||
_attach_model_to_data_sparsifier,
|
||||
_get_valid_name,
|
||||
_log_sparsified_level,
|
||||
)
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import torch
|
||||
|
||||
|
||||
class PostTrainingDataSparsity(pl.callbacks.Callback):
|
||||
"""Lightning callback that enables post-training sparsity.
|
||||
|
||||
This callback aims to sparsify the model inside lightning module after training.
|
||||
**Note that the model is copied and then sparsified, so the existing model is not modified**
|
||||
|
||||
The sparsified model can be used for comparison and can be accessed using
|
||||
<callback_obj>.sparsified
|
||||
|
||||
Args:
|
||||
data_sparsifier_class (some implemented class of BaseDataSparsifier)
|
||||
The data sparsifier object of this class is created when the
|
||||
training starts.
|
||||
Note: Objects should not be passed in here as they are created
|
||||
once the training completes.
|
||||
|
||||
data_sparsifier_args (Dict)
|
||||
Dictionary of args to be passed to the data sparsifier.
|
||||
Note: data_list arg should be ignored
|
||||
|
||||
Hooks implemented:
|
||||
on_fit_end()
|
||||
1. copies the model and attaches it to the sparsifier
|
||||
2. sparsier step() is called
|
||||
3. squashes the mask()
|
||||
"""
|
||||
|
||||
def __init__(self, data_sparsifier_class, data_sparsifier_args):
|
||||
super().__init__()
|
||||
self.data_sparsifier_class = data_sparsifier_class
|
||||
self.data_sparsifier_args = data_sparsifier_args
|
||||
self.data_sparsifier: Any = None
|
||||
self.sparsified: torch.nn.Module | None = None
|
||||
|
||||
def on_fit_end(self, trainer, pl_module) -> None:
|
||||
self.sparsified = deepcopy(pl_module.model).eval()
|
||||
self.data_sparsifier = self.data_sparsifier_class(**self.data_sparsifier_args)
|
||||
|
||||
_attach_model_to_data_sparsifier(self.sparsified, self.data_sparsifier)
|
||||
|
||||
self.data_sparsifier.step()
|
||||
|
||||
self.data_sparsifier.squash_mask() # currently squashes params for all mask
|
||||
|
||||
_log_sparsified_level(self.sparsified, self.data_sparsifier)
|
||||
|
||||
|
||||
class TrainingAwareDataSparsity(pl.callbacks.Callback):
|
||||
"""Lightning callback that enables in-training sparsity.
|
||||
|
||||
This callback aims to sparsify the model inside lightning module during training.
|
||||
**Note that the model is copied and then sparsified, so the existing model is not modified**
|
||||
|
||||
The sparsified model can be used for comparison and can be accessed using
|
||||
<callback_obj>.sparsified
|
||||
|
||||
Args:
|
||||
data_sparsifier_class (some implemented class of BaseDataSparsifier)
|
||||
The data sparsifier object of this class is created when the
|
||||
training starts.
|
||||
Note: Objects should not be passed in here as they are created
|
||||
when the training starts.
|
||||
|
||||
data_sparsifier_args (Dict)
|
||||
Dictionary of args to be passed to the data sparsifier.
|
||||
Note: data_list arg should be ignored
|
||||
|
||||
data_scheduler_class (some implemented class of BaseDataScheduler)
|
||||
The data scheduler of this class is created when the training starts
|
||||
Note: Objects should not be passed in here as they are created
|
||||
when the training starts.
|
||||
|
||||
data_scheduler_args(Dict)
|
||||
Dictionary of args to be passed to the data scheduler.
|
||||
**Note: data_sparsifier arg should be ignored as the recipe
|
||||
creates and pass sparsifier object into the class**
|
||||
|
||||
Hooks implemented:
|
||||
on_train_start()
|
||||
Data sparsifier and scheduler objects are created.
|
||||
Pytorch model attached to the sparsifier
|
||||
|
||||
on_train_epoch_start()
|
||||
Loads the state_dict of the data sparsifier
|
||||
|
||||
on_train_epoch_end()
|
||||
1. Copies the model and attaches it to the sparsifier
|
||||
2. sparsifier step() and scheduler step()
|
||||
3. Dump state_dict of the current sparsifier
|
||||
|
||||
on_train_end()
|
||||
squash mask
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
data_sparsifier_class,
|
||||
data_sparsifier_args,
|
||||
data_scheduler_class,
|
||||
data_scheduler_args,
|
||||
):
|
||||
super().__init__()
|
||||
# data sparsifier objects
|
||||
self.data_sparsifier_class = data_sparsifier_class
|
||||
self.data_sparsifier_args = data_sparsifier_args
|
||||
|
||||
# scheduler objects
|
||||
self.data_scheduler_class = data_scheduler_class
|
||||
self.data_scheduler_args = data_scheduler_args
|
||||
|
||||
# fields
|
||||
self.data_sparsifier: Any = None
|
||||
self.data_scheduler: Any = None
|
||||
self.sparsified: torch.nn.Module | None = None
|
||||
|
||||
self.data_sparsifier_state_dict: Any = None
|
||||
|
||||
def on_train_start(self, trainer, pl_module) -> None:
|
||||
# create sparsifier
|
||||
self.data_sparsifier = self.data_sparsifier_class(**self.data_sparsifier_args)
|
||||
self.sparsified = deepcopy(pl_module.model)
|
||||
|
||||
_attach_model_to_data_sparsifier(
|
||||
self.sparsified, self.data_sparsifier
|
||||
) # just to populate the base_sl in the scheduler
|
||||
|
||||
# create scheduler
|
||||
args = deepcopy(self.data_scheduler_args)
|
||||
args["data_sparsifier"] = self.data_sparsifier
|
||||
self.data_scheduler = self.data_scheduler_class(**args)
|
||||
|
||||
def on_train_epoch_start(self, trainer, pl_module):
|
||||
if self.data_sparsifier_state_dict is None:
|
||||
return # probably first epoch
|
||||
|
||||
# load the existing config for each data
|
||||
self.data_sparsifier.load_state_dict(self.data_sparsifier_state_dict)
|
||||
|
||||
def __create_config_based_on_state(self, pl_module):
|
||||
config: dict = defaultdict()
|
||||
if self.data_sparsifier_state_dict is None:
|
||||
return config
|
||||
for name, _ in pl_module.model.named_parameters():
|
||||
valid_name = _get_valid_name(name)
|
||||
config[valid_name] = self.data_sparsifier.data_groups[valid_name]
|
||||
|
||||
return config
|
||||
|
||||
def on_train_epoch_end(self, trainer, pl_module):
|
||||
self.sparsified = deepcopy(pl_module.model)
|
||||
config = self.__create_config_based_on_state(pl_module)
|
||||
|
||||
# attach model to the data sparsifier
|
||||
_attach_model_to_data_sparsifier(
|
||||
self.sparsified, self.data_sparsifier, config=config
|
||||
)
|
||||
self.data_sparsifier.step()
|
||||
self.data_scheduler.step()
|
||||
|
||||
self.data_sparsifier_state_dict = self.data_sparsifier.state_dict()
|
||||
|
||||
def on_train_end(self, trainer, pl_module):
|
||||
self.data_sparsifier.squash_mask()
|
||||
+156
@@ -0,0 +1,156 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.ao.pruning.sparsifier.utils import fqn_to_module, module_to_fqn
|
||||
|
||||
|
||||
SUPPORTED_MODULES = {nn.Embedding, nn.EmbeddingBag}
|
||||
|
||||
|
||||
def _fetch_all_embeddings(model):
|
||||
"""Fetches Embedding and EmbeddingBag modules from the model"""
|
||||
embedding_modules = []
|
||||
stack = [model]
|
||||
while stack:
|
||||
module = stack.pop()
|
||||
for _, child in module.named_children():
|
||||
fqn_name = module_to_fqn(model, child)
|
||||
if type(child) in SUPPORTED_MODULES:
|
||||
embedding_modules.append((fqn_name, child))
|
||||
else:
|
||||
stack.append(child)
|
||||
return embedding_modules
|
||||
|
||||
|
||||
def post_training_sparse_quantize(
|
||||
model,
|
||||
data_sparsifier_class,
|
||||
sparsify_first=True,
|
||||
select_embeddings: list[nn.Module] | None = None,
|
||||
**sparse_config,
|
||||
):
|
||||
"""Takes in a model and applies sparsification and quantization to only embeddings & embeddingbags.
|
||||
The quantization step can happen before or after sparsification depending on the `sparsify_first` argument.
|
||||
|
||||
Args:
|
||||
- model (nn.Module)
|
||||
model whose embeddings needs to be sparsified
|
||||
- data_sparsifier_class (type of data sparsifier)
|
||||
Type of sparsification that needs to be applied to model
|
||||
- sparsify_first (bool)
|
||||
if true, sparsifies first and then quantizes
|
||||
otherwise, quantizes first and then sparsifies.
|
||||
- select_embeddings (List of Embedding modules)
|
||||
List of embedding modules to in the model to be sparsified & quantized.
|
||||
If None, all embedding modules with be sparsified
|
||||
- sparse_config (Dict)
|
||||
config that will be passed to the constructor of data sparsifier object.
|
||||
|
||||
Note:
|
||||
1. When `sparsify_first=False`, quantization occurs first followed by sparsification.
|
||||
- before sparsifying, the embedding layers are dequantized.
|
||||
- scales and zero-points are saved
|
||||
- embedding layers are sparsified and `squash_mask` is applied
|
||||
- embedding weights are requantized using the saved scales and zero-points
|
||||
2. When `sparsify_first=True`, sparsification occurs first followed by quantization.
|
||||
- embeddings are sparsified first
|
||||
- quantization is applied on the sparsified embeddings
|
||||
"""
|
||||
data_sparsifier = data_sparsifier_class(**sparse_config)
|
||||
|
||||
# if select_embeddings is None, perform it on all embeddings
|
||||
if select_embeddings is None:
|
||||
embedding_modules = _fetch_all_embeddings(model)
|
||||
|
||||
else:
|
||||
embedding_modules = []
|
||||
if not isinstance(select_embeddings, list):
|
||||
raise AssertionError(
|
||||
"the embedding_modules must be a list of embedding modules"
|
||||
)
|
||||
for emb in select_embeddings:
|
||||
if type(emb) not in SUPPORTED_MODULES:
|
||||
raise AssertionError(
|
||||
"the embedding_modules list must be an embedding or embedding bags"
|
||||
)
|
||||
fqn_name = module_to_fqn(model, emb)
|
||||
if fqn_name is None:
|
||||
raise AssertionError(
|
||||
"the embedding modules must be part of input model"
|
||||
)
|
||||
embedding_modules.append((fqn_name, emb))
|
||||
|
||||
if sparsify_first:
|
||||
# sparsify
|
||||
for name, emb_module in embedding_modules:
|
||||
valid_name = name.replace(".", "_")
|
||||
data_sparsifier.add_data(name=valid_name, data=emb_module)
|
||||
|
||||
data_sparsifier.step()
|
||||
data_sparsifier.squash_mask()
|
||||
|
||||
# quantize
|
||||
for _, emb_module in embedding_modules:
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
emb_module.qconfig = torch.ao.quantization.float_qparams_weight_only_qconfig
|
||||
|
||||
torch.ao.quantization.prepare(model, inplace=True)
|
||||
torch.ao.quantization.convert(model, inplace=True)
|
||||
|
||||
else:
|
||||
# quantize
|
||||
for _, emb_module in embedding_modules:
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
emb_module.qconfig = torch.ao.quantization.float_qparams_weight_only_qconfig
|
||||
|
||||
torch.ao.quantization.prepare(model, inplace=True)
|
||||
torch.ao.quantization.convert(model, inplace=True)
|
||||
|
||||
# retrieve scale & zero_points
|
||||
quantize_params: dict[str, dict] = {
|
||||
"scales": {},
|
||||
"zero_points": {},
|
||||
"dequant_weights": {},
|
||||
"axis": {},
|
||||
"dtype": {},
|
||||
}
|
||||
|
||||
for name, _ in embedding_modules:
|
||||
quantized_emb = fqn_to_module(model, name)
|
||||
if quantized_emb is None:
|
||||
raise AssertionError(f"quantized embedding {name} not found in model")
|
||||
|
||||
quantized_weight = quantized_emb.weight() # type: ignore[operator]
|
||||
quantize_params["scales"][name] = quantized_weight.q_per_channel_scales()
|
||||
quantize_params["zero_points"][name] = (
|
||||
quantized_weight.q_per_channel_zero_points()
|
||||
)
|
||||
quantize_params["dequant_weights"][name] = torch.dequantize(
|
||||
quantized_weight
|
||||
)
|
||||
quantize_params["axis"][name] = quantized_weight.q_per_channel_axis()
|
||||
quantize_params["dtype"][name] = quantized_weight.dtype
|
||||
|
||||
# attach data to sparsifier
|
||||
data_sparsifier.add_data(
|
||||
name=name.replace(".", "_"),
|
||||
data=quantize_params["dequant_weights"][name],
|
||||
)
|
||||
|
||||
data_sparsifier.step()
|
||||
data_sparsifier.squash_mask()
|
||||
|
||||
for name, _ in embedding_modules:
|
||||
quantized_emb = fqn_to_module(model, name)
|
||||
if quantized_emb is None:
|
||||
raise AssertionError(f"quantized embedding {name} not found in model")
|
||||
requantized_vector = torch.quantize_per_channel(
|
||||
quantize_params["dequant_weights"][name],
|
||||
scales=quantize_params["scales"][name],
|
||||
zero_points=quantize_params["zero_points"][name],
|
||||
dtype=quantize_params["dtype"][name],
|
||||
axis=quantize_params["axis"][name],
|
||||
)
|
||||
|
||||
quantized_emb.set_weight(requantized_vector) # type: ignore[operator]
|
||||
+95
@@ -0,0 +1,95 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
|
||||
from .base_structured_sparsifier import BaseStructuredSparsifier
|
||||
|
||||
|
||||
__all__ = ["FPGMPruner"]
|
||||
|
||||
|
||||
class FPGMPruner(BaseStructuredSparsifier):
|
||||
r"""Filter Pruning via Geometric Median (FPGM) Structured Pruner
|
||||
This sparsifier prune filter (row) in a tensor according to distances among filters according to
|
||||
`Filter Pruning via Geometric Median for Deep Convolutional Neural Networks Acceleration <https://arxiv.org/abs/1811.00250>`_.
|
||||
|
||||
This sparsifier is controlled by three variables:
|
||||
1. `sparsity_level` defines the number of filters (rows) that are zeroed-out.
|
||||
2. `dist` defines the distance measurement type. Default: 3 (L2 distance).
|
||||
Available options are: [1, 2, (custom callable distance function)].
|
||||
|
||||
Note::
|
||||
Inputs should be a 4D convolutional tensor of shape (N, C, H, W).
|
||||
- N: output channels size
|
||||
- C: input channels size
|
||||
- H: height of kernel
|
||||
- W: width of kernel
|
||||
"""
|
||||
|
||||
def __init__(self, sparsity_level: float = 0.5, dist: Callable | int | None = None):
|
||||
defaults = {
|
||||
"sparsity_level": sparsity_level,
|
||||
}
|
||||
|
||||
if dist is None:
|
||||
dist = 2
|
||||
|
||||
if callable(dist):
|
||||
self.dist_fn = dist
|
||||
elif dist == 1:
|
||||
self.dist_fn = lambda x: torch.cdist(x, x, p=1)
|
||||
elif dist == 2:
|
||||
self.dist_fn = lambda x: torch.cdist(x, x, p=2)
|
||||
else:
|
||||
raise NotImplementedError("Distance function is not yet implemented.")
|
||||
super().__init__(defaults=defaults)
|
||||
|
||||
def _compute_distance(self, t):
|
||||
r"""Compute distance across all entries in tensor `t` along all dimension
|
||||
except for the one identified by dim.
|
||||
Args:
|
||||
t (torch.Tensor): tensor representing the parameter to prune
|
||||
Returns:
|
||||
distance (torch.Tensor): distance computed across filtters
|
||||
"""
|
||||
dim = 0 # prune filter (row)
|
||||
|
||||
size = t.size(dim)
|
||||
slc = [slice(None)] * t.dim()
|
||||
|
||||
# flatten the tensor along the dimension
|
||||
t_flatten = [
|
||||
t[tuple(slc[:dim] + [slice(i, i + 1)] + slc[dim + 1 :])].reshape(-1)
|
||||
for i in range(size)
|
||||
]
|
||||
t_flatten = torch.stack(t_flatten)
|
||||
|
||||
# distance measurement
|
||||
dist_matrix = self.dist_fn(t_flatten)
|
||||
|
||||
# more similar with other filter indicates large in the sum of row
|
||||
distance = torch.sum(torch.abs(dist_matrix), 1)
|
||||
|
||||
return distance
|
||||
|
||||
def update_mask( # type: ignore[override]
|
||||
self, module, tensor_name, sparsity_level, **kwargs
|
||||
):
|
||||
tensor_weight = getattr(module, tensor_name)
|
||||
mask = getattr(module.parametrizations, tensor_name)[0].mask
|
||||
|
||||
if sparsity_level <= 0:
|
||||
mask.data = torch.ones_like(mask).bool()
|
||||
elif sparsity_level >= 1.0:
|
||||
mask.data = torch.zeros_like(mask).bool()
|
||||
else:
|
||||
distance = self._compute_distance(tensor_weight)
|
||||
|
||||
tensor_size = tensor_weight.shape[0] # prune filter (row)
|
||||
nparams_toprune = round(sparsity_level * tensor_size)
|
||||
nparams_toprune = min(
|
||||
max(nparams_toprune, 0), tensor_size
|
||||
) # clamp to [0, tensor_size]
|
||||
topk = torch.topk(distance, k=nparams_toprune, largest=False)
|
||||
mask[topk.indices] = False
|
||||
+5
@@ -0,0 +1,5 @@
|
||||
from .base_structured_sparsifier import BaseStructuredSparsifier
|
||||
from .FPGM_pruner import FPGMPruner
|
||||
from .lstm_saliency_pruner import LSTMSaliencyPruner
|
||||
from .parametrization import BiasHook, FakeStructuredSparsity
|
||||
from .saliency_pruner import SaliencyPruner
|
||||
+314
@@ -0,0 +1,314 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from collections.abc import Callable
|
||||
from itertools import chain
|
||||
from operator import getitem
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
from torch.ao.pruning.sparsifier.base_sparsifier import BaseSparsifier
|
||||
from torch.fx import symbolic_trace
|
||||
from torch.nn.utils import parametrize
|
||||
|
||||
from .match_utils import apply_match, MatchAllNode
|
||||
from .parametrization import BiasHook, FakeStructuredSparsity, module_contains_param
|
||||
from .prune_functions import (
|
||||
prune_conv2d,
|
||||
prune_conv2d_activation_conv2d,
|
||||
prune_conv2d_activation_pool_conv2d,
|
||||
prune_conv2d_conv2d,
|
||||
prune_conv2d_pool_activation_conv2d,
|
||||
prune_conv2d_pool_flatten_linear,
|
||||
prune_linear,
|
||||
prune_linear_activation_linear,
|
||||
prune_linear_linear,
|
||||
prune_lstm_output_layernorm_linear,
|
||||
prune_lstm_output_linear,
|
||||
)
|
||||
|
||||
|
||||
def _get_supported_structured_pruning_modules():
|
||||
SUPPORTED_STRUCTURED_PRUNING_MODULES = { # added to config if None given
|
||||
nn.Linear,
|
||||
nn.Conv2d,
|
||||
nn.LSTM,
|
||||
}
|
||||
return SUPPORTED_STRUCTURED_PRUNING_MODULES
|
||||
|
||||
|
||||
def _get_supported_activation_functions():
|
||||
SUPPORTED_ACTIVATION_FUNCTIONS = {
|
||||
F.relu,
|
||||
F.rrelu,
|
||||
F.hardtanh,
|
||||
F.relu6,
|
||||
F.sigmoid,
|
||||
F.hardsigmoid,
|
||||
F.tanh,
|
||||
F.silu,
|
||||
F.mish,
|
||||
F.hardswish,
|
||||
F.elu,
|
||||
F.celu,
|
||||
F.selu,
|
||||
F.hardshrink,
|
||||
F.leaky_relu,
|
||||
F.logsigmoid,
|
||||
F.softplus,
|
||||
F.prelu,
|
||||
F.softsign,
|
||||
F.tanhshrink,
|
||||
F.gelu,
|
||||
}
|
||||
return SUPPORTED_ACTIVATION_FUNCTIONS
|
||||
|
||||
|
||||
def _get_supported_activation_modules():
|
||||
SUPPORTED_ACTIVATION_MODULES = {
|
||||
nn.ReLU,
|
||||
nn.RReLU,
|
||||
nn.Hardtanh,
|
||||
nn.ReLU6,
|
||||
nn.Sigmoid,
|
||||
nn.Hardsigmoid,
|
||||
nn.Tanh,
|
||||
nn.SiLU,
|
||||
nn.Mish,
|
||||
nn.Hardswish,
|
||||
nn.ELU,
|
||||
nn.CELU,
|
||||
nn.SELU,
|
||||
nn.Hardshrink,
|
||||
nn.LeakyReLU,
|
||||
nn.LogSigmoid,
|
||||
nn.Softplus,
|
||||
nn.PReLU,
|
||||
nn.Softsign,
|
||||
nn.Tanhshrink,
|
||||
nn.GELU,
|
||||
}
|
||||
return SUPPORTED_ACTIVATION_MODULES
|
||||
|
||||
|
||||
def _get_default_structured_pruning_patterns() -> dict[
|
||||
tuple[type[nn.Module] | Callable | MatchAllNode | str, ...],
|
||||
Callable[..., None],
|
||||
]:
|
||||
"""
|
||||
Returns the patterns for conv2d / linear conversion for each element in the activation functions/modules defined above.
|
||||
"""
|
||||
patterns: dict[
|
||||
tuple[type[nn.Module] | Callable | MatchAllNode | str, ...],
|
||||
Callable[..., None],
|
||||
] = {
|
||||
# linear -> linear
|
||||
(nn.Linear, "output"): prune_linear,
|
||||
(nn.Linear, nn.Linear): prune_linear_linear,
|
||||
# conv2d -> conv2d
|
||||
(nn.Conv2d, "output"): prune_conv2d,
|
||||
(nn.Conv2d, nn.Conv2d): prune_conv2d_conv2d,
|
||||
# TODO LSTM Structured pruning does not support returned state currently.
|
||||
# Should find a way to explicitly match getitem(0) instead of getitem.
|
||||
# This will also require changing the pruning function.
|
||||
# lstm -> getitem(0) -> linear
|
||||
(nn.LSTM, getitem, nn.Linear): prune_lstm_output_linear,
|
||||
# lstm -> getitem(0) -> layernorm -> linear
|
||||
(nn.LSTM, getitem, nn.LayerNorm, nn.Linear): prune_lstm_output_layernorm_linear,
|
||||
}
|
||||
|
||||
for activation in chain(
|
||||
_get_supported_activation_functions(), _get_supported_activation_modules()
|
||||
):
|
||||
patterns.update(
|
||||
{
|
||||
# linear -> activation -> linear
|
||||
(nn.Linear, activation, nn.Linear): prune_linear_activation_linear,
|
||||
# conv2d -> activation -> conv2d
|
||||
(nn.Conv2d, activation, nn.Conv2d): prune_conv2d_activation_conv2d,
|
||||
# conv2d -> activation -> pool -> conv2d
|
||||
(
|
||||
nn.Conv2d,
|
||||
activation,
|
||||
nn.AvgPool2d,
|
||||
nn.Conv2d,
|
||||
): prune_conv2d_activation_pool_conv2d,
|
||||
(
|
||||
nn.Conv2d,
|
||||
activation,
|
||||
F.avg_pool2d,
|
||||
nn.Conv2d,
|
||||
): prune_conv2d_activation_pool_conv2d,
|
||||
(
|
||||
nn.Conv2d,
|
||||
activation,
|
||||
nn.MaxPool2d,
|
||||
nn.Conv2d,
|
||||
): prune_conv2d_activation_pool_conv2d,
|
||||
(
|
||||
nn.Conv2d,
|
||||
activation,
|
||||
F.max_pool2d,
|
||||
nn.Conv2d,
|
||||
): prune_conv2d_activation_pool_conv2d,
|
||||
# conv2d -> pool -> activation -> conv2d
|
||||
(
|
||||
nn.Conv2d,
|
||||
nn.AvgPool2d,
|
||||
activation,
|
||||
nn.Conv2d,
|
||||
): prune_conv2d_pool_activation_conv2d,
|
||||
(
|
||||
nn.Conv2d,
|
||||
F.avg_pool2d,
|
||||
activation,
|
||||
nn.Conv2d,
|
||||
): prune_conv2d_pool_activation_conv2d,
|
||||
(
|
||||
nn.Conv2d,
|
||||
nn.MaxPool2d,
|
||||
activation,
|
||||
nn.Conv2d,
|
||||
): prune_conv2d_pool_activation_conv2d,
|
||||
(
|
||||
nn.Conv2d,
|
||||
F.max_pool2d,
|
||||
activation,
|
||||
nn.Conv2d,
|
||||
): prune_conv2d_pool_activation_conv2d,
|
||||
# conv2d -> adaptive pool -> flatten -> linear
|
||||
(
|
||||
nn.Conv2d,
|
||||
nn.AdaptiveAvgPool2d,
|
||||
nn.Flatten,
|
||||
nn.Linear,
|
||||
): prune_conv2d_pool_flatten_linear,
|
||||
(
|
||||
nn.Conv2d,
|
||||
nn.AdaptiveAvgPool2d,
|
||||
torch.flatten,
|
||||
nn.Linear,
|
||||
): prune_conv2d_pool_flatten_linear,
|
||||
(
|
||||
nn.Conv2d,
|
||||
nn.AdaptiveMaxPool2d,
|
||||
nn.Flatten,
|
||||
nn.Linear,
|
||||
): prune_conv2d_pool_flatten_linear,
|
||||
(
|
||||
nn.Conv2d,
|
||||
nn.AdaptiveMaxPool2d,
|
||||
torch.flatten,
|
||||
nn.Linear,
|
||||
): prune_conv2d_pool_flatten_linear,
|
||||
}
|
||||
)
|
||||
return patterns
|
||||
|
||||
|
||||
class BaseStructuredSparsifier(BaseSparsifier):
|
||||
r"""Base class for structured pruning.
|
||||
|
||||
Abstract methods that need to be implemented:
|
||||
- update_mask: Function to compute a new mask for all keys in the
|
||||
`groups` attribute.
|
||||
|
||||
Args:
|
||||
- defaults [dict]: default configurations will be attached to the
|
||||
configuration. Only the keys that don't exist in the `config` will
|
||||
be updated.
|
||||
"""
|
||||
|
||||
def __init__(self, defaults, patterns=None):
|
||||
super().__init__(defaults)
|
||||
if patterns is None:
|
||||
patterns = _get_default_structured_pruning_patterns()
|
||||
self.patterns = patterns
|
||||
|
||||
def make_config_from_model(
|
||||
self,
|
||||
model: nn.Module,
|
||||
SUPPORTED_MODULES: set[type] | None = None,
|
||||
) -> None:
|
||||
if SUPPORTED_MODULES is None:
|
||||
SUPPORTED_MODULES = _get_supported_structured_pruning_modules()
|
||||
super().make_config_from_model(model, SUPPORTED_MODULES=SUPPORTED_MODULES)
|
||||
|
||||
def _prepare(self, *args, **kwargs) -> None:
|
||||
r"""This function will attach the FakeStructuredSparsity parameterizations
|
||||
and BiasHooks at the appropriate points in the model.
|
||||
"""
|
||||
for config in self.groups:
|
||||
module = config["module"]
|
||||
tensor_name = config["tensor_name"]
|
||||
parametrization = config.get("parametrization", FakeStructuredSparsity)
|
||||
tensor = getattr(module, tensor_name)
|
||||
|
||||
mask = config.get(
|
||||
"mask",
|
||||
torch.ones(tensor.shape[0], dtype=torch.bool, device=tensor.device),
|
||||
)
|
||||
self.state[config["tensor_fqn"]]["mask"] = mask
|
||||
parametrize.register_parametrization(
|
||||
module, tensor_name, parametrization(mask)
|
||||
)
|
||||
|
||||
# if linear / conv, we add in bias hooks
|
||||
if isinstance(module, (nn.Linear, nn.Conv2d)):
|
||||
prune_bias = config.get("prune_bias", True)
|
||||
if module.bias is not None:
|
||||
module.register_parameter(
|
||||
"_bias", nn.Parameter(module.bias.detach())
|
||||
)
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
module.bias = None
|
||||
module.prune_bias = prune_bias
|
||||
|
||||
module.register_forward_hook(
|
||||
BiasHook(module.parametrizations.weight[0], prune_bias) # type: ignore[union-attr, index]
|
||||
)
|
||||
|
||||
def prune(self) -> None:
|
||||
r"""
|
||||
This function will FX symbolically trace the model and then find instances of the patterns
|
||||
defined in self.patterns (by default SUPPORTED_STRUCTURED_PRUNING_PATTERNS ).
|
||||
|
||||
For each pattern, it will apply to corresponding conversion function, which will modify the output
|
||||
and input size expected by the modules within the pattern
|
||||
"""
|
||||
|
||||
self.traced = symbolic_trace(self.model)
|
||||
modules = dict(self.traced.named_modules())
|
||||
|
||||
# Right now we check for matches simply by iterating across all the patterns
|
||||
# if this is slow we can store patterns in a trie-structure and modify this code for faster lookup
|
||||
for node in self.traced.graph.nodes:
|
||||
for pattern, convert_fn in self.patterns.items():
|
||||
matched = apply_match(modules, pattern, node, [])
|
||||
if matched is None:
|
||||
continue
|
||||
|
||||
# pyrefly: ignore [no-matching-overload]
|
||||
first_module = modules.get(node.target)
|
||||
# check if first module exists and has appropriate parameterization, otherwise skip
|
||||
if (
|
||||
first_module is not None
|
||||
and parametrize.is_parametrized(first_module)
|
||||
and module_contains_param(first_module, FakeStructuredSparsity)
|
||||
):
|
||||
convert_block = []
|
||||
for node in matched:
|
||||
if node.op == "call_module":
|
||||
convert_block.append(modules.get(node.target))
|
||||
elif node.op == "call_function":
|
||||
convert_block.append(node.target)
|
||||
convert_fn(*convert_block)
|
||||
|
||||
for module in self.traced.modules():
|
||||
if module_contains_param(module, FakeStructuredSparsity):
|
||||
raise Exception( # noqa: TRY002
|
||||
f"Error: {module} still contains FakeStructuredSparsity parametrizations!"
|
||||
)
|
||||
|
||||
self.traced.graph.lint()
|
||||
self.traced.recompile()
|
||||
return self.traced # type: ignore[return-value]
|
||||
+54
@@ -0,0 +1,54 @@
|
||||
from typing import Any, cast
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from .base_structured_sparsifier import BaseStructuredSparsifier
|
||||
from .parametrization import FakeStructuredSparsity
|
||||
|
||||
|
||||
class LSTMSaliencyPruner(BaseStructuredSparsifier):
|
||||
"""
|
||||
Prune packed LSTM weights based on saliency.
|
||||
For each layer {k} inside a LSTM, we have two packed weight matrices
|
||||
- weight_ih_l{k}
|
||||
- weight_hh_l{k}
|
||||
|
||||
These tensors pack the weights for the 4 linear layers together for efficiency.
|
||||
|
||||
[W_ii | W_if | W_ig | W_io]
|
||||
|
||||
Pruning this tensor directly will lead to weights being misassigned when unpacked.
|
||||
To ensure that each packed linear layer is pruned the same amount:
|
||||
1. We split the packed weight into the 4 constituent linear parts
|
||||
2. Update the mask for each individual piece using saliency individually
|
||||
|
||||
This applies to both weight_ih_l{k} and weight_hh_l{k}.
|
||||
"""
|
||||
|
||||
def update_mask(self, module: nn.Module, tensor_name: str, **kwargs: Any) -> None:
|
||||
weights = getattr(module, tensor_name)
|
||||
|
||||
for p in getattr(module.parametrizations, tensor_name):
|
||||
if isinstance(p, FakeStructuredSparsity):
|
||||
mask = cast(torch.Tensor, p.mask)
|
||||
|
||||
# select weights based on magnitude
|
||||
if weights.dim() <= 1:
|
||||
raise Exception( # noqa: TRY002
|
||||
"Structured pruning can only be applied to a 2+dim weight tensor!"
|
||||
)
|
||||
# take norm over all but first dim
|
||||
dims = tuple(range(1, weights.dim()))
|
||||
saliency = weights.norm(dim=dims, p=1)
|
||||
|
||||
# handle weights in 4 groups
|
||||
split_size = len(mask) // 4
|
||||
masks = torch.split(mask, split_size)
|
||||
saliencies = torch.split(saliency, split_size)
|
||||
|
||||
for keep_mask, sal in zip(masks, saliencies):
|
||||
# mask smallest k values to be removed
|
||||
k = int(len(keep_mask) * kwargs["sparsity_level"])
|
||||
prune = sal.topk(k, largest=False, sorted=False).indices
|
||||
keep_mask.data[prune] = False # modifies underlying p.mask directly
|
||||
+65
@@ -0,0 +1,65 @@
|
||||
"""
|
||||
Contains utility functions to check if a pattern is in the graph and return the matching nodes
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.ao.quantization.utils import MatchAllNode
|
||||
from torch.fx import Node
|
||||
from torch.nn.utils import parametrize
|
||||
|
||||
|
||||
def _match(
|
||||
modules: dict[str, nn.ModuleDict],
|
||||
node: Node,
|
||||
current: nn.Module | Any,
|
||||
) -> bool:
|
||||
r"""
|
||||
checks to see if a single node of a pattern matches
|
||||
"""
|
||||
if isinstance(current, type) and issubclass(current, MatchAllNode):
|
||||
return True
|
||||
if not isinstance(node, Node):
|
||||
return False
|
||||
if isinstance(current, type) and issubclass(current, torch.nn.Module):
|
||||
return (
|
||||
node.op == "call_module"
|
||||
and parametrize.type_before_parametrizations(modules[node.target]) # type: ignore[index]
|
||||
== current
|
||||
)
|
||||
elif callable(current):
|
||||
return node.op == "call_function" and node.target is current
|
||||
elif isinstance(current, str):
|
||||
return node.target == current
|
||||
return False
|
||||
|
||||
|
||||
def apply_match(
|
||||
modules: dict[str, nn.ModuleDict],
|
||||
pattern: tuple[Any] | Any,
|
||||
node: Node,
|
||||
matched_node_pattern: list[Node],
|
||||
) -> list[Node] | None:
|
||||
r"""
|
||||
This function will return the matched nodes if the pattern matches the node given
|
||||
If there is no match, it will return None
|
||||
"""
|
||||
if isinstance(pattern, tuple):
|
||||
if len(pattern) == 1:
|
||||
if _match(modules, node, pattern[0]):
|
||||
return matched_node_pattern + [node]
|
||||
|
||||
first, *rest = pattern
|
||||
if _match(modules, node, first):
|
||||
if rest is None:
|
||||
return matched_node_pattern + [node]
|
||||
|
||||
for user in node.users:
|
||||
return apply_match(
|
||||
modules, tuple(rest), user, matched_node_pattern + [node]
|
||||
)
|
||||
elif _match(modules, node, pattern):
|
||||
return [node]
|
||||
return None
|
||||
+63
@@ -0,0 +1,63 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn.utils.parametrize import is_parametrized
|
||||
|
||||
|
||||
def module_contains_param(module, parametrization):
|
||||
if is_parametrized(module):
|
||||
# see if any of the module tensors have a parametriztion attached that matches the one passed in
|
||||
return any(
|
||||
any(isinstance(param, parametrization) for param in param_list)
|
||||
for key, param_list in module.parametrizations.items()
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
# Structured Pruning Parameterizations
|
||||
class FakeStructuredSparsity(nn.Module):
|
||||
r"""
|
||||
Parametrization for Structured Pruning. Like FakeSparsity, this should be attached to
|
||||
the 'weight' or any other parameter that requires a mask.
|
||||
|
||||
Instead of an element-wise bool mask, this parameterization uses a row-wise bool mask.
|
||||
"""
|
||||
|
||||
def __init__(self, mask):
|
||||
super().__init__()
|
||||
self.register_buffer("mask", mask)
|
||||
|
||||
def forward(self, x):
|
||||
if not isinstance(self.mask, torch.Tensor):
|
||||
raise AssertionError("mask must be a torch.Tensor")
|
||||
if self.mask.shape[0] != x.shape[0]:
|
||||
raise AssertionError(
|
||||
f"mask shape[0] ({self.mask.shape[0]}) must match x shape[0] ({x.shape[0]})"
|
||||
)
|
||||
shape = [1] * len(x.shape)
|
||||
shape[0] = -1
|
||||
return self.mask.reshape(shape) * x
|
||||
|
||||
def state_dict(self, *args, **kwargs):
|
||||
# avoid double saving masks
|
||||
return {}
|
||||
|
||||
|
||||
class BiasHook:
|
||||
def __init__(self, parametrization, prune_bias):
|
||||
self.param = parametrization
|
||||
self.prune_bias = prune_bias
|
||||
|
||||
def __call__(self, module, input, output):
|
||||
if getattr(module, "_bias", None) is not None:
|
||||
bias = module._bias.data
|
||||
if self.prune_bias:
|
||||
bias[~self.param.mask] = 0
|
||||
|
||||
# reshape bias to broadcast over output dimensions
|
||||
idx = [1] * len(output.shape)
|
||||
idx[1] = -1
|
||||
bias = bias.reshape(idx)
|
||||
|
||||
output += bias
|
||||
return output
|
||||
+485
@@ -0,0 +1,485 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""
|
||||
Collection of conversion functions for linear / conv2d structured pruning
|
||||
Also contains utilities for bias propagation
|
||||
"""
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import cast
|
||||
|
||||
import torch
|
||||
from torch import nn, Tensor
|
||||
from torch.nn.utils import parametrize
|
||||
from torch.nn.utils.parametrize import ParametrizationList
|
||||
|
||||
from .parametrization import BiasHook, FakeStructuredSparsity
|
||||
|
||||
|
||||
# BIAS PROPAGATION
|
||||
def _remove_bias_handles(module: nn.Module) -> None:
|
||||
if hasattr(module, "_forward_hooks"):
|
||||
bias_hooks: list[int] = []
|
||||
for key, hook in module._forward_hooks.items():
|
||||
if isinstance(hook, BiasHook):
|
||||
bias_hooks.append(key)
|
||||
|
||||
for key in bias_hooks:
|
||||
del module._forward_hooks[key]
|
||||
|
||||
|
||||
def _get_adjusted_next_layer_bias(
|
||||
next_layer: nn.Module, pruned_biases: Tensor, mask: Tensor
|
||||
) -> nn.Parameter:
|
||||
r"""Returns new adjusted bias for the second supported module"""
|
||||
if parametrize.is_parametrized(next_layer):
|
||||
# need to access original weight
|
||||
parametrization_dict = cast(nn.ModuleDict, next_layer.parametrizations)
|
||||
weight_parameterizations = cast(
|
||||
ParametrizationList, parametrization_dict.weight
|
||||
)
|
||||
next_weight = weight_parameterizations.original
|
||||
else:
|
||||
next_weight = cast(Tensor, next_layer.weight)
|
||||
|
||||
scaling_weight = next_weight[:, ~mask]
|
||||
if isinstance(next_layer, nn.Conv2d): # checking for Conv2d
|
||||
# Propagating first layer pruned biases and calculating the new second layer bias
|
||||
# involves more steps since the Conv2d scaling weight has extra dimensions,
|
||||
# so adding bias involves broadcasting, logically:
|
||||
# for each channel k in range(oC):
|
||||
# scaled_biases = sum(first_bias[pruned_idx] @ next_weight[k, pruned_idx, :, :].T)
|
||||
# new_next_bias[k] = old_next_bias[k] + scaled_biases
|
||||
scaling_product = torch.matmul(
|
||||
pruned_biases.reshape(1, -1), torch.transpose(scaling_weight, 1, 2)
|
||||
)
|
||||
sum_range = list(range(len(scaling_product.shape)))[
|
||||
1:
|
||||
] # all but the first dimension
|
||||
scaled_biases = torch.sum(scaling_product, sum_range)
|
||||
elif isinstance(next_layer, nn.Linear): # Linear
|
||||
scaled_biases = torch.matmul(
|
||||
pruned_biases, torch.transpose(scaling_weight, 0, 1)
|
||||
) # recall b2_new = b1 @ w2.T + b2
|
||||
else:
|
||||
raise NotImplementedError(f"Type {type(next_layer)} not supported yet.")
|
||||
|
||||
if (
|
||||
parametrize.is_parametrized(next_layer)
|
||||
and getattr(next_layer, "_bias", None) is not None
|
||||
): # next_layer is parametrized & has original bias ._bias
|
||||
adjusted_bias = nn.Parameter(scaled_biases + next_layer._bias) # type: ignore[operator]
|
||||
elif (
|
||||
not parametrize.is_parametrized(next_layer) and next_layer.bias is not None
|
||||
): # next_layer not parametrized & has .bias
|
||||
adjusted_bias = nn.Parameter(scaled_biases + next_layer.bias) # type: ignore[operator]
|
||||
else: # next_layer has no bias
|
||||
adjusted_bias = nn.Parameter(scaled_biases)
|
||||
return adjusted_bias
|
||||
|
||||
|
||||
def _prune_module_bias(module: nn.Module, mask: Tensor) -> None:
|
||||
r"""Applies mask to given modules bias"""
|
||||
# prune bias along with weights, discard pruned indices of bias
|
||||
original_bias = cast(Tensor, getattr(module, "_bias", module.bias))
|
||||
if original_bias is not None:
|
||||
module.bias = nn.Parameter(original_bias[mask])
|
||||
|
||||
# remove _bias parameter
|
||||
if hasattr(module, "_bias"):
|
||||
delattr(module, "_bias")
|
||||
|
||||
|
||||
def _propagate_module_bias(module: nn.Module, mask: Tensor) -> Tensor | None:
|
||||
r"""
|
||||
In the case that we need to propagate biases, this function will return the biases we need
|
||||
"""
|
||||
# set current module bias
|
||||
if module.bias is not None:
|
||||
module.bias = nn.Parameter(cast(Tensor, module.bias)[mask])
|
||||
elif getattr(module, "_bias", None) is not None:
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
module.bias = nn.Parameter(cast(Tensor, module._bias)[mask])
|
||||
|
||||
# get pruned biases to propagate to subsequent layer
|
||||
if getattr(module, "_bias", None) is not None:
|
||||
pruned_biases = cast(Tensor, module._bias)[~mask]
|
||||
else:
|
||||
pruned_biases = None
|
||||
|
||||
if hasattr(module, "_bias"):
|
||||
delattr(module, "_bias")
|
||||
|
||||
return pruned_biases
|
||||
|
||||
|
||||
# LINEAR
|
||||
def _prune_linear_helper(linear: nn.Linear) -> Tensor:
|
||||
# expects linear to be a parameterized linear module
|
||||
parametrization_dict = cast(nn.ModuleDict, linear.parametrizations)
|
||||
weight_parameterizations = cast(ParametrizationList, parametrization_dict.weight)
|
||||
for p in weight_parameterizations:
|
||||
if isinstance(p, FakeStructuredSparsity):
|
||||
mask = cast(Tensor, p.mask)
|
||||
|
||||
with torch.no_grad():
|
||||
parametrize.remove_parametrizations(linear, "weight", leave_parametrized=True)
|
||||
linear.weight = nn.Parameter(linear.weight[mask]) # type: ignore[possibly-undefined]
|
||||
linear.out_features = linear.weight.shape[0]
|
||||
_remove_bias_handles(linear)
|
||||
|
||||
# pyrefly: ignore [unbound-name]
|
||||
return mask
|
||||
|
||||
|
||||
def prune_linear(linear: nn.Linear) -> None:
|
||||
mask = _prune_linear_helper(linear)
|
||||
if getattr(linear, "prune_bias", False):
|
||||
_prune_module_bias(linear, mask)
|
||||
|
||||
|
||||
def prune_linear_linear(linear1: nn.Linear, linear2: nn.Linear) -> None:
|
||||
prune_linear_activation_linear(linear1, None, linear2)
|
||||
|
||||
|
||||
def prune_linear_activation_linear(
|
||||
linear1: nn.Linear,
|
||||
activation: Callable[[Tensor], Tensor] | None,
|
||||
linear2: nn.Linear,
|
||||
):
|
||||
mask = _prune_linear_helper(linear1)
|
||||
if getattr(linear1, "prune_bias", False):
|
||||
_prune_module_bias(linear1, mask)
|
||||
else:
|
||||
pruned_biases = _propagate_module_bias(linear1, mask)
|
||||
if pruned_biases is not None:
|
||||
if activation:
|
||||
pruned_biases = activation(pruned_biases)
|
||||
linear2.bias = _get_adjusted_next_layer_bias(linear2, pruned_biases, mask)
|
||||
|
||||
with torch.no_grad():
|
||||
if parametrize.is_parametrized(linear2):
|
||||
parametrization_dict = cast(nn.ModuleDict, linear2.parametrizations)
|
||||
weight_parameterizations = cast(
|
||||
ParametrizationList, parametrization_dict.weight
|
||||
)
|
||||
|
||||
weight_parameterizations.original = nn.Parameter(
|
||||
weight_parameterizations.original[:, mask]
|
||||
)
|
||||
linear2.in_features = weight_parameterizations.original.shape[1]
|
||||
else:
|
||||
linear2.weight = nn.Parameter(linear2.weight[:, mask])
|
||||
linear2.in_features = linear2.weight.shape[1]
|
||||
|
||||
|
||||
# CONV2D
|
||||
def _prune_conv2d_helper(conv2d: nn.Conv2d) -> Tensor:
|
||||
parametrization_dict = cast(nn.ModuleDict, conv2d.parametrizations)
|
||||
weight_parameterizations = cast(ParametrizationList, parametrization_dict.weight)
|
||||
for p in weight_parameterizations:
|
||||
if isinstance(p, FakeStructuredSparsity):
|
||||
mask = cast(Tensor, p.mask)
|
||||
|
||||
with torch.no_grad():
|
||||
parametrize.remove_parametrizations(conv2d, "weight", leave_parametrized=True)
|
||||
conv2d.weight = nn.Parameter(conv2d.weight[mask]) # type: ignore[possibly-undefined]
|
||||
conv2d.out_channels = conv2d.weight.shape[0]
|
||||
|
||||
_remove_bias_handles(conv2d)
|
||||
# pyrefly: ignore [unbound-name]
|
||||
return mask
|
||||
|
||||
|
||||
def prune_conv2d_padded(conv2d_1: nn.Conv2d) -> None:
|
||||
parametrization_dict = cast(nn.ModuleDict, conv2d_1.parametrizations)
|
||||
weight_parameterizations = cast(ParametrizationList, parametrization_dict.weight)
|
||||
for p in weight_parameterizations:
|
||||
if isinstance(p, FakeStructuredSparsity):
|
||||
mask = cast(Tensor, p.mask)
|
||||
|
||||
with torch.no_grad():
|
||||
parametrize.remove_parametrizations(conv2d_1, "weight", leave_parametrized=True)
|
||||
|
||||
if getattr(conv2d_1, "_bias", None) is not None:
|
||||
if (
|
||||
conv2d_1.bias is not None
|
||||
): # conv2d_1 has original bias and bias propagated from previous layer
|
||||
new_bias = torch.zeros(conv2d_1.bias.shape)
|
||||
new_bias[mask] = conv2d_1.bias[mask] # type: ignore[possibly-undefined]
|
||||
# adjusted bias that to keep in conv2d_1
|
||||
# pyrefly: ignore [unbound-name]
|
||||
new_bias[~mask] = cast(Tensor, conv2d_1._bias)[~mask]
|
||||
# pruned biases that are kept instead of propagated
|
||||
conv2d_1.bias = nn.Parameter(new_bias)
|
||||
else: # conv2d_1 has only original bias
|
||||
conv2d_1.bias = nn.Parameter(cast(Tensor, conv2d_1._bias))
|
||||
else:
|
||||
# no original bias, only propagated bias
|
||||
if (
|
||||
conv2d_1.bias is not None
|
||||
): # conv2d_1 has bias propagated from previous layer
|
||||
conv2d_1.bias.data[~mask] = 0 # type: ignore[possibly-undefined]
|
||||
|
||||
if hasattr(conv2d_1, "_bias"):
|
||||
delattr(conv2d_1, "_bias")
|
||||
|
||||
|
||||
def prune_conv2d(conv2d: nn.Conv2d) -> None:
|
||||
mask = _prune_conv2d_helper(conv2d)
|
||||
if getattr(conv2d, "prune_bias", False):
|
||||
_prune_module_bias(conv2d, mask)
|
||||
|
||||
|
||||
def prune_conv2d_conv2d(conv2d_1: nn.Conv2d, conv2d_2: nn.Conv2d) -> None:
|
||||
prune_conv2d_activation_conv2d(conv2d_1, None, conv2d_2)
|
||||
|
||||
|
||||
def prune_conv2d_activation_conv2d(
|
||||
conv2d_1: nn.Conv2d,
|
||||
activation: Callable[[Tensor], Tensor] | None,
|
||||
conv2d_2: nn.Conv2d,
|
||||
):
|
||||
r"""
|
||||
Fusion Pattern for conv2d -> some activation module / function -> conv2d layers
|
||||
"""
|
||||
parametrization_dict = cast(nn.ModuleDict, conv2d_1.parametrizations)
|
||||
weight_parameterizations = cast(ParametrizationList, parametrization_dict.weight)
|
||||
for p in weight_parameterizations:
|
||||
if isinstance(p, FakeStructuredSparsity):
|
||||
mask = cast(Tensor, p.mask)
|
||||
|
||||
prune_bias = getattr(conv2d_1, "prune_bias", False)
|
||||
if (
|
||||
hasattr(conv2d_2, "padding")
|
||||
and cast(tuple[int], conv2d_2.padding) > (0, 0)
|
||||
and (conv2d_1.bias is not None or getattr(conv2d_1, "_bias", None) is not None)
|
||||
):
|
||||
prune_conv2d_padded(conv2d_1)
|
||||
else:
|
||||
mask = _prune_conv2d_helper(conv2d_1)
|
||||
if prune_bias:
|
||||
_prune_module_bias(conv2d_1, mask)
|
||||
else:
|
||||
pruned_biases = _propagate_module_bias(conv2d_1, mask)
|
||||
if pruned_biases is not None:
|
||||
if activation:
|
||||
pruned_biases = activation(pruned_biases)
|
||||
conv2d_2.bias = _get_adjusted_next_layer_bias(
|
||||
conv2d_2, pruned_biases, mask
|
||||
)
|
||||
|
||||
if (
|
||||
not (
|
||||
hasattr(conv2d_2, "padding")
|
||||
and cast(tuple[int], conv2d_2.padding) > (0, 0)
|
||||
)
|
||||
or conv2d_1.bias is None
|
||||
):
|
||||
with torch.no_grad():
|
||||
if parametrize.is_parametrized(conv2d_2):
|
||||
parametrization_dict = cast(
|
||||
nn.ModuleDict, conv2d_2.parametrizations
|
||||
)
|
||||
weight_parameterizations = cast(
|
||||
ParametrizationList, parametrization_dict.weight
|
||||
)
|
||||
weight_parameterizations.original = nn.Parameter(
|
||||
weight_parameterizations.original[:, mask]
|
||||
)
|
||||
conv2d_2.in_channels = weight_parameterizations.original.shape[1]
|
||||
else:
|
||||
conv2d_2.weight = nn.Parameter(conv2d_2.weight[:, mask])
|
||||
conv2d_2.in_channels = conv2d_2.weight.shape[1]
|
||||
|
||||
|
||||
def prune_conv2d_pool_activation_conv2d(
|
||||
c1: nn.Conv2d,
|
||||
pool: nn.Module,
|
||||
activation: Callable[[Tensor], Tensor] | None,
|
||||
c2: nn.Conv2d,
|
||||
) -> None:
|
||||
prune_conv2d_activation_conv2d(c1, activation, c2)
|
||||
|
||||
|
||||
def prune_conv2d_activation_pool_conv2d(
|
||||
c1: nn.Conv2d,
|
||||
activation: Callable[[Tensor], Tensor] | None,
|
||||
pool: nn.Module,
|
||||
c2: nn.Conv2d,
|
||||
) -> None:
|
||||
prune_conv2d_activation_conv2d(c1, activation, c2)
|
||||
|
||||
|
||||
def prune_conv2d_pool_flatten_linear(
|
||||
conv2d: nn.Conv2d,
|
||||
pool: nn.Module,
|
||||
flatten: Callable[[Tensor], Tensor] | None,
|
||||
linear: nn.Linear,
|
||||
) -> None:
|
||||
mask = _prune_conv2d_helper(conv2d)
|
||||
|
||||
# We map the pruned indices of the Conv2d output to the flattened indices of the Linear following the Flatten layer.
|
||||
# we determine the flattening scale (h * w), and readjust `first_pruned_indices`
|
||||
# (each idx maps to range idx * h * w to (idx+1) * h * w), `first_valid_indices`,
|
||||
# and `pruned_biases` (repeat each bias by h * w).
|
||||
if parametrize.is_parametrized(linear):
|
||||
parametrization_dict = cast(nn.ModuleDict, linear.parametrizations)
|
||||
weight_parameterizations = cast(
|
||||
ParametrizationList, parametrization_dict.weight
|
||||
)
|
||||
linear_ic = weight_parameterizations.original.shape[1]
|
||||
else:
|
||||
linear_ic = linear.weight.shape[1]
|
||||
|
||||
conv2d_oc = len(mask)
|
||||
if linear_ic % conv2d_oc != 0:
|
||||
raise AssertionError(
|
||||
f"Flattening from dimensions {conv2d_oc} to {linear_ic} not supported"
|
||||
)
|
||||
|
||||
flatten_scale = linear_ic // conv2d_oc
|
||||
flattened_mask = torch.tensor(
|
||||
[[val] * flatten_scale for val in mask], dtype=torch.bool, device=mask.device
|
||||
).flatten()
|
||||
|
||||
if getattr(conv2d, "prune_bias", False):
|
||||
_prune_module_bias(conv2d, mask)
|
||||
else:
|
||||
pruned_biases = cast(Tensor, _propagate_module_bias(conv2d, mask))
|
||||
flattened_pruned_biases = torch.tensor(
|
||||
[[bias] * flatten_scale for bias in pruned_biases], device=mask.device
|
||||
).flatten()
|
||||
linear.bias = _get_adjusted_next_layer_bias(
|
||||
linear, flattened_pruned_biases, flattened_mask
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
if parametrize.is_parametrized(linear):
|
||||
parametrization_dict = cast(nn.ModuleDict, linear.parametrizations)
|
||||
weight_parameterizations = cast(
|
||||
ParametrizationList, parametrization_dict.weight
|
||||
)
|
||||
weight_parameterizations.original = nn.Parameter(
|
||||
weight_parameterizations.original[:, flattened_mask]
|
||||
)
|
||||
linear.in_features = weight_parameterizations.original.shape[1]
|
||||
else:
|
||||
linear.weight = nn.Parameter(linear.weight[:, flattened_mask])
|
||||
linear.in_features = linear.weight.shape[1]
|
||||
|
||||
|
||||
def prune_lstm_output_linear(
|
||||
lstm: nn.LSTM, getitem: Callable, linear: nn.Linear
|
||||
) -> None:
|
||||
prune_lstm_output_layernorm_linear(lstm, getitem, None, linear)
|
||||
|
||||
|
||||
def prune_lstm_output_layernorm_linear(
|
||||
lstm: nn.LSTM,
|
||||
getitem: Callable,
|
||||
layernorm: nn.LayerNorm | None,
|
||||
linear: nn.Linear,
|
||||
) -> None:
|
||||
for i in range(lstm.num_layers):
|
||||
if parametrize.is_parametrized(lstm, f"weight_ih_l{i}"):
|
||||
parametrization_dict = cast(nn.ModuleDict, lstm.parametrizations)
|
||||
weight_parameterizations = cast(
|
||||
ParametrizationList, parametrization_dict[f"weight_ih_l{i}"]
|
||||
)
|
||||
mask = weight_parameterizations[0].mask
|
||||
|
||||
with torch.no_grad():
|
||||
parametrize.remove_parametrizations(
|
||||
lstm, f"weight_ih_l{i}", leave_parametrized=True
|
||||
)
|
||||
setattr(
|
||||
lstm,
|
||||
f"weight_ih_l{i}",
|
||||
nn.Parameter(getattr(lstm, f"weight_ih_l{i}")[mask]),
|
||||
)
|
||||
setattr(
|
||||
lstm,
|
||||
f"bias_ih_l{i}",
|
||||
nn.Parameter(getattr(lstm, f"bias_ih_l{i}")[mask]),
|
||||
)
|
||||
|
||||
if parametrize.is_parametrized(lstm, f"weight_hh_l{i}"):
|
||||
parametrization_dict = cast(nn.ModuleDict, lstm.parametrizations)
|
||||
weight_parameterizations = cast(
|
||||
ParametrizationList, parametrization_dict[f"weight_hh_l{i}"]
|
||||
)
|
||||
mask = weight_parameterizations[0].mask
|
||||
|
||||
with torch.no_grad():
|
||||
parametrize.remove_parametrizations(
|
||||
lstm, f"weight_hh_l{i}", leave_parametrized=True
|
||||
)
|
||||
# splitting out hidden-hidden masks
|
||||
W_hi, W_hf, W_hg, W_ho = torch.split(
|
||||
getattr(lstm, f"weight_hh_l{i}"), lstm.hidden_size
|
||||
)
|
||||
M_hi, M_hf, M_hg, M_ho = torch.split(mask, lstm.hidden_size) # type: ignore[arg-type]
|
||||
|
||||
# resize each individual weight separately
|
||||
W_hi = W_hi[M_hi][:, M_hi]
|
||||
W_hf = W_hf[M_hf][:, M_hf]
|
||||
W_hg = W_hg[M_hg][:, M_hg]
|
||||
W_ho = W_ho[M_ho][:, M_ho]
|
||||
|
||||
# concat, use this as new weight
|
||||
new_weight = torch.cat((W_hi, W_hf, W_hg, W_ho))
|
||||
setattr(lstm, f"weight_hh_l{i}", nn.Parameter(new_weight))
|
||||
setattr(
|
||||
lstm,
|
||||
f"bias_hh_l{i}",
|
||||
nn.Parameter(getattr(lstm, f"bias_hh_l{i}")[mask]),
|
||||
)
|
||||
|
||||
# If this is the final layer, then we need to prune linear layer columns
|
||||
if i + 1 == lstm.num_layers:
|
||||
lstm.hidden_size = int(M_hi.sum())
|
||||
with torch.no_grad():
|
||||
if parametrize.is_parametrized(linear):
|
||||
parametrization_dict = cast(
|
||||
nn.ModuleDict, linear.parametrizations
|
||||
)
|
||||
weight_parameterizations = cast(
|
||||
ParametrizationList, parametrization_dict.weight
|
||||
)
|
||||
|
||||
weight_parameterizations.original = nn.Parameter(
|
||||
weight_parameterizations.original[:, M_ho]
|
||||
)
|
||||
linear.in_features = weight_parameterizations.original.shape[1]
|
||||
else:
|
||||
linear.weight = nn.Parameter(linear.weight[:, M_ho])
|
||||
linear.in_features = linear.weight.shape[1]
|
||||
|
||||
# if layernorm module, prune weight and bias
|
||||
if layernorm is not None:
|
||||
layernorm.normalized_shape = (linear.in_features,)
|
||||
layernorm.weight = nn.Parameter(layernorm.weight[M_ho])
|
||||
layernorm.bias = nn.Parameter(layernorm.bias[M_ho])
|
||||
|
||||
# otherwise need to prune the columns of the input of the next LSTM layer
|
||||
else:
|
||||
with torch.no_grad():
|
||||
if parametrize.is_parametrized(lstm, f"weight_ih_l{i + 1}"):
|
||||
parametrization_dict = cast(
|
||||
nn.ModuleDict, lstm.parametrizations
|
||||
)
|
||||
weight_parameterizations = cast(
|
||||
ParametrizationList,
|
||||
getattr(parametrization_dict, f"weight_ih_l{i + 1}"),
|
||||
)
|
||||
|
||||
weight_parameterizations.original = nn.Parameter(
|
||||
weight_parameterizations.original[:, M_ho]
|
||||
)
|
||||
else:
|
||||
next_layer_weight = getattr(lstm, f"weight_ih_l{i + 1}")
|
||||
setattr(
|
||||
lstm,
|
||||
f"weight_ih_l{i + 1}",
|
||||
nn.Parameter(next_layer_weight[:, M_ho]),
|
||||
)
|
||||
+35
@@ -0,0 +1,35 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from .base_structured_sparsifier import BaseStructuredSparsifier
|
||||
|
||||
|
||||
class SaliencyPruner(BaseStructuredSparsifier):
|
||||
"""
|
||||
Prune rows based on the saliency (L1 norm) of each row.
|
||||
|
||||
This pruner works on N-Dimensional weight tensors.
|
||||
For each row, we will calculate the saliency, which is the sum the L1 norm of all weights in that row.
|
||||
We expect that the resulting saliency vector has the same shape as our mask.
|
||||
We then pick elements to remove until we reach the target sparsity_level.
|
||||
"""
|
||||
|
||||
def update_mask(self, module, tensor_name, **kwargs):
|
||||
# tensor_name will give you the FQN, all other entries in sparse config is present in kwargs
|
||||
weights = getattr(module, tensor_name)
|
||||
mask = getattr(module.parametrizations, tensor_name)[0].mask
|
||||
|
||||
# use negative weights so we can use topk (we prune out the smallest)
|
||||
if weights.dim() <= 1:
|
||||
raise Exception( # noqa: TRY002
|
||||
"Structured pruning can only be applied to a 2+dim weight tensor!"
|
||||
)
|
||||
saliency = -weights.norm(dim=tuple(range(1, weights.dim())), p=1)
|
||||
if saliency.shape != mask.shape:
|
||||
raise AssertionError(
|
||||
f"saliency shape ({saliency.shape}) must match mask shape ({mask.shape})"
|
||||
)
|
||||
|
||||
num_to_pick = int(len(mask) * kwargs["sparsity_level"])
|
||||
prune = saliency.topk(num_to_pick).indices
|
||||
|
||||
# Set the mask to be false for the rows we want to prune
|
||||
mask.data[prune] = False
|
||||
@@ -0,0 +1,23 @@
|
||||
# mypy: allow-untyped-defs
|
||||
__all__ = [
|
||||
"get_static_sparse_quantized_mapping",
|
||||
"get_dynamic_sparse_quantized_mapping",
|
||||
]
|
||||
|
||||
|
||||
def get_static_sparse_quantized_mapping():
|
||||
import torch.ao.nn.sparse
|
||||
|
||||
_static_sparse_quantized_mapping = {
|
||||
torch.nn.Linear: torch.ao.nn.sparse.quantized.Linear,
|
||||
}
|
||||
return _static_sparse_quantized_mapping
|
||||
|
||||
|
||||
def get_dynamic_sparse_quantized_mapping():
|
||||
import torch.ao.nn.sparse
|
||||
|
||||
_dynamic_sparse_quantized_mapping = {
|
||||
torch.nn.Linear: torch.ao.nn.sparse.quantized.dynamic.Linear,
|
||||
}
|
||||
return _dynamic_sparse_quantized_mapping
|
||||
+173
@@ -0,0 +1,173 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import warnings
|
||||
import weakref
|
||||
from functools import wraps
|
||||
|
||||
from torch.ao.pruning.sparsifier.base_sparsifier import BaseSparsifier
|
||||
|
||||
|
||||
__all__ = ["BaseScheduler"]
|
||||
|
||||
|
||||
class BaseScheduler:
|
||||
def __init__(self, sparsifier, last_epoch=-1, verbose=False):
|
||||
# Attach sparsifier
|
||||
if not isinstance(sparsifier, BaseSparsifier):
|
||||
raise TypeError(
|
||||
f"{type(sparsifier).__name__} is not an instance of torch.ao.pruning.BaseSparsifier"
|
||||
)
|
||||
self.sparsifier = sparsifier
|
||||
|
||||
# Initialize epoch and base sparsity levels
|
||||
|
||||
self.base_sl = [group["sparsity_level"] for group in sparsifier.groups]
|
||||
self.last_epoch = last_epoch
|
||||
|
||||
# Following https://github.com/pytorch/pytorch/issues/20124
|
||||
# We would like to ensure that `scheduler.step()` is called after
|
||||
# `sparsifier.step()`
|
||||
def with_counter(method):
|
||||
if getattr(method, "_with_counter", False):
|
||||
# `sparsifier.step()` has already been replaced, return.
|
||||
return method
|
||||
|
||||
# Keep a weak reference to the sparsifier instance to prevent
|
||||
# cyclic references.
|
||||
instance_ref = weakref.ref(method.__self__)
|
||||
# Get the unbound method for the same purpose.
|
||||
func = method.__func__
|
||||
cls = instance_ref().__class__
|
||||
del method
|
||||
|
||||
@wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
instance = instance_ref()
|
||||
instance._step_count += 1 # type: ignore[union-attr]
|
||||
wrapped = func.__get__(instance, cls)
|
||||
return wrapped(*args, **kwargs)
|
||||
|
||||
# Note that the returned function here is no longer a bound method,
|
||||
# so attributes like `__func__` and `__self__` no longer exist.
|
||||
wrapper._with_counter = True # type: ignore[attr-defined]
|
||||
return wrapper
|
||||
|
||||
self.sparsifier.step = with_counter(self.sparsifier.step) # type: ignore[assignment]
|
||||
self.sparsifier._step_count = 0 # type: ignore[attr-defined]
|
||||
self._step_count: int = 0
|
||||
self.verbose = verbose
|
||||
|
||||
# Housekeeping
|
||||
self._get_sl_called_within_step: bool = False
|
||||
|
||||
self.step()
|
||||
|
||||
def state_dict(self):
|
||||
"""Returns the state of the scheduler as a :class:`dict`.
|
||||
|
||||
It contains an entry for every variable in self.__dict__ which
|
||||
is not the sparsifier.
|
||||
"""
|
||||
return {
|
||||
key: value for key, value in self.__dict__.items() if key != "sparsifier"
|
||||
}
|
||||
|
||||
def load_state_dict(self, state_dict):
|
||||
"""Loads the schedulers state.
|
||||
|
||||
Args:
|
||||
state_dict (dict): scheduler state. Should be an object returned
|
||||
from a call to :meth:`state_dict`.
|
||||
"""
|
||||
self.__dict__.update(state_dict)
|
||||
|
||||
def get_last_sl(self):
|
||||
"""Return last computed sparsity level by current scheduler."""
|
||||
return self._last_sl
|
||||
|
||||
def get_sl(self):
|
||||
# Compute sparsity level using chainable form of the scheduler
|
||||
# Note: This method is not intended to be called directly, and is only
|
||||
# used by the ".step" method. Use .get_last_sl() instead.
|
||||
if not self._get_sl_called_within_step:
|
||||
warnings.warn(
|
||||
"To get the last sparsity level computed by the scheduler, "
|
||||
"please use `get_last_sl()`.",
|
||||
stacklevel=2,
|
||||
)
|
||||
raise NotImplementedError
|
||||
|
||||
def print_sl(self, is_verbose, group, sl, epoch=None):
|
||||
"""Display the current sparsity level."""
|
||||
if is_verbose:
|
||||
if epoch is None:
|
||||
print(f"Adjusting sparsity level of group {group} to {sl:.4e}.")
|
||||
else:
|
||||
print(
|
||||
f"Epoch {epoch:5d}: adjusting sparsity level of group {group} to {sl:.4e}."
|
||||
)
|
||||
|
||||
def __repr__(self):
|
||||
format_string = self.__class__.__name__ + " ("
|
||||
format_string += "\n"
|
||||
format_string += f"Sparsifier {self.sparsifier}\n"
|
||||
format_string += f" base_sl: {self.base_sl}\n"
|
||||
format_string += ")"
|
||||
return format_string
|
||||
|
||||
def step(self, epoch=None):
|
||||
# Raise warning if trying to call scheduler step before the sparsifier.
|
||||
# https://github.com/pytorch/pytorch/issues/20124
|
||||
if self._step_count == 1:
|
||||
if not hasattr(self.sparsifier.step, "_with_counter"):
|
||||
warnings.warn(
|
||||
"Seems like `sparsifier.step()` has been overridden after sparsity scheduler "
|
||||
"initialization. Please, make sure to call `sparsifier.step()` before "
|
||||
"`scheduler.step()`.",
|
||||
UserWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
# Just check if there were two first scheduler.step() calls before sparsifier.step()
|
||||
elif self.sparsifier._step_count < 1: # type: ignore[attr-defined]
|
||||
warnings.warn(
|
||||
"Detected call of `scheduler.step()` before `sparsifier.step()`. "
|
||||
"You have to make sure you run the sparsifier.step() BEFORE any "
|
||||
"calls to the scheduler.step().",
|
||||
UserWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
self._step_count += 1
|
||||
|
||||
class _enable_get_sl_call:
|
||||
def __init__(self, o):
|
||||
self.o = o
|
||||
|
||||
def __enter__(self):
|
||||
self.o._get_sl_called_within_step = True
|
||||
return self
|
||||
|
||||
def __exit__(self, type, value, traceback):
|
||||
self.o._get_sl_called_within_step = False
|
||||
|
||||
with _enable_get_sl_call(self):
|
||||
self.last_epoch += 1
|
||||
values = self.get_sl()
|
||||
|
||||
for i, data in enumerate(zip(self.sparsifier.groups, values)):
|
||||
param_group, sl = data
|
||||
param_group["sparsity_level"] = sl
|
||||
self.print_sl(self.verbose, i, sl, epoch)
|
||||
|
||||
self._last_sl = [group["sparsity_level"] for group in self.sparsifier.groups]
|
||||
self.sparsifier.enable_mask_update = True
|
||||
|
||||
def _make_sure_a_list(self, var):
|
||||
r"""Utility that extends it to the same length as the .groups, ensuring it is a list"""
|
||||
n = len(self.sparsifier.groups)
|
||||
if not isinstance(var, (list, tuple)):
|
||||
return [var] * n
|
||||
else:
|
||||
if len(var) != n:
|
||||
raise ValueError(f"Expected variable of length {n}, but got {len(var)}")
|
||||
return list(var) # We want the result to be in a list, not tuple
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user