Google Git
Sign in
android / platform / external / pytorch / 013b51f8cc / . / test / quantization / fx / test_quantize_fx.py
blob: 49c79cd9663e0c77e38dafce32bac604d4b5df25 [file] [log] [blame]
# Owner(s): ["oncall: quantization"]
from collections import OrderedDict
import contextlib
import torch
import torch.nn.functional as F
import torch.nn as nn
import torch.ao.nn.quantized as nnq
import torch.ao.nn.quantized.reference as nnqr
import torch.ao.nn.quantized.dynamic as nnqd
import torch.ao.nn.intrinsic as nni
import torch.ao.nn.intrinsic.quantized as nniq
import torch.ao.nn.intrinsic.quantized.dynamic as nniqd
import torch.multiprocessing as mp
from torch.fx.graph_module import _USER_PRESERVED_ATTRIBUTES_KEY
# graph mode quantization based on fx
from torch.ao.quantization.quantize_fx import (
prepare_fx,
convert_fx,
convert_to_reference_fx,
_convert_to_reference_decomposed_fx,
prepare_qat_fx,
fuse_fx,
)
from torch.ao.quantization.fx.quantize_handler import DefaultNodeQuantizeHandler
from torch.ao.quantization.fx.match_utils import (
_is_match,
MatchAllNode,
)
from torch.ao.quantization import (
QuantType,
)
from torch.ao.quantization.quant_type import _get_quant_type_to_str
from torch.ao.quantization import (
QuantStub,
DeQuantStub,
QuantWrapper,
default_qconfig,
default_dynamic_qconfig,
default_per_channel_qconfig,
default_qat_qconfig,
default_reuse_input_qconfig,
default_symmetric_qnnpack_qconfig,
default_symmetric_qnnpack_qat_qconfig,
per_channel_dynamic_qconfig,
float16_dynamic_qconfig,
float16_static_qconfig,
float_qparams_weight_only_qconfig,
float_qparams_weight_only_qconfig_4bit,
get_default_qconfig,
get_default_qat_qconfig,
get_default_qconfig_mapping,
get_default_qat_qconfig_mapping,
fuse_modules,
fuse_modules_qat,
prepare,
prepare_qat,
convert,
quantize_dynamic,
default_placeholder_observer,
default_weight_observer,
PerChannelMinMaxObserver,
FixedQParamsFakeQuantize,
FixedQParamsObserver,
FusedMovingAvgObsFakeQuantize,
FakeQuantize,
MovingAverageMinMaxObserver,
HistogramObserver,
ReuseInputObserver,
QConfig,
default_embedding_qat_qconfig,
)
from torch.ao.quantization.backend_config import (
get_fbgemm_backend_config,
get_qnnpack_backend_config,
BackendConfig,
BackendPatternConfig,
DTypeConfig,
DTypeWithConstraints,
ObservationType
)
from torch.ao.quantization.backend_config.native import (
get_test_only_legacy_native_backend_config,
)
from torch.ao.quantization.qconfig_mapping import (
_get_symmetric_qnnpack_qconfig_mapping,
_get_symmetric_qnnpack_qat_qconfig_mapping,
_GLOBAL_DICT_KEY,
_MODULE_NAME_DICT_KEY,
_MODULE_NAME_OBJECT_TYPE_ORDER_DICT_KEY,
_MODULE_NAME_REGEX_DICT_KEY,
_OBJECT_TYPE_DICT_KEY,
QConfigMapping,
)
from torch.ao.quantization.fx.qconfig_mapping_utils import (
_get_object_type_qconfig,
_get_module_name_qconfig,
_get_module_name_regex_qconfig,
_maybe_adjust_qconfig_for_module_name_object_type_order,
)
from torch.ao.quantization.fx.pattern_utils import (
_DEFAULT_FUSION_PATTERNS,
_DEFAULT_QUANTIZATION_PATTERNS,
_DEFAULT_OUTPUT_FAKE_QUANTIZE_MAP,
_DEFAULT_OUTPUT_OBSERVER_MAP,
_register_fusion_pattern,
_register_quant_pattern,
get_default_output_activation_post_process_map
)
from torch.ao.quantization.fx.custom_config import (
STANDALONE_MODULE_NAME_DICT_KEY,
STANDALONE_MODULE_CLASS_DICT_KEY,
FLOAT_TO_OBSERVED_DICT_KEY,
OBSERVED_TO_QUANTIZED_DICT_KEY,
NON_TRACEABLE_MODULE_NAME_DICT_KEY,
NON_TRACEABLE_MODULE_CLASS_DICT_KEY,
INPUT_QUANTIZED_INDEXES_DICT_KEY,
OUTPUT_QUANTIZED_INDEXES_DICT_KEY,
PRESERVED_ATTRIBUTES_DICT_KEY,
FuseCustomConfig,
ConvertCustomConfig,
PrepareCustomConfig,
StandaloneModuleConfigEntry,
)
import torch.ao.quantization.fx.lstm_utils
from torch.ao.quantization.fx.utils import (
_reroute_tuple_getitem_pattern,
NodeInfo,
)
from torch.ao.quantization.fake_quantize import (
default_fixed_qparams_range_0to1_fake_quant,
default_fixed_qparams_range_neg1to1_fake_quant,
)
from torch.ao.quantization.observer import (
default_fixed_qparams_range_0to1_observer,
default_fixed_qparams_range_neg1to1_observer,
MinMaxObserver,
_is_activation_post_process,
)
# test utils
from hypothesis import given, settings
from hypothesis import strategies as st
from torch.testing._internal.common_cuda import TEST_MULTIGPU, TEST_CUDA
from torch.testing._internal.common_quantization import (
LinearReluLinearModel,
LinearReluModel,
LinearBnLeakyReluModel,
LinearTanhModel,
ConvBnAddReluModel,
QuantizationTestCase,
skipIfNoFBGEMM,
skipIfNoQNNPACK,
skip_if_no_torchvision,
train_one_epoch,
run_ddp,
test_only_eval_fn,
test_only_train_fn,
ModelForConvTransposeBNFusion,
get_supported_device_types,
skipIfNoONEDNN,
)
from torch.testing._internal.common_quantization import (
LinearModelWithSubmodule,
ResNetBase,
RNNDynamicModel,
RNNCellDynamicModel,
)
from torch.testing._internal.common_quantized import (
supported_qengines,
override_qengines,
override_quantized_engine,
)
from torch.testing._internal.common_utils import (
TemporaryFileName,
IS_ARM64,
)
from torch.testing._internal.common_quantization import NodeSpec as ns
from torch.testing import FileCheck
import copy
import itertools
import operator
import unittest
import io
from typing import Callable, Optional, List, Tuple
class BinaryOp(torch.nn.Module):
def __init__(self, binary_op, ibinary_op, is_inplace, is_scalar):
""" ibinary_op means inplace binary op
"""
super().__init__()
self.conv1 = torch.nn.Conv2d(1, 1, 1).float()
self.conv2 = torch.nn.Conv2d(1, 1, 1).float()
self.is_scalar = is_scalar
self.op = ibinary_op if ibinary_op and is_inplace else binary_op
def forward(self, x, y):
x = self.conv1(x)
y = 3 if self.is_scalar else self.conv2(y)
# x = x + y
x = self.op(x, y)
# x = y + x
x = self.op(y, x)
return x
class BinaryOpNonQuantizedInput(torch.nn.Module):
def __init__(self, binary_op, ibinary_op, is_inplace, is_scalar):
""" ibinary_op means inplace binary op
"""
super().__init__()
self.is_scalar = is_scalar
self.op = ibinary_op if ibinary_op and is_inplace else binary_op
def forward(self, x, y):
y = 3 if self.is_scalar else y
x = self.op(x, y)
return x
class BinaryOpRelu(torch.nn.Module):
def __init__(self, binary_op, ibinary_op, is_inplace, relu_callable,
is_scalar):
""" ibinary_op means inplace binary op
"""
super().__init__()
self.conv1 = torch.nn.Conv2d(1, 1, 1).float()
self.conv2 = torch.nn.Conv2d(1, 1, 1).float()
self.op = ibinary_op if ibinary_op and is_inplace else binary_op
self.relu_callable = relu_callable
self.is_scalar = is_scalar
if relu_callable is torch.nn.ReLU:
self.relu = torch.nn.ReLU()
else:
self.relu = relu_callable
def forward(self, x, y):
x = self.conv1(x)
y = 3 if self.is_scalar else self.conv2(y)
x = self.op(x, y)
x = self.relu(x)
x = self.op(y, x)
x = self.relu(x)
return x
@torch.fx.wrap
def _user_func_with_complex_return_type(x):
return list(torch.split(x, 1, 1))
class TestFuseFx(QuantizationTestCase):
def test_fuse_conv_bn_relu(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1d = nn.Conv1d(1, 1, 1)
self.conv2d = nn.Conv2d(1, 1, 1)
self.conv3d = nn.Conv3d(1, 1, 1)
self.bn1d = nn.BatchNorm1d(1)
self.bn2d = nn.BatchNorm2d(1)
self.bn3d = nn.BatchNorm3d(1)
self.conv1d2 = nn.Conv1d(1, 1, 1)
self.conv2d2 = nn.Conv2d(1, 1, 1)
self.conv3d2 = nn.Conv3d(1, 1, 1)
self.bn1d2 = nn.BatchNorm1d(1)
self.bn2d2 = nn.BatchNorm2d(1)
self.bn3d2 = nn.BatchNorm3d(1)
self.relu = nn.ReLU()
def forward(self, x):
x = self.conv1d(x)
x = self.bn1d(x)
x = self.conv2d(x)
x = self.bn2d(x)
x = self.conv3d(x)
x = self.bn3d(x)
x = self.conv1d2(x)
x = self.bn1d2(x)
x = self.relu(x)
x = self.conv2d2(x)
x = self.bn2d2(x)
x = self.relu(x)
x = self.conv3d2(x)
x = self.bn3d2(x)
x = self.relu(x)
return x
# test train mode
m = M().train()
# currently we don't check if the module are configured with qconfig before fusion
# TODO: if we decide to do that in the future, this test needs to
# be updated
# train mode fuse_fx is called in prepare_qat_fx
m = prepare_qat_fx(m, {}, example_inputs=(torch.randn(1, 1, 1, 1),))
expected_nodes = [
ns.call_module(nni.ConvBn1d),
ns.call_module(nni.ConvBn2d),
ns.call_module(nni.ConvBn3d),
ns.call_module(nni.ConvBnReLU1d),
ns.call_module(nni.ConvBnReLU2d),
ns.call_module(nni.ConvBnReLU3d),
]
expected_occurrence = {
ns.call_module(nn.ReLU): 0
}
self.checkGraphModuleNodes(
m,
expected_node_list=expected_nodes,
expected_node_occurrence=expected_occurrence)
# test eval mode
m = M().eval()
# fuse_fx is a top level api and only supports eval mode
m = fuse_fx(m)
expected_nodes = [
ns.call_module(nn.Conv1d),
ns.call_module(nn.Conv2d),
ns.call_module(nn.Conv3d),
ns.call_module(nni.ConvReLU1d),
ns.call_module(nni.ConvReLU2d),
ns.call_module(nni.ConvReLU3d),
]
# ConvBnRelu1d is not fused
expected_occurrence = {
ns.call_module(nn.ReLU): 0
}
self.checkGraphModuleNodes(
m,
expected_node_list=expected_nodes,
expected_node_occurrence=expected_occurrence)
def test_fuse_linear_bn_eval(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(1, 1)
self.bn1d = nn.BatchNorm1d(1)
def forward(self, x):
x = self.linear(x)
x = self.bn1d(x)
return x
# test eval mode
m = M().eval()
# fuse_fx is a top level api and only supports eval mode
m = fuse_fx(m)
expected_nodes = [
ns.call_module(nn.Linear),
]
expected_occurrence = {
ns.call_module(nn.BatchNorm1d): 0,
}
self.checkGraphModuleNodes(
m,
expected_node_list=expected_nodes,
expected_node_occurrence=expected_occurrence)
@skipIfNoONEDNN
def test_fuse_linear_bn_leaky_relu_onednn(self):
# linear - bn - leaky_relu is fused for onednn backend only
from torch.ao.quantization.backend_config import get_onednn_backend_config
expected_nodes = [
ns.call_module(nni.LinearLeakyReLU),
]
expected_occurrence = {
ns.call_module(nn.BatchNorm1d): 0,
ns.call_module(nn.LeakyReLU): 0,
}
for with_bn in [True, False]:
# test eval mode
m = LinearBnLeakyReluModel(with_bn).eval()
# fuse_fx is a top level api and only supports eval mode
m = fuse_fx(m,
backend_config=get_onednn_backend_config())
self.checkGraphModuleNodes(
m,
expected_node_list=expected_nodes,
expected_node_occurrence=expected_occurrence)
def test_linear_bn_leaky_relu_not_fused_by_default(self):
# Make sure linear - bn - leaky_relu is not fused by default
for with_bn in [True, False]:
# test eval mode
m = LinearBnLeakyReluModel(with_bn).eval()
# fuse_fx is a top level api and only supports eval mode
m = fuse_fx(m)
expected_nodes = [
ns.call_module(nn.Linear),
ns.call_module(nn.LeakyReLU),
]
expected_occurrence = {
ns.call_module(nni.LinearLeakyReLU): 0,
}
self.checkGraphModuleNodes(
m,
expected_node_list=expected_nodes,
expected_node_occurrence=expected_occurrence)
@skipIfNoONEDNN
def test_fuse_linear_tanh_for_onednn_backend(self):
# linear - tanh is fused for onednn backend only
from torch.ao.quantization.backend_config import get_onednn_backend_config
expected_nodes = [
ns.call_module(nni.LinearTanh),
]
expected_occurrence = {
ns.call_module(nn.Linear): 0,
ns.call_module(nn.Tanh): 0,
}
# test eval mode
m = LinearTanhModel().eval()
# fuse_fx is a top level api and only supports eval mode
m = fuse_fx(m,
backend_config=get_onednn_backend_config())
self.checkGraphModuleNodes(
m,
expected_node_list=expected_nodes,
expected_node_occurrence=expected_occurrence)
def test_linear_tanh_not_fused_by_default(self):
# Make sure linear - tanh is not fused by default
# test eval mode
m = LinearTanhModel().eval()
# fuse_fx is a top level api and only supports eval mode
m = fuse_fx(m)
expected_nodes = [
ns.call_module(nn.Linear),
ns.call_module(nn.Tanh),
]
expected_occurrence = {
ns.call_module(nni.LinearTanh): 0,
}
self.checkGraphModuleNodes(
m,
expected_node_list=expected_nodes,
expected_node_occurrence=expected_occurrence)
def test_fuse_conv_bn_add_relu_onednn(self):
# conv - bn - add - relu is fused for onednn backend only
from torch.ao.quantization.backend_config import get_onednn_backend_config
options = itertools.product(
[True, False], # with_bn
[True, False], # with_relu
[True, False], # conv in the left
[True, False], # with_two_conv
[True, False], # use_torch_add
)
for with_bn, with_relu, left_conv, two_conv, use_torch_add in options:
expected_nodes = [
ns.call_module(nni.ConvAddReLU2d if with_relu else nni.ConvAdd2d),
]
expected_occurrence = {
ns.call_module(nni.ConvAddReLU2d if with_relu else nni.ConvAdd2d): 1,
ns.call_module(nn.BatchNorm2d): 0,
}
# test eval mode
m = ConvBnAddReluModel(
with_bn=with_bn,
with_relu=with_relu,
left_conv=left_conv,
two_conv=two_conv,
use_torch_add=use_torch_add).eval()
m = fuse_fx(m,
backend_config=get_onednn_backend_config())
self.checkGraphModuleNodes(
m,
expected_node_list=expected_nodes,
expected_node_occurrence=expected_occurrence)
def test_fuse_conv_bn_add_relu_by_default(self):
options = itertools.product(
[True, False], # with_bn
[True, False], # with_relu
[True, False], # conv in the left
[True, False], # with_two_conv
[True, False], # use_torch_add
)
for with_bn, with_relu, left_conv, two_conv, use_torch_add in options:
# test eval mode
expected_nodes = [
ns.call_module(nn.Conv2d),
]
expected_occurrence = {
ns.call_module(nni.ConvAdd2d): 0,
}
m = ConvBnAddReluModel(
with_bn=with_bn,
with_relu=with_relu,
left_conv=left_conv,
two_conv=two_conv,
use_torch_add=use_torch_add).eval()
m = fuse_fx(m)
self.checkGraphModuleNodes(
m,
expected_node_list=expected_nodes,
expected_node_occurrence=expected_occurrence)
@skipIfNoONEDNN
def test_fuse_conv_bn_add_relu_lowering(self):
""" Test fusion and lowering of Conv2d - (bn -) ReLU
by FX. For onednn backedn only.
"""
from torch.ao.quantization.backend_config import get_onednn_backend_config
qconfig_mapping = get_default_qconfig_mapping('onednn')
with override_quantized_engine('onednn'):
options = itertools.product(
[True, False], # with_bn
[True, False], # with_relu
[True, False], # conv in the left
[True, False], # two_conv
[True, False], # use_torch_add
)
for with_bn, with_relu, left_conv, two_conv, use_torch_add in options:
node_occurrence = {
ns.call_function(torch.quantize_per_tensor): 1 if two_conv else 2,
ns.call_method("dequantize"): 1,
ns.call_module(nniq.ConvAddReLU2d if with_relu else nniq.ConvAdd2d): 1,
ns.call_module(nn.Conv2d): 0,
ns.call_module(nn.ReLU): 0,
}
node_occurrence_ref = {
ns.call_function(torch.quantize_per_tensor): 3,
ns.call_method("dequantize"): 3,
}
# test eval mode
m = ConvBnAddReluModel(
with_bn=with_bn,
with_relu=with_relu,
left_conv=left_conv,
two_conv=two_conv,
use_torch_add=use_torch_add).eval()
example_x = m.get_example_inputs()
m = prepare_fx(m, qconfig_mapping,
example_inputs=example_x,
backend_config=get_onednn_backend_config())
m_copy = copy.deepcopy(m)
m = convert_fx(m, backend_config=get_onednn_backend_config())
m_ref = convert_to_reference_fx(m_copy)
self.checkGraphModuleNodes(m, expected_node_occurrence=node_occurrence)
self.checkGraphModuleNodes(m_ref, expected_node_occurrence=node_occurrence_ref)
m(*example_x)
def test_fuse_convtranspose_bn_eval(self):
m = ModelForConvTransposeBNFusion().eval()
m = fuse_fx(m)
expected_nodes = [
ns.call_module(nn.ConvTranspose1d),
ns.call_module(nn.ConvTranspose2d),
ns.call_module(nn.ConvTranspose3d),
]
expected_occurrence = {
ns.call_module(nn.BatchNorm1d): 0,
ns.call_module(nn.BatchNorm2d): 0,
ns.call_module(nn.BatchNorm3d): 0,
}
self.checkGraphModuleNodes(
m,
expected_node_list=expected_nodes,
expected_node_occurrence=expected_occurrence)
def test_fuse_module_relu(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1d = nn.Conv1d(1, 1, 1)
self.conv2d = nn.Conv2d(1, 1, 1)
self.conv3d = nn.Conv3d(1, 1, 1)
self.bn1d = nn.BatchNorm1d(1)
self.bn2d = nn.BatchNorm2d(1)
self.bn3d = nn.BatchNorm3d(1)
self.relu = nn.ReLU()
def forward(self, x):
x = self.conv1d(x)
x = self.relu(x)
x = self.conv2d(x)
x = self.relu(x)
x = self.conv3d(x)
x = self.relu(x)
x = self.bn1d(x)
x = self.relu(x)
x = self.bn2d(x)
x = self.relu(x)
x = self.bn3d(x)
x = self.relu(x)
return x
m = M().eval()
m = fuse_fx(m)
expected_nodes = [
ns.call_module(nni.ConvReLU1d),
ns.call_module(nni.ConvReLU2d),
ns.call_module(nni.ConvReLU3d),
ns.call_module(nni.BNReLU2d),
ns.call_module(nni.BNReLU3d),
]
self.checkGraphModuleNodes(m, expected_node_list=expected_nodes)
@skipIfNoFBGEMM
def test_qconfig_fused_module(self):
""" TODO: add test for all fused modules
"""
qconfig_dict = {
"": None,
"object_type": [(nn.Linear, default_qconfig),
(nn.ReLU, default_qconfig),
(F.relu, default_qconfig)]
}
linearRelu_node_list = [
ns.call_function(torch.quantize_per_tensor),
ns.call_module(nniq.LinearReLU),
ns.call_method('dequantize')
]
linearReluLinear_node_list = [
ns.call_function(torch.quantize_per_tensor),
ns.call_module(nniq.LinearReLU),
ns.call_module(nnq.Linear),
ns.call_method('dequantize')
]
tests = [(LinearReluModel, linearRelu_node_list),
(LinearReluLinearModel, linearReluLinear_node_list)]
for M, node_list in tests:
m = M().eval()
example_inputs = (torch.rand(5, 5),)
prepared = prepare_fx(m, qconfig_dict, example_inputs=example_inputs)
prepared(*example_inputs)
quantized = convert_fx(prepared)
self.checkGraphModuleNodes(quantized, expected_node_list=node_list)
def test_problematic_fuse_example(self):
class LinearRelu(nn.Sequential):
def __init__(self):
super().__init__(
nn.Linear(5, 5),
nn.ReLU(),
)
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.lin_relu = LinearRelu()
self.linear = nn.Linear(5, 5)
def forward(self, x):
x = self.lin_relu(x)
x = self.linear(x)
return x
model = M().eval()
# these qconfigs somehow fail equality where default_qconfig does not
qconfig_dict = {
"": None,
"object_type": [
(torch.nn.Linear, get_default_qconfig('fbgemm')),
(torch.nn.ReLU, get_default_qconfig('fbgemm')),
],
}
m = prepare_fx(model, qconfig_dict, example_inputs=(torch.randn(1, 5),))
self.checkGraphModuleNodes(m, expected_node=ns.call_module(torch.ao.nn.intrinsic.modules.fused.LinearReLU))
@unittest.skip("Temporarily skipping the test case, will enable after the simple"
"pattern format is supported")
def test_fuse_addtional_fuser_method(self):
class MyConvReLU(torch.nn.Module):
pass
def my_conv_relu_fuser(conv, relu):
return MyConvReLU()
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv2d(3, 3, 3)
self.relu = torch.nn.ReLU()
def forward(self, x):
return self.relu(self.conv(x))
m = M().eval()
m = fuse_fx(m, fuse_custom_config={
"additional_fuser_method_mapping": {
(torch.nn.Conv2d, torch.nn.ReLU): my_conv_relu_fuser
}
})
self.checkGraphModuleNodes(m, expected_node=ns.call_module(MyConvReLU))
def test_fuse_custom_pattern(self):
class M(torch.nn.Module):
def __init__(self, use_torch_add=True):
super().__init__()
self.conv = torch.nn.Conv2d(3, 3, 3)
self.bn = torch.nn.BatchNorm2d(3)
self.relu = torch.nn.ReLU()
self.maxpool = torch.nn.MaxPool2d(3)
if use_torch_add:
self.add = torch.add
else:
self.add = operator.add
def forward(self, x):
y = x
y = self.maxpool(x)
x = self.conv(x)
x = self.bn(x)
x = self.add(y, x)
x = self.relu(x)
return x
for use_torch_add in [True, False]:
m = M(use_torch_add).eval()
def fuse_conv_bn_relu(is_qat, relu, add_pattern):
_, _, bn_pattern = add_pattern
bn, conv = bn_pattern
return conv
conv_bn_res_relu_config1 = BackendPatternConfig() \
._set_pattern_complex_format((nn.ReLU, (torch.add, MatchAllNode, (nn.BatchNorm2d, nn.Conv2d)))) \
.set_fuser_method(fuse_conv_bn_relu)
conv_bn_res_relu_config2 = BackendPatternConfig() \
._set_pattern_complex_format((nn.ReLU, (operator.add, MatchAllNode, (nn.BatchNorm2d, nn.Conv2d)))) \
.set_fuser_method(fuse_conv_bn_relu)
backend_config = BackendConfig() \
.set_backend_pattern_config(conv_bn_res_relu_config1) \
.set_backend_pattern_config(conv_bn_res_relu_config2)
m = fuse_fx(m, backend_config=backend_config)
self.assertEqual(type(m.conv), torch.nn.Conv2d)
# check bn and relu are gone since we replaced the whole pattern to conv
self.assertFalse(hasattr(m, "bn"))
self.assertFalse(hasattr(m, "relu"))
def test_fusion_pattern_with_multiple_inputs(self):
""" This test tests two keys in backend_config: root_node_getter and
extra_inputs_getter,
root_node_getter is used to identify a "root" module in the node pattern,
the node that we'll keep after fusion.
extra_inputs_getter will return a list of node that needs to be added to the
fused node as extra inputs.
"""
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv2d(3, 3, 3)
self.bn = torch.nn.BatchNorm2d(3)
self.relu = torch.nn.ReLU()
self.maxpool = torch.nn.MaxPool2d(3)
def forward(self, x):
y = x
y = self.maxpool(x)
x = self.conv(x)
x = self.bn(x)
x = torch.add(x, y)
x = self.relu(x)
return x
m = M().eval()
def fuse_conv_bn_relu(is_qat, relu, add_pattern):
_, bn_pattern, _ = add_pattern
bn, conv = bn_pattern
return conv
def conv_bn_res_relu_root_node_getter(pattern):
relu, add_pattern = pattern
_, bn_pattern, _ = add_pattern
bn, conv = bn_pattern
return conv
def conv_bn_res_relu_extra_inputs_getter(pattern):
""" get inputs pattern for extra inputs, inputs for root node
are assumed to be copied over from root node to the fused node
"""
relu, add_pattern = pattern
_, bn_pattern, extra_input = add_pattern
bn, conv = bn_pattern
return [extra_input]
conv_bn_res_relu_config = BackendPatternConfig() \
._set_pattern_complex_format((nn.ReLU, (torch.add, (nn.BatchNorm2d, nn.Conv2d), MatchAllNode))) \
.set_fuser_method(fuse_conv_bn_relu) \
._set_root_node_getter(conv_bn_res_relu_root_node_getter) \
._set_extra_inputs_getter(conv_bn_res_relu_extra_inputs_getter)
backend_config = BackendConfig().set_backend_pattern_config(conv_bn_res_relu_config)
m = fuse_fx(m, backend_config=backend_config)
self.assertEqual(type(m.conv), torch.nn.Conv2d)
# check bn and relu are gone since we replaced the whole pattern to conv
self.assertFalse(hasattr(m, "bn"))
self.assertFalse(hasattr(m, "relu"))
# check conv module has two inputs
named_modules = dict(m.named_modules())
for node in m.graph.nodes:
if node.op == "call_module" and type(named_modules[node.target]) == torch.nn.Conv2d:
self.assertTrue(len(node.args) == 2), "Expecting the fused op to have two arguments"
def test_fusion_pattern_with_matchallnode(self):
"""This test tests that the node matched by MatchAllNode will be regared as an input
instead of a module to be fused. For instance, we have two patterns:
(nn.ReLU, (torch.add, MatchAllNode, nn.Conv2d))
(nn.ReLU, nn.Conv2d)
And we wanna fuse the following model
Conv2d -> ReLU +
Conv2d ------ Add -> ReLU
ReLU in the first row is matched as MatchAllNode in the residual pattern. But it won't be
fused as part of that pattnern. It needs to be properly fused with the upstream Conv2d.
"""
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1 = torch.nn.Conv2d(3, 3, 3)
self.relu1 = torch.nn.ReLU()
self.conv2 = torch.nn.Conv2d(3, 3, 3)
self.relu2 = torch.nn.ReLU()
def forward(self, x):
y = self.conv1(x)
y = self.relu1(y)
x = self.conv2(x)
x = torch.add(x, y)
x = self.relu2(x)
return x
m = M().eval()
def fuse_conv_relu(is_qat, conv, relu):
return conv
def fuse_conv_res_relu(is_qat, relu, add_pattern):
_, conv, _ = add_pattern
return conv
def conv_res_relu_root_node_getter(pattern):
relu, (_, conv, _) = pattern
return conv
def conv_res_relu_extra_inputs_getter(pattern):
relu, (_, _, extra_input) = pattern
return [extra_input]
conv_relu_config = BackendPatternConfig((nn.Conv2d, nn.ReLU)) \
.set_fuser_method(fuse_conv_relu)
conv_res_relu_config = BackendPatternConfig() \
._set_pattern_complex_format((nn.ReLU, (torch.add, nn.Conv2d, MatchAllNode))) \
.set_fuser_method(fuse_conv_res_relu) \
._set_root_node_getter(conv_res_relu_root_node_getter) \
._set_extra_inputs_getter(conv_res_relu_extra_inputs_getter)
backend_config = BackendConfig() \
.set_backend_pattern_config(conv_relu_config) \
.set_backend_pattern_config(conv_res_relu_config)
m = fuse_fx(m, backend_config=backend_config)
self.assertEqual(type(m.conv1), torch.nn.Conv2d)
self.assertEqual(type(m.conv2), torch.nn.Conv2d)
# check relu are gone since we replaced both patterns to conv
self.assertFalse(hasattr(m, "relu1"))
self.assertFalse(hasattr(m, "relu2"))
@skipIfNoFBGEMM
class TestQuantizeFx(QuantizationTestCase):
def test_pattern_match(self):
""" test MatchAllNode with
conv - bn - add - relu pattern
"""
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(1, 1, 1)
self.bn = nn.BatchNorm2d(1)
self.relu = nn.ReLU()
def forward(self, x, y):
x = self.conv(x)
x = self.bn(x)
x = x + y
x = self.relu(x)
return x
pattern = (nn.ReLU, (operator.add, (nn.BatchNorm2d, nn.Conv2d), MatchAllNode))
m = torch.fx.symbolic_trace(M())
modules = dict(m.named_modules())
for n in m.graph.nodes:
if n.op == 'call_module' and type(modules[n.target]) == nn.ReLU:
self.assertTrue(_is_match(modules, n, pattern))
def test_pattern_match_constant(self):
class M(torch.nn.Module):
def forward(self, x):
x, _ = torch.ops.aten.max_pool2d_with_indices.default(x)
return x
pattern = (operator.getitem, torch.ops.aten.max_pool2d_with_indices.default, 0)
m = torch.fx.symbolic_trace(M())
# eliminate the code that get the second output of maxpool, so that the pattern
# can be matched
m.graph.eliminate_dead_code()
modules = dict(m.named_modules())
for n in m.graph.nodes:
if n.op == "call_function" and n.target == operator.getitem:
self.assertTrue(_is_match(modules, n, pattern))
def test_fused_module_qat_swap(self):
class Tmp(torch.nn.Module):
def __init__(self):
super().__init__()
self.tmp = torch.nn.Linear(5, 5)
self.relu = torch.nn.ReLU()
def forward(self, x):
x = self.tmp(x)
return self.relu(x)
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.mods1 = torch.nn.Sequential(Tmp(), torch.nn.Linear(5, 5))
self.mods2 = torch.nn.Linear(5, 5)
def forward(self, x):
a = self.mods1(x)
x = torch.add(x, 5)
x = self.mods2(x)
x = torch.add(x, 5)
return a, x
model = M().train()
qconfig_dict = {
"": None,
"object_type": [
(torch.nn.Linear, default_qat_qconfig),
(torch.nn.ReLU, default_qat_qconfig),
],
}
prepared = prepare_qat_fx(model, qconfig_dict, example_inputs=(torch.randn(1, 5),))
self.assertTrue(isinstance(getattr(prepared.mods1, "0").tmp, torch.ao.nn.intrinsic.qat.LinearReLU))
def _get_conv_linear_test_cases(self, is_reference):
""" Returns a list of test cases, with format:
is_dynamic, ModuleClass, module_constructor_inputs,
inputs, quantized_node, weight_prepack_op
"""
class FunctionalConv1d(torch.nn.Module):
def __init__(self, weight):
super().__init__()
self.weight = torch.nn.Parameter(weight)
self.stride = 1
self.padding = 0
self.dilation = 1
self.groups = 1
def forward(self, x):
return F.conv1d(x, self.weight, None, self.stride, self.padding, self.dilation, self.groups)
class Conv1d(torch.nn.Module):
def __init__(self, *args):
super().__init__()
self.conv = torch.nn.Conv1d(*args)
def forward(self, x):
return self.conv(x)
conv1d_input = torch.rand(1, 3, 224)
conv1d_weight = torch.rand(3, 3, 3)
conv1d_module_args = (3, 3, 3)
class FunctionalConv2d(torch.nn.Module):
def __init__(self, weight):
super().__init__()
self.weight = torch.nn.Parameter(weight)
self.stride = (1, 1)
self.padding = (0, 0)
self.dilation = (1, 1)
self.groups = 1
def forward(self, x):
return F.conv2d(x, self.weight, None, self.stride, self.padding, self.dilation, self.groups)
class Conv2d(torch.nn.Module):
def __init__(self, *args):
super().__init__()
self.conv = torch.nn.Conv2d(*args)
def forward(self, x):
return self.conv(x)
conv2d_input = torch.rand(1, 3, 224, 224)
conv2d_weight = torch.rand(3, 3, 3, 3)
conv2d_module_args = (3, 3, 3)
class FunctionalConv3d(torch.nn.Module):
def __init__(self, weight):
super().__init__()
self.weight = torch.nn.Parameter(weight)
self.stride = (1, 1, 1)
self.padding = (0, 0, 0)
self.dilation = (1, 1, 1)
self.groups = 1
def forward(self, x):
return F.conv3d(
x,
self.weight,
None,
self.stride,
self.padding,
self.dilation,
self.groups,
)
class Conv3d(torch.nn.Module):
def __init__(self, *args):
super().__init__()
self.conv = torch.nn.Conv3d(*args)
def forward(self, x):
return self.conv(x)
conv3d_input = torch.rand(1, 3, 32, 224, 224)
conv3d_weight = torch.rand(3, 3, 3, 3, 3)
conv3d_module_args = (3, 3, 3)
class Linear(torch.nn.Module):
def __init__(self, weight):
super().__init__()
self.weight = torch.nn.Parameter(weight)
def forward(self, x):
return F.linear(x, self.weight)
linear_input = torch.rand(8, 5)
linear_weight = torch.rand(10, 5)
class LinearModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(5, 10)
def forward(self, x):
return self.linear(x)
linear_module_input = torch.rand(8, 5)
# is_dynamic, ModuleClass, module_constructor_inputs,
# inputs, quantized_node, weight_prepack_node
tests = [
(
False,
FunctionalConv1d,
(conv1d_weight,),
(conv1d_input,),
ns.call_function(torch.nn.functional.conv1d if is_reference else torch.ops.quantized.conv1d) ,
ns.call_function(torch.ops.quantized.conv1d_prepack),
),
(
False,
FunctionalConv2d,
(conv2d_weight,),
(conv2d_input,),
ns.call_function(torch.nn.functional.conv2d if is_reference else torch.ops.quantized.conv2d),
ns.call_function(torch.ops.quantized.conv2d_prepack),
),
(
False,
FunctionalConv3d,
(conv3d_weight,),
(conv3d_input,),
ns.call_function(torch.nn.functional.conv3d if is_reference else torch.ops.quantized.conv3d),
ns.call_function(torch.ops.quantized.conv3d_prepack),
),
(
False,
Conv1d,
conv1d_module_args,
(conv1d_input,),
ns.call_module(nnqr.Conv1d if is_reference else nnq.Conv1d),
None
),
(
False,
Conv2d,
conv2d_module_args,
(conv2d_input,),
ns.call_module(nnqr.Conv2d if is_reference else nnq.Conv2d),
None
),
(
False,
Conv3d,
conv3d_module_args,
(conv3d_input,),
ns.call_module(nnqr.Conv3d if is_reference else nnq.Conv3d),
None
),
(
True,
Linear,
(linear_weight,),
(linear_input,),
None if is_reference else ns.call_function(torch.ops.quantized.linear_dynamic),
ns.call_function(torch.ops.quantized.linear_prepack),
),
(
False,
Linear,
(linear_weight,),
(linear_input,),
ns.call_function(torch.nn.functional.linear if is_reference else torch.ops.quantized.linear),
ns.call_function(torch.ops.quantized.linear_prepack),
),
(
True,
LinearModule,
(),
(linear_module_input,),
ns.call_module(nnqr.Linear) if is_reference else ns.call_module(nnqd.Linear),
None,
),
(
False,
LinearModule,
(),
(linear_module_input,),
ns.call_module(nnqr.Linear if is_reference else nnq.Linear),
None,
),
]
return tests
@skipIfNoFBGEMM
def test_conv_linear_not_reference(self):
""" Test quantizing conv and linear
"""
tests = self._get_conv_linear_test_cases(is_reference=False)
for (is_dynamic, ModuleClass, module_constructor_inputs,
inputs, quantized_node, weight_prepack_node) in tests:
quant_type = QuantType.DYNAMIC if is_dynamic else QuantType.STATIC
node_occurrence = {}
if weight_prepack_node:
node_occurrence[weight_prepack_node] = 0
self.checkGraphModeFxOp(
ModuleClass(*module_constructor_inputs),
inputs, quant_type,
expected_node=quantized_node,
expected_node_occurrence=node_occurrence,
is_reference=False)
@skipIfNoFBGEMM
def test_conv_linear_reference(self):
""" Test quantizing functional conv and linear with reference option
"""
tests = self._get_conv_linear_test_cases(is_reference=True)
def _get_keys(prefix, is_dynamic):
all_keys = [prefix + "." + k for k in ["weight_qscheme", "weight_dtype"]]
if not is_dynamic:
all_keys.extend([prefix + "." + k for k in ["weight_scale", "weight_zero_point"]])
return all_keys
for (is_dynamic, ModuleClass, module_constructor_inputs,
inputs, quantized_node, weight_prepack_node) in tests:
quant_type = QuantType.DYNAMIC if is_dynamic else QuantType.STATIC
node_occurrence = {}
if weight_prepack_node:
node_occurrence[weight_prepack_node] = 0
result_dict = self.checkGraphModeFxOp(
ModuleClass(*module_constructor_inputs),
inputs, quant_type,
expected_node=quantized_node,
expected_node_occurrence=node_occurrence,
is_reference=True)
qr = result_dict["quantized_reference"]
def checkWeightQParams(model):
for module_name in ("linear", "conv"):
if hasattr(model, module_name):
self.assertTrue(hasattr(qr.get_submodule(module_name), "weight_qscheme"))
self.assertTrue(hasattr(qr.get_submodule(module_name), "weight_scale"))
self.assertTrue(hasattr(qr.get_submodule(module_name), "weight_zero_point"))
self.assertTrue("Reference" in qr.get_submodule(module_name)._get_name())
def checkSerDeser(model, is_dynamic):
for module_name in ("linear", "conv"):
if hasattr(model, module_name):
# make sure seralization works
state_dict = copy.deepcopy(model.state_dict())
all_keys = _get_keys(module_name, is_dynamic)
for key in all_keys:
self.assertTrue(key in state_dict)
# check load_state_dict restores states
module = getattr(model, module_name)
prev_scale = module.weight_scale
module.weight_scale = None
model.load_state_dict(state_dict)
module = getattr(model, module_name)
self.assertTrue(torch.equal(prev_scale, module.weight_scale))
checkWeightQParams(qr)
qr = copy.deepcopy(qr)
# make sure the qparams are preserved after copy
checkWeightQParams(qr)
checkSerDeser(qr, is_dynamic)
def _get_conv_transpose_test_cases(self, use_relu, is_reference):
""" Returns a list of test cases, with format:
is_dynamic, ModuleClass, module_constructor_inputs,
inputs, quantized_node, weight_prepack_op
"""
class FunctionalConvTranspose1d(torch.nn.Module):
def __init__(self, weight):
super().__init__()
self.weight = torch.nn.Parameter(weight)
self.stride = 1
self.padding = 0
self.output_padding = 0
self.dilation = 1
self.groups = 1
def forward(self, x):
y = F.conv_transpose1d(
x,
self.weight,
None,
self.stride,
self.padding,
self.output_padding,
self.groups,
self.dilation
)
if use_relu:
y = F.relu(y)
return y
class ConvTranspose1d(torch.nn.Module):
def __init__(self, *args):
super().__init__()
self.deconv = torch.nn.ConvTranspose1d(*args)
self.relu = torch.nn.ReLU()
def forward(self, x):
y = self.deconv(x)
if use_relu:
y = self.relu(y)
return y
conv_transpose1d_input = torch.rand(1, 3, 224)
conv_transpose1d_weight = torch.rand(3, 3, 3)
conv_transpose1d_module_args = (3, 3, 3)
class FunctionalConvTranspose2d(torch.nn.Module):
def __init__(self, weight):
super().__init__()
self.weight = torch.nn.Parameter(weight)
self.stride = (1, 1)
self.padding = (0, 0)
self.output_padding = (0, 0)
self.dilation = (1, 1)
self.groups = 1
def forward(self, x):
y = F.conv_transpose2d(
x,
self.weight,
None,
self.stride,
self.padding,
self.output_padding,
self.groups,
self.dilation
)
if use_relu:
y = F.relu(y)
return y
class ConvTranspose2d(torch.nn.Module):
def __init__(self, *args):
super().__init__()
self.deconv = torch.nn.ConvTranspose2d(*args)
self.relu = torch.nn.ReLU()
def forward(self, x):
y = self.deconv(x)
if use_relu:
y = self.relu(y)
return y
conv_transpose2d_input = torch.rand(1, 3, 224, 224)
conv_transpose2d_weight = torch.rand(3, 3, 3, 3)
conv_transpose2d_module_args = (3, 3, 3)
class FunctionalConvTranspose3d(torch.nn.Module):
def __init__(self, weight):
super().__init__()
self.weight = torch.nn.Parameter(weight)
self.stride = (1, 1, 1)
self.padding = (0, 0, 0)
self.output_padding = (0, 0, 0)
self.dilation = (1, 1, 1)
self.groups = 1
def forward(self, x):
y = F.conv_transpose3d(
x,
self.weight,
None,
self.stride,
self.padding,
self.output_padding,
self.groups,
self.dilation
)
if use_relu:
y = F.relu(y)
return y
class ConvTranspose3d(torch.nn.Module):
def __init__(self, *args):
super().__init__()
self.deconv = torch.nn.ConvTranspose3d(*args)
self.relu = torch.nn.ReLU()
def forward(self, x):
y = self.deconv(x)
if use_relu:
y = self.relu(y)
return y
conv_transpose3d_input = torch.rand(1, 3, 32, 224, 224)
conv_transpose3d_weight = torch.rand(3, 3, 3, 3, 3)
conv_transpose3d_module_args = (3, 3, 3)
# is_dynamic, ModuleClass, module_constructor_inputs,
# inputs, quantized_node, weight_prepack_node
tests = [
(
False,
FunctionalConvTranspose1d,
(conv_transpose1d_weight,),
(conv_transpose1d_input,),
ns.call_function(
torch.nn.functional.conv_transpose1d if is_reference else torch.ops.quantized.conv_transpose1d
),
ns.call_function(torch.ops.quantized.conv_transpose1d_prepack),
),
(
False,
FunctionalConvTranspose2d,
(conv_transpose2d_weight,),
(conv_transpose2d_input,),
ns.call_function(
torch.nn.functional.conv_transpose2d if is_reference else torch.ops.quantized.conv_transpose2d
),
ns.call_function(torch.ops.quantized.conv_transpose2d_prepack),
),
(
False,
FunctionalConvTranspose3d,
(conv_transpose3d_weight,),
(conv_transpose3d_input,),
ns.call_function(
torch.nn.functional.conv_transpose3d if is_reference else torch.ops.quantized.conv_transpose3d),
ns.call_function(torch.ops.quantized.conv_transpose3d_prepack),
),
(
False,
ConvTranspose1d,
conv_transpose1d_module_args,
(conv_transpose1d_input,),
ns.call_module(nnqr.ConvTranspose1d if is_reference else nnq.ConvTranspose1d),
None
),
(
False,
ConvTranspose2d,
conv_transpose2d_module_args,
(conv_transpose2d_input,),
ns.call_module(nnqr.ConvTranspose2d if is_reference else nnq.ConvTranspose2d),
None
),
(
False,
ConvTranspose3d,
conv_transpose3d_module_args,
(conv_transpose3d_input,),
ns.call_module(nnqr.ConvTranspose3d if is_reference else nnq.ConvTranspose3d),
None
),
]
return tests
@skipIfNoFBGEMM
def test_conv_transpose_not_reference(self):
""" Test quantizing transposed conv
"""
tests = self._get_conv_transpose_test_cases(use_relu=False, is_reference=False)
for (is_dynamic, ModuleClass, module_constructor_inputs,
inputs, quantized_node, weight_prepack_node) in tests:
quant_type = QuantType.DYNAMIC if is_dynamic else QuantType.STATIC
node_occurrence = {}
if weight_prepack_node:
node_occurrence[weight_prepack_node] = 0
self.checkGraphModeFxOp(
ModuleClass(*module_constructor_inputs),
inputs, quant_type,
expected_node=quantized_node,
expected_node_occurrence=node_occurrence,
is_reference=False)
@skipIfNoFBGEMM
def test_conv_transpose_reference(self):
""" Test quantizing transposed conv with reference option
"""
tests = self._get_conv_transpose_test_cases(use_relu=False, is_reference=True)
def _get_keys(prefix, is_dynamic):
all_keys = [prefix + "." + k for k in ["weight_qscheme", "weight_dtype"]]
if not is_dynamic:
all_keys.extend([prefix + "." + k for k in ["weight_scale", "weight_zero_point"]])
return all_keys
for (is_dynamic, ModuleClass, module_constructor_inputs,
inputs, quantized_node, weight_prepack_node) in tests:
quant_type = QuantType.DYNAMIC if is_dynamic else QuantType.STATIC
node_occurrence = {}
if weight_prepack_node:
node_occurrence[weight_prepack_node] = 0
result_dict = self.checkGraphModeFxOp(
ModuleClass(*module_constructor_inputs),
inputs, quant_type,
expected_node=quantized_node,
expected_node_occurrence=node_occurrence,
is_reference=True)
qr = result_dict["quantized_reference"]
def checkWeightQParams(model):
module_name = "deconv"
if hasattr(model, module_name):
self.assertTrue(hasattr(qr.get_submodule(module_name), "weight_qscheme"))
self.assertTrue(hasattr(qr.get_submodule(module_name), "weight_scale"))
self.assertTrue(hasattr(qr.get_submodule(module_name), "weight_zero_point"))
self.assertTrue("Reference" in qr.get_submodule(module_name)._get_name())
def checkSerDeser(model, is_dynamic):
module_name = "deconv"
if hasattr(model, module_name):
# make sure seralization works
state_dict = copy.deepcopy(model.state_dict())
all_keys = _get_keys(module_name, is_dynamic)
for key in all_keys:
self.assertTrue(key in state_dict)
# check load_state_dict restores states
module = getattr(model, module_name)
prev_scale = module.weight_scale
module.weight_scale = None
model.load_state_dict(state_dict)
module = getattr(model, module_name)
self.assertTrue(torch.equal(prev_scale, module.weight_scale))
checkWeightQParams(qr)
qr = copy.deepcopy(qr)
# make sure the qparams are preserved after copy
checkWeightQParams(qr)
checkSerDeser(qr, is_dynamic)
def test_conv_transpose_relu_not_reference(self):
""" Test quantizing transposed conv + relu
Fusion with relu is not supported.
"""
tests = self._get_conv_transpose_test_cases(use_relu=True, is_reference=False)
for (is_dynamic, ModuleClass, module_constructor_inputs,
inputs, quantized_node, weight_prepack_node) in tests:
quant_type = QuantType.DYNAMIC if is_dynamic else QuantType.STATIC
node_occurrence = {}
if weight_prepack_node:
node_occurrence[weight_prepack_node] = 0
if quantized_node.op == 'call_module':
node_occurrence[ns.call_module(nn.ReLU)] = 1
else:
node_occurrence[ns.call_function(F.relu)] = 1
self.checkGraphModeFxOp(
ModuleClass(*module_constructor_inputs),
inputs, quant_type,
expected_node=quantized_node,
expected_node_occurrence=node_occurrence,
is_reference=False)
@skipIfNoFBGEMM
def test_conv_transpose_relu_reference(self):
""" Test quantizing transposed conv with reference option
Fusion with relu is not supported.
"""
tests = self._get_conv_transpose_test_cases(use_relu=True, is_reference=True)
def _get_keys(prefix, is_dynamic):
all_keys = [prefix + "." + k for k in ["weight_qscheme", "weight_dtype"]]
if not is_dynamic:
all_keys.extend([prefix + "." + k for k in ["weight_scale", "weight_zero_point"]])
return all_keys
for (is_dynamic, ModuleClass, module_constructor_inputs,
inputs, quantized_node, weight_prepack_node) in tests:
quant_type = QuantType.DYNAMIC if is_dynamic else QuantType.STATIC
node_occurrence = {}
if weight_prepack_node:
node_occurrence[weight_prepack_node] = 0
if quantized_node.op == 'call_module':
node_occurrence[ns.call_module(nn.ReLU)] = 1
else:
node_occurrence[ns.call_function(F.relu)] = 1
result_dict = self.checkGraphModeFxOp(
ModuleClass(*module_constructor_inputs),
inputs, quant_type,
expected_node=quantized_node,
expected_node_occurrence=node_occurrence,
is_reference=True)
qr = result_dict["quantized_reference"]
def checkWeightQParams(model):
module_name = "deconv"
if hasattr(model, module_name):
self.assertTrue(hasattr(qr.get_submodule(module_name), "weight_qscheme"))
self.assertTrue(hasattr(qr.get_submodule(module_name), "weight_scale"))
self.assertTrue(hasattr(qr.get_submodule(module_name), "weight_zero_point"))
self.assertTrue("Reference" in qr.get_submodule(module_name)._get_name())
def checkSerDeser(model, is_dynamic):
module_name = "deconv"
if hasattr(model, module_name):
# make sure seralization works
state_dict = copy.deepcopy(model.state_dict())
all_keys = _get_keys(module_name, is_dynamic)
for key in all_keys:
self.assertTrue(key in state_dict)
# check load_state_dict restores states
module = getattr(model, module_name)
prev_scale = module.weight_scale
module.weight_scale = None
model.load_state_dict(state_dict)
module = getattr(model, module_name)
self.assertTrue(torch.equal(prev_scale, module.weight_scale))
checkWeightQParams(qr)
qr = copy.deepcopy(qr)
# make sure the qparams are preserved after copy
checkWeightQParams(qr)
checkSerDeser(qr, is_dynamic)
@skipIfNoFBGEMM
def test_dynamic_quant_weight_observer(self):
''' Test that weight observer is run in convert step
'''
class M(torch.nn.Module):
def __init__(self, weight):
super().__init__()
self.weight = torch.nn.Parameter(weight)
def forward(self, x):
return F.linear(x, self.weight)
m = M(torch.rand(1, 1)).eval()
qconfig = default_dynamic_qconfig
qconfig_dict = {'': qconfig}
example_inputs = (torch.rand(1, 1),)
prepared = prepare_fx(m, qconfig_dict, example_inputs=example_inputs)
quantized = convert_to_reference_fx(prepared)
qparams = (quantized._scale_0, quantized._zero_point_0)
weight_obs = qconfig.weight()
weight_obs(quantized.weight)
# Get the actual value to avoid tensor size mismatch error, torch.Size([]) vs torch.Size([1])
ref_qparams = (weight_obs.calculate_qparams()[0].item(), weight_obs.calculate_qparams()[1].item())
self.assertEqual(qparams, ref_qparams)
def test_conv_bn_relu(self):
""" Tests fusion and quantization for "Conv - Bn" and "Conv - Bn - ReLU"
"""
convs = {
1: nn.Conv1d,
2: nn.Conv2d,
3: nn.Conv3d,
}
bns = {
1: nn.BatchNorm1d,
2: nn.BatchNorm2d,
3: nn.BatchNorm3d,
}
quantized_convs = {
1: nnq.Conv1d,
2: nnq.Conv2d,
3: nnq.Conv3d,
}
quantized_conv_relus = {
1: nniq.ConvReLU1d,
2: nniq.ConvReLU2d,
3: nniq.ConvReLU3d,
}
class M(torch.nn.Module):
def __init__(self, dim, has_relu):
super().__init__()
self.conv = convs[dim](3, 3, 3)
self.bn = bns[dim](3)
self.relu = nn.ReLU() if has_relu else nn.Identity()
self.has_relu = has_relu
self.quant = QuantStub()
self.dequant = DeQuantStub()
def forward(self, x):
x = self.quant(x)
x = self.conv(x)
x = self.bn(x)
if self.has_relu:
x = self.relu(x)
x = self.dequant(x)
return x
options = itertools.product([1, 2, 3], [True, False], self.static_quant_types)
for dim, has_relu, quant_type in options:
expected_node = ns.call_module(
quantized_conv_relus[dim] if has_relu
else quantized_convs[dim])
m = M(dim, has_relu)
m_eager = copy.deepcopy(m)
result_dict = self.checkGraphModeFxOp(
m,
self.img_data_dict[dim],
quant_type,
expected_node=expected_node,
)
result = result_dict["quantized_output"]
# check numerics
qengine = torch.backends.quantized.engine
if quant_type == QuantType.STATIC:
m_eager.eval()
qconfig = get_default_qconfig(qengine)
prepare_fn = prepare
is_qat = False
else:
m_eager.train()
qconfig = get_default_qat_qconfig(qengine)
prepare_fn = prepare_qat
is_qat = True
fuse_list = ["conv", "bn"]
if has_relu:
fuse_list.append("relu")
if is_qat:
fuse_modules_qat(m_eager, fuse_list, inplace=True)
else:
fuse_modules(m_eager, fuse_list, inplace=True)
m_eager.qconfig = qconfig
m_eager = prepare_fn(m_eager)
prepared_fx = result_dict["prepared"]
m_eager(*self.img_data_dict[dim][0])
m_eager = convert(m_eager)
result_eager = m_eager(*self.img_data_dict[dim][0])
self.assertEqual(result, result_eager)
def test_linear_bn(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(4, 4)
self.bn = nn.BatchNorm1d(4)
self.quant = QuantStub()
self.dequant = DeQuantStub()
def forward(self, x):
x = self.quant(x)
x = self.linear(x)
x = self.bn(x)
x = self.dequant(x)
return x
data = (torch.randn(4, 4),)
for quant_type in self.static_quant_types:
expected_node = ns.call_module(nnq.Linear)
m = M()
m_eager = copy.deepcopy(m)
result_dict = self.checkGraphModeFxOp(m, data, quant_type, expected_node=expected_node)
result = result_dict["quantized_output"]
# check numerics vs eager mode
fuse_list = ["linear", "bn"]
qengine = torch.backends.quantized.engine
if quant_type == QuantType.STATIC:
m_eager.eval()
qconfig = get_default_qconfig(qengine)
prepare_fn = prepare
fuse_modules(m_eager, fuse_list, inplace=True)
else:
m_eager.train()
qconfig = get_default_qat_qconfig(qengine)
prepare_fn = prepare_qat
fuse_modules_qat(m_eager, fuse_list, inplace=True)
m_eager.qconfig = qconfig
m_eager = prepare_fn(m_eager)
m_eager(*data)
m_eager = convert(m_eager)
result_eager = m_eager(*data)
self.assertEqual(result, result_eager)
@skipIfNoFBGEMM
def test_dynamic_quant_fp16(self):
with override_quantized_engine('fbgemm'):
class Linear(torch.nn.Module):
def __init__(self, weight):
super().__init__()
self.weight = torch.nn.Parameter(weight)
def forward(self, x):
return F.linear(x, self.weight)
linear_input = torch.rand(8, 5)
linear_weight = torch.rand(10, 5)
class LinearModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(5, 10)
def forward(self, x):
return self.linear(x)
linear_module_input = torch.rand(8, 5)
tests = [
(Linear, (linear_weight,), (linear_input,),
ns.call_function(torch.ops.quantized.linear_dynamic_fp16),
ns.call_function(torch.ops.quantized.linear_prepack_fp16)),
(LinearModule, (), (linear_module_input,),
ns.call_module(nnqd.Linear),
None),
]
for (ModuleClass, module_constructor_inputs,
inputs, quantized_node, weight_prepack_node) in tests:
for is_reference in [True, False]:
node_occurrence = {}
if weight_prepack_node:
node_occurrence[weight_prepack_node] = 0
m = ModuleClass(*module_constructor_inputs).eval()
qconfig_dict = {"": float16_dynamic_qconfig}
m = prepare_fx(m, qconfig_dict, example_inputs=inputs)
convert_fn = convert_to_reference_fx if is_reference else convert_fx
m = convert_fn(m)
self.checkGraphModuleNodes(m, expected_node_occurrence=node_occurrence)
@unittest.skipIf(not TEST_MULTIGPU, "multi-GPU not supported")
@unittest.skipIf(not TEST_CUDA, "CUDA unavailable")
@override_qengines
def test_qat_prepare_device_affinity(self):
"""
Tests that FX QAT prepare pass respects device affinity
"""
class Model(nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(1, 1, 1)
self.bn = nn.BatchNorm2d(1)
self.relu = nn.ReLU()
def forward(self, x):
x = self.conv(x)
x = self.bn(x)
x = self.relu(x)
return x
model = Model()
qengine = torch.backends.quantized.engine
qconfig_dict = {'': torch.ao.quantization.get_default_qat_qconfig(qengine)}
device = torch.device('cuda:0')
model.to(device)
example_inputs = (torch.randn(4, 1, 4, 4, device=device),)
# QAT prepare
model = prepare_qat_fx(model, qconfig_dict, example_inputs=example_inputs)
# ensure that running an input on CUDA works without any needed changes
model(*example_inputs)
# ensure all buffers and parameters are on the device we expect
model_devices = {p.device for p in model.parameters()} | \
{p.device for p in model.buffers()}
self.assertEqual(len(model_devices), 1)
model_device = next(iter(model_devices))
self.assertEqual(model_device, device)
@skipIfNoFBGEMM
def test_dict_output(self):
""" Make sure quantization runs for models with dictionary output
"""
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv2d(1, 1, 1)
def forward(self, x):
return {"output": self.conv(x["input"])}
example_inputs = ({"input": torch.randn(1, 1, 1, 1)},)
m = M().eval()
qconfig_dict = {"": default_qconfig}
m = prepare_fx(m, qconfig_dict, example_inputs=example_inputs)
m(*example_inputs)
m = convert_fx(m)
m(*example_inputs)
@override_qengines
def test_attention(self):
""" Make sure quantization runs for a corner case in attention module
"""
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv2d(1, 1, 1)
def forward(self, x):
x = self.conv(x)
q, k, v = x.chunk(3, dim=0)
q = q.contiguous().view(-1, 1).transpose(0, 1)
k = k.contiguous().view(-1, 1).transpose(0, 1)
v = v.contiguous().view(-1, 1).transpose(0, 1)
torch._assert(
k.size(1) == 1, "key size should be equal to 1"
)
r = torch.mm(k, v)
return q * k + r
example_inputs = (torch.randn(3, 1, 1, 1),)
m = M().eval()
qconfig_dict = {
"": None,
"object_type": [
(nn.Conv2d, default_qconfig),
]
}
# make sure it runs
m = prepare_fx(m, qconfig_dict, example_inputs=example_inputs)
m(*example_inputs)
m = convert_fx(m)
m(*example_inputs)
def _test_standalone_module(
self,
interface_config,
prepare_count_check,
standalone_prepare_count_check,
convert_count_check,
standalone_convert_count_check):
""" Test standalone module with different quantized input/quantized output
configurations
"""
class StandaloneModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv2d(1, 1, 1)
def forward(self, x):
return self.conv(x)
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv2d(1, 1, 1)
self.standalone = StandaloneModule()
def forward(self, x):
x = self.conv(x)
x = self.standalone(x)
return x
class RefM(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1 = torch.nn.Conv2d(1, 1, 1)
self.conv2 = torch.nn.Conv2d(1, 1, 1)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
return x
example_inputs = (torch.randn(1, 1, 1, 1),)
# instantiate M and RefM and align the parameters
original_m = M().eval()
original_ref_m = RefM().eval()
original_ref_m.conv1.weight = torch.nn.Parameter(original_m.conv.weight.detach())
original_ref_m.conv1.bias = torch.nn.Parameter(original_m.conv.bias.detach())
original_ref_m.conv2.weight = torch.nn.Parameter(original_m.standalone.conv.weight.detach())
original_ref_m.conv2.bias = torch.nn.Parameter(original_m.standalone.conv.bias.detach())
for is_name in [True, False]:
sm_example_inputs = example_inputs
if is_name:
prepare_config = {
"standalone_module_name": [("standalone", None, sm_example_inputs, interface_config, None)]
}
else:
prepare_config = {
"standalone_module_class": [(StandaloneModule, None, sm_example_inputs, interface_config, None)]
}
original_m_copy = copy.deepcopy(original_m)
original_ref_m_copy = copy.deepcopy(original_ref_m)
qconfig_dict = {"": default_qconfig}
# check prepared model
m = prepare_fx(
original_m_copy,
qconfig_dict,
example_inputs=example_inputs,
prepare_custom_config=prepare_config)
# calibration
m(*example_inputs)
self.checkGraphModuleNodes(m, expected_node_occurrence=prepare_count_check)
self.checkGraphModuleNodes(m.standalone, expected_node_occurrence=standalone_prepare_count_check)
# check converted/quantized model
m = convert_fx(m)
self.checkGraphModuleNodes(m, expected_node_occurrence=convert_count_check)
self.checkGraphModuleNodes(m.standalone, expected_node_occurrence=standalone_convert_count_check)
res = m(*example_inputs)
# quantize the reference model
ref_m = prepare_fx(
original_ref_m_copy,
qconfig_dict,
example_inputs=example_inputs,
)
ref_m(*example_inputs)
ref_m = convert_fx(ref_m)
ref_res = ref_m(*example_inputs)
self.assertEqual(res, ref_res)
def test_standalone_module_float_interface(self):
float_interface_config = {
"input_quantized_idxs": [], # float input
"output_quantized_idxs": [], # float output
}
interface_config = float_interface_config
# input and output of first conv, observer for standalone module
# will be inserted in the standalone module itself
prepare_count_check = {
ns.call_module(torch.ao.quantization.MinMaxObserver): 2
}
# for input and output of conv in the standalone module
standalone_prepare_count_check = {
ns.call_module(torch.ao.quantization.MinMaxObserver): 2
}
convert_count_check = {
ns.call_function(torch.quantize_per_tensor) : 1,
ns.call_module(nnq.Conv2d) : 1,
ns.call_method("dequantize") : 1,
}
standalone_convert_count_check = {
# standalone module will take float as input and output
# so we'll see quantize and dequantize in the modoule
ns.call_function(torch.quantize_per_tensor) : 1,
ns.call_module(nnq.Conv2d): 1,
ns.call_method("dequantize") : 1,
}
self._test_standalone_module(
interface_config,
prepare_count_check,
standalone_prepare_count_check,
convert_count_check,
standalone_convert_count_check)
def test_standalone_module_quantized_interface(self):
quantized_interface_config = {
"input_quantized_idxs": [0], # quantized input
"output_quantized_idxs": [0], # quantized output
}
interface_config = quantized_interface_config
# observer for input and output of first conv
prepare_count_check = {
ns.call_module(torch.ao.quantization.MinMaxObserver): 2
}
# for output of conv in the standalone module
standalone_prepare_count_check = {
ns.call_module(torch.ao.quantization.MinMaxObserver): 1
}
convert_count_check = {
# quantizing input for conv
ns.call_function(torch.quantize_per_tensor) : 1,
ns.call_module(nnq.Conv2d) : 1,
# dequantizing output of standalone module
ns.call_method("dequantize") : 1,
}
standalone_convert_count_check = {
# quantization of input happens in parent module
# quantization of output happens in the quantized conv module
ns.call_function(torch.quantize_per_tensor) : 0,
ns.call_module(nnq.Conv2d): 1,
# dequantization for output happens in parent module
ns.call_method("dequantize") : 0,
}
self._test_standalone_module(
interface_config,
prepare_count_check,
standalone_prepare_count_check,
convert_count_check,
standalone_convert_count_check)
@skipIfNoFBGEMM
def test_qconfig_none(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 1, 1)
self.conv2 = nn.Conv2d(1, 1, 1)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
return x
m = M().eval()
qconfig_dict = {"": default_qconfig,
"module_name": [("conv2", None)]}
example_inputs = (torch.randn(1, 1, 1, 1),)
m = prepare_fx(m, qconfig_dict, example_inputs=example_inputs)
m(*example_inputs)
m = convert_fx(m)
m(*example_inputs)
# first conv is quantized, second conv is not quantized
node_list = [
ns.call_function(torch.quantize_per_tensor),
ns.call_module(nnq.Conv2d),
ns.call_method("dequantize"),
ns.call_module(nn.Conv2d),
]
self.checkGraphModuleNodes(m, expected_node_list=node_list)
def test_qconfig_module_type(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(1, 1, 1)
self.linear = nn.Linear(9, 3)
def forward(self, x):
x = self.conv(x)
x = x.reshape((1, -1))
x = self.linear(x)
return x
m = M().eval()
qconfig_dict = {"object_type": [(torch.nn.Conv2d, default_qconfig)]}
example_inputs = (torch.randn(1, 1, 3, 3),)
m = prepare_fx(m, qconfig_dict, example_inputs=example_inputs)
m(*example_inputs)
m = convert_fx(m)
m(*example_inputs)
# conv is quantized, linear is not quantized
node_list = [
ns.call_function(torch.quantize_per_tensor),
ns.call_module(nnq.Conv2d),
ns.call_method("dequantize"),
ns.call_module(nn.Linear),
]
self.checkGraphModuleNodes(m, expected_node_list=node_list)
def test_qconfig_qat_module_type(self):
class LinearRelu(nn.Sequential):
def __init__(self):
super().__init__(
nn.Linear(5, 5),
nn.ReLU(),
)
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.lin_relu = LinearRelu()
self.linear = nn.Linear(5, 5)
def forward(self, x):
x = self.lin_relu(x)
x = self.linear(x)
return x
model = M().train()
qconfig_dict = {
"": None,
"object_type": [
(torch.nn.Linear, default_qat_qconfig),
(torch.nn.ReLU, default_qat_qconfig),
],
}
example_inputs = (torch.rand(5, 5),)
m = prepare_qat_fx(model, qconfig_dict, example_inputs=example_inputs)
m(*example_inputs)
m = convert_fx(m)
m(*example_inputs)
node_list = [
ns.call_function(torch.quantize_per_tensor),
ns.call_module(nniq.LinearReLU),
ns.call_module(nnq.Linear),
ns.call_method("dequantize"),
]
self.checkGraphModuleNodes(m, expected_node_list=node_list)
def test_qconfig_function(self):
class M(torch.nn.Module):
def forward(self, x, y):
return x + y
m = M().eval()
qconfig_dict = {"object_type": [(operator.add, default_qconfig)]}
data = torch.randn(1, 1, 1, 1)
example_inputs = (data, data)
m = prepare_fx(m, qconfig_dict, example_inputs)
m(*example_inputs)
m = convert_fx(m)
m(*example_inputs)
# first conv is quantized, second conv is not quantized
node_list = [
ns.call_function(torch.quantize_per_tensor),
ns.call_function(torch.ops.quantized.add),
ns.call_method("dequantize"),
]
self.checkGraphModuleNodes(m, expected_node_list=node_list)
def test_qconfig_module_name_regex(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 1, 1)
self.conv2 = nn.Conv2d(1, 1, 1)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
return x
m = M().eval()
qconfig_dict = {"module_name_regex": [("conv*", default_qconfig)]}
example_inputs = (torch.randn(1, 1, 1, 1),)
m = prepare_fx(m, qconfig_dict, example_inputs=example_inputs)
m(*example_inputs)
m = convert_fx(m)
m(*example_inputs)
# first conv is quantized, second conv is not quantized
node_list = [
ns.call_function(torch.quantize_per_tensor),
ns.call_module(nnq.Conv2d),
ns.call_module(nnq.Conv2d),
ns.call_method("dequantize"),
]
self.checkGraphModuleNodes(m, expected_node_list=node_list)
def test_qconfig_precedence(self):
for device in get_supported_device_types():
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(1, 1)
self.conv = nn.Conv2d(1, 1, 1)
self.module_conv1 = nn.Conv2d(1, 1, 1)
self.module_conv2 = nn.Conv2d(1, 1, 1)
def forward(self, x):
# global
x = self.linear(x)
# global + object_type --> object_type
x = self.conv(x)
# global + object_type + module_name_regex --> module_name_regex
x = self.module_conv1(x)
# global + object_type + module_name_regex + module_name --> module_name
x = self.module_conv2(x)
return x
m = M().to(device).eval()
global_qconfig = default_qconfig
object_type_qconfig = default_dynamic_qconfig
module_name_regex_qconfig = float16_dynamic_qconfig
module_name_qconfig = default_qat_qconfig
qconfig_dict = {
"": global_qconfig,
"object_type": [(nn.Conv2d, object_type_qconfig)],
"module_name_regex": [("module_conv*", module_name_regex_qconfig)],
"module_name": [("module_conv2", module_name_qconfig)]}
m_prep = prepare_fx(m, qconfig_dict, example_inputs=(torch.randn(1, 1),))
self.assertEqual(m_prep.linear.qconfig.activation.p.func, global_qconfig.activation.p.func)
self.assertEqual(m_prep.linear.qconfig.weight.p.func, global_qconfig.weight.p.func)
self.assertEqual(m_prep.conv.qconfig.activation.p.func, object_type_qconfig.activation.p.func)
self.assertEqual(m_prep.conv.qconfig.weight.p.func, object_type_qconfig.weight.p.func)
self.assertEqual(m_prep.module_conv1.qconfig.activation.p.func, module_name_regex_qconfig.activation.p.func)
self.assertEqual(m_prep.module_conv1.qconfig.weight.p.func, module_name_regex_qconfig.weight.p.func)
self.assertEqual(m_prep.module_conv2.qconfig.activation.p.func, module_name_qconfig.activation.p.func)
self.assertEqual(m_prep.module_conv2.qconfig.weight.p.func, module_name_qconfig.weight.p.func)
def test_qconfig_module_name_object_type_order(self):
class M1(torch.nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(1, 1)
self.fc2 = nn.Linear(1, 1)
def forward(self, x):
x = self.fc1(x)
x = self.fc2(x)
x = torch.add(x, x)
x = torch.add(x, x)
return x
class M2(torch.nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(1, 1)
self.fc2 = nn.Linear(1, 1)
self.m1 = M1()
def forward(self, x):
x = self.fc1(x)
x = self.fc2(x)
x = torch.add(x, x)
x = torch.add(x, x)
x = self.m1(x)
return x
class M3(torch.nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(1, 1)
self.fc2 = nn.Linear(1, 1)
self.m2 = M2()
def forward(self, x):
x = self.fc1(x)
x = self.fc2(x)
x = torch.add(x, x)
x = torch.add(x, x)
x = self.m2(x)
return x
m = M3().eval()
qconfig_dict = {
"module_name_object_type_order": [
# test various FQNs: global, single child, multiple children
("", nn.Linear, 0, torch.ao.quantization.default_qconfig),
("", torch.add, 0, torch.ao.quantization.default_qconfig),
("m2", nn.Linear, 1, torch.ao.quantization.default_qconfig),
("m2", torch.add, 1, torch.ao.quantization.default_qconfig),
("m2.m1", nn.Linear, 0, torch.ao.quantization.default_qconfig),
("m2.m1", torch.add, 0, torch.ao.quantization.default_qconfig),
],
}
example_inputs = (torch.randn(1, 1, 1, 1),)
m = prepare_fx(m, qconfig_dict, example_inputs=example_inputs)
m(*example_inputs)
m = convert_fx(m)
m(*example_inputs)
node_list = [
# m3
ns.call_function(torch.quantize_per_tensor),
ns.call_module(nnq.Linear),
ns.call_method("dequantize"),
ns.call_module(nn.Linear),
ns.call_function(torch.quantize_per_tensor),
ns.call_function(torch.ops.quantized.add),
ns.call_method("dequantize"),
ns.call_function(torch.add),
# m2
ns.call_module(nn.Linear),
ns.call_function(torch.quantize_per_tensor),
ns.call_module(nnq.Linear),
ns.call_method("dequantize"),
ns.call_function(torch.add),
ns.call_function(torch.quantize_per_tensor),
ns.call_function(torch.ops.quantized.add),
# m1
ns.call_module(nnq.Linear),
ns.call_method("dequantize"),
ns.call_module(nn.Linear),
ns.call_function(torch.quantize_per_tensor),
ns.call_function(torch.ops.quantized.add),
ns.call_method("dequantize"),
ns.call_function(torch.add),
]
self.checkGraphModuleNodes(m, expected_node_list=node_list)
# test that function order overrides global qconfig
class M4(torch.nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(1, 1)
self.fc2 = nn.Linear(1, 1)
def forward(self, x):
x = self.fc1(x)
x = self.fc2(x)
x = torch.add(x, x)
x = torch.add(x, x)
return x
m = M4().eval()
qconfig_dict = {
"": torch.ao.quantization.default_qconfig,
"module_name_object_type_order": [
("", nn.Linear, 1, None),
("", torch.add, 1, None),
],
}
example_inputs = (torch.randn(1, 1, 1, 1),)
m = prepare_fx(m, qconfig_dict, example_inputs=example_inputs)
m(*example_inputs)
m = convert_fx(m)
m(*example_inputs)
node_list = [
ns.call_function(torch.quantize_per_tensor),
ns.call_module(nnq.Linear),
ns.call_method("dequantize"),
ns.call_module(nn.Linear),
ns.call_function(torch.quantize_per_tensor),
ns.call_function(torch.ops.quantized.add),
ns.call_method("dequantize"),
ns.call_function(torch.add),
]
self.checkGraphModuleNodes(m, expected_node_list=node_list)
@override_qengines
def test_qconfig_dict_with_fused_modules(self):
class LinearReLUModel(torch.nn.Module):
def __init__(self, relu):
super().__init__()
self.linear = torch.nn.Linear(3, 3)
self.relu = relu
def forward(self, x):
x = self.linear(x)
x = self.relu(x)
return x
class ConvReLUModel(torch.nn.Module):
def __init__(self, relu):
super().__init__()
self.conv = torch.nn.Conv1d(3, 3, 3)
self.relu = relu
def forward(self, x):
x = self.conv(x)
x = self.relu(x)
return x
class ConvBnReLUModel(torch.nn.Module):
def __init__(self, relu):
super().__init__()
self.conv = torch.nn.Conv1d(3, 3, 3)
self.bn = torch.nn.BatchNorm1d(3)
self.relu = relu
def forward(self, x):
x = self.conv(x)
x = self.bn(x)
x = self.relu(x)
return x
for model in [LinearReLUModel, ConvReLUModel, ConvBnReLUModel]:
for relu in [torch.nn.ReLU(), torch.nn.functional.relu, torch.relu]:
m = model(relu).eval()
qengine = torch.backends.quantized.engine
qconfig_dict = torch.ao.quantization.get_default_qconfig_mapping(qengine)
# should not crash as in https://github.com/pytorch/pytorch/issues/75825
prepare_fx(m, qconfig_dict, example_inputs=(torch.randn(1, 3, 3, 3),))
# TODO: move QConfigMapping tests to test/quantization/core
def test_qconfig_mapping_set_global(self):
qconfig = get_default_qconfig()
qconfig_mapping = QConfigMapping()
self.assertEqual(qconfig_mapping.global_qconfig, None)
qconfig_mapping.set_global(qconfig)
self.assertEqual(qconfig_mapping.global_qconfig, qconfig)
def test_qconfig_mapping_set_object_type(self):
qconfig1 = get_default_qconfig()
qconfig2 = get_default_qconfig()
qconfig3 = get_default_qconfig()
self.assertNotEqual(qconfig1, qconfig2)
self.assertNotEqual(qconfig1, qconfig3)
qconfig_mapping = QConfigMapping()
self.assertEqual(len(qconfig_mapping.object_type_qconfigs), 0)
# Insert some entries
qconfig_mapping.set_object_type(torch.nn.Linear, qconfig1)
qconfig_mapping.set_object_type(torch.nn.ReLU, qconfig2)
self.assertEqual(len(qconfig_mapping.object_type_qconfigs), 2)
self.assertEqual(qconfig_mapping.object_type_qconfigs[torch.nn.Linear], qconfig1)
self.assertEqual(qconfig_mapping.object_type_qconfigs[torch.nn.ReLU], qconfig2)
# Override existing key
qconfig_mapping.set_object_type(torch.nn.Linear, qconfig3)
self.assertEqual(qconfig_mapping.object_type_qconfigs[torch.nn.Linear], qconfig3)
self.assertEqual(qconfig_mapping.object_type_qconfigs[torch.nn.ReLU], qconfig2)
self.assertEqual(_get_object_type_qconfig(qconfig_mapping, torch.nn.Linear, None), qconfig3)
self.assertEqual(_get_object_type_qconfig(qconfig_mapping, torch.nn.ReLU, None), qconfig2)
self.assertEqual(_get_object_type_qconfig(qconfig_mapping, "nomatch", None), None)
def test_qconfig_mapping_set_module_name_regex(self):
qconfig1 = get_default_qconfig()
qconfig2 = get_default_qconfig()
qconfig3 = get_default_qconfig()
self.assertNotEqual(qconfig1, qconfig2)
self.assertNotEqual(qconfig1, qconfig3)
qconfig_mapping = QConfigMapping()
self.assertEqual(len(qconfig_mapping.module_name_regex_qconfigs), 0)
# Insert some entries
qconfig_mapping.set_module_name_regex("foo.*bar", qconfig1)
qconfig_mapping.set_module_name_regex("foo.*", qconfig2)
self.assertEqual(len(qconfig_mapping.module_name_regex_qconfigs), 2)
self.assertEqual(qconfig_mapping.module_name_regex_qconfigs["foo.*bar"], qconfig1)
self.assertEqual(qconfig_mapping.module_name_regex_qconfigs["foo.*"], qconfig2)
# Override existing key
qconfig_mapping.set_module_name_regex("foo.*bar", qconfig3)
self.assertEqual(qconfig_mapping.module_name_regex_qconfigs["foo.*bar"], qconfig3)
self.assertEqual(qconfig_mapping.module_name_regex_qconfigs["foo.*"], qconfig2)
self.assertEqual(_get_module_name_regex_qconfig(qconfig_mapping, "foo123bar", None), qconfig3)
self.assertEqual(_get_module_name_regex_qconfig(qconfig_mapping, "foobar", None), qconfig3)
self.assertEqual(_get_module_name_regex_qconfig(qconfig_mapping, "foobaz", None), qconfig2)
self.assertEqual(_get_module_name_regex_qconfig(qconfig_mapping, "foo", None), qconfig2)
self.assertEqual(_get_module_name_regex_qconfig(qconfig_mapping, "nomatch", None), None)
def test_qconfig_mapping_set_module_name(self):
qconfig1 = get_default_qconfig()
qconfig2 = get_default_qconfig()
qconfig3 = get_default_qconfig()
self.assertNotEqual(qconfig1, qconfig2)
self.assertNotEqual(qconfig1, qconfig3)
qconfig_mapping = QConfigMapping()
self.assertEqual(len(qconfig_mapping.module_name_qconfigs), 0)
# Insert some entries
qconfig_mapping.set_module_name("mod1", qconfig1)
qconfig_mapping.set_module_name("mod2", qconfig2)
self.assertEqual(len(qconfig_mapping.module_name_qconfigs), 2)
self.assertEqual(qconfig_mapping.module_name_qconfigs["mod1"], qconfig1)
self.assertEqual(qconfig_mapping.module_name_qconfigs["mod2"], qconfig2)
# Override existing key
qconfig_mapping.set_module_name("mod1", qconfig3)
self.assertEqual(qconfig_mapping.module_name_qconfigs["mod1"], qconfig3)
self.assertEqual(qconfig_mapping.module_name_qconfigs["mod2"], qconfig2)
self.assertEqual(_get_module_name_qconfig(qconfig_mapping, "mod1", None), qconfig3)
self.assertEqual(_get_module_name_qconfig(qconfig_mapping, "mod2", None), qconfig2)
self.assertEqual(_get_module_name_qconfig(qconfig_mapping, "nomatch", None), None)
def test_qconfig_mapping_set_module_name_object_type_order(self):
qconfig1 = get_default_qconfig()
qconfig2 = get_default_qconfig()
qconfig3 = get_default_qconfig()
self.assertNotEqual(qconfig1, qconfig2)
self.assertNotEqual(qconfig1, qconfig3)
qconfig_mapping = QConfigMapping()
self.assertEqual(len(qconfig_mapping.module_name_object_type_order_qconfigs), 0)
# Insert some entries
qconfig_mapping.set_module_name_object_type_order("mod1", torch.nn.Linear, 0, qconfig1)
qconfig_mapping.set_module_name_object_type_order("mod2", torch.nn.ReLU, 1, qconfig2)
self.assertEqual(len(qconfig_mapping.module_name_object_type_order_qconfigs), 2)
key1 = ("mod1", torch.nn.Linear, 0)
key2 = ("mod2", torch.nn.ReLU, 1)
self.assertEqual(list(qconfig_mapping.module_name_object_type_order_qconfigs)[0], key1)
self.assertEqual(list(qconfig_mapping.module_name_object_type_order_qconfigs)[1], key2)
self.assertEqual(qconfig_mapping.module_name_object_type_order_qconfigs[key1], qconfig1)
self.assertEqual(qconfig_mapping.module_name_object_type_order_qconfigs[key2], qconfig2)
self.assertEqual(_maybe_adjust_qconfig_for_module_name_object_type_order(
qconfig_mapping, "mod1", torch.nn.Linear, 0, None), qconfig1)
self.assertEqual(_maybe_adjust_qconfig_for_module_name_object_type_order(
qconfig_mapping, "mod2", torch.nn.ReLU, 1, None), qconfig2)
# Override existing key
qconfig_mapping.set_module_name_object_type_order("mod1", torch.nn.Linear, 0, qconfig3)
self.assertEqual(len(qconfig_mapping.module_name_object_type_order_qconfigs), 2)
self.assertEqual(list(qconfig_mapping.module_name_object_type_order_qconfigs)[0], key1)
self.assertEqual(list(qconfig_mapping.module_name_object_type_order_qconfigs)[1], key2)
self.assertEqual(qconfig_mapping.module_name_object_type_order_qconfigs[key1], qconfig3)
self.assertEqual(qconfig_mapping.module_name_object_type_order_qconfigs[key2], qconfig2)
self.assertEqual(_maybe_adjust_qconfig_for_module_name_object_type_order(
qconfig_mapping, "mod1", torch.nn.Linear, 0, None), qconfig3)
self.assertEqual(_maybe_adjust_qconfig_for_module_name_object_type_order(
qconfig_mapping, "mod2", torch.nn.ReLU, 1, None), qconfig2)
# No match
self.assertEqual(_maybe_adjust_qconfig_for_module_name_object_type_order(
qconfig_mapping, "mod123", torch.nn.Linear, 0, None), None)
self.assertEqual(_maybe_adjust_qconfig_for_module_name_object_type_order(
qconfig_mapping, "mod1", torch.nn.Linear, 35, None), None)
self.assertEqual(_maybe_adjust_qconfig_for_module_name_object_type_order(
qconfig_mapping, "mod2", torch.nn.Conv2d, 1, None), None)
def _get_qconfig_dict_for_qconfig_mapping_test(self, global_qconfig, qconfig1, qconfig2):
"""
Return a dummy qconfig_dict to test QConfigMapping's to_dict and from_dict methods.
"""
return {
_GLOBAL_DICT_KEY: global_qconfig,
_OBJECT_TYPE_DICT_KEY: [
(torch.nn.Linear, qconfig1),
(torch.nn.ReLU, qconfig2),
],
_MODULE_NAME_REGEX_DICT_KEY: [
("foo.*bar", qconfig1),
("foo.*", qconfig2),
],
_MODULE_NAME_DICT_KEY: [
("bazbaz", qconfig1),
("borbor", qconfig2),
],
_MODULE_NAME_OBJECT_TYPE_ORDER_DICT_KEY: [
("bazbaz", torch.nn.Linear, 0, qconfig1),
("foofoo", torch.nn.ReLU, 1, qconfig2),
],
}
with self.assertRaises(ValueError) as context:
m = prepare_fx(m, qconfig_dict, example_inputs=(torch.randn(1, 3, 3, 3),))
self.assertTrue(
'Expected qconfig_dict to have the following keys:' in str(context.exception)
)
self.assertTrue('But found \'object_typo\' instead.' in str(context.exception))
def test_qconfig_mapping_from_dict(self):
global_qconfig = QConfig(123, "global")
qconfig1 = QConfig(1, "one")
qconfig2 = QConfig(2, "two")
qconfig_dict = self._get_qconfig_dict_for_qconfig_mapping_test(global_qconfig, qconfig1, qconfig2)
qconfig_dict["undefined_dict_key"] = [(123, qconfig1), (234, qconfig2)]
qconfig_mapping = QConfigMapping.from_dict(qconfig_dict)
self.assertEqual(qconfig_mapping.global_qconfig, global_qconfig)
self.assertEqual(qconfig_mapping.object_type_qconfigs, OrderedDict({
torch.nn.Linear: qconfig1,
torch.nn.ReLU: qconfig2,
}))
self.assertEqual(qconfig_mapping.module_name_regex_qconfigs, OrderedDict({
"foo.*bar": qconfig1,
"foo.*": qconfig2,
}))
self.assertEqual(qconfig_mapping.module_name_qconfigs, OrderedDict({
"bazbaz": qconfig1,
"borbor": qconfig2,
}))
self.assertEqual(qconfig_mapping.module_name_object_type_order_qconfigs, OrderedDict({
("bazbaz", torch.nn.Linear, 0): qconfig1,
("foofoo", torch.nn.ReLU, 1): qconfig2,
}))
def test_qconfig_mapping_to_dict(self):
global_qconfig = QConfig(123, "global")
qconfig1 = QConfig(1, "one")
qconfig2 = QConfig(2, "two")
qconfig_mapping = QConfigMapping().set_global(global_qconfig) \
.set_object_type(torch.nn.Linear, qconfig1) \
.set_object_type(torch.nn.ReLU, qconfig2) \
.set_module_name_regex("foo.*bar", qconfig1) \
.set_module_name_regex("foo.*", qconfig2) \
.set_module_name("bazbaz", qconfig1) \
.set_module_name("borbor", qconfig2) \
.set_module_name_object_type_order("bazbaz", torch.nn.Linear, 0, qconfig1) \
.set_module_name_object_type_order("foofoo", torch.nn.ReLU, 1, qconfig2)
qconfig_dict = self._get_qconfig_dict_for_qconfig_mapping_test(global_qconfig, qconfig1, qconfig2)
self.assertEqual(qconfig_mapping.to_dict(), qconfig_dict)
def test_qconfig_mapping_repr(self):
self.assertTrue(isinstance(get_default_qconfig_mapping().__repr__(), str))
def test_default_qconfig_mapping_override_global(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv2d(1, 1, 1)
def forward(self, x):
return self.conv(x)
m = M().eval()
my_qconfig = QConfig(activation=MinMaxObserver, weight=default_weight_observer)
qconfig_mapping = get_default_qconfig_mapping()
# Override global qconfig
old_global_qconfig = qconfig_mapping.global_qconfig
qconfig_mapping.set_global(my_qconfig)
# Verify the correct qconfig was used
example_inputs = (torch.randn(1, 1, 1, 1),)
m = prepare_fx(m, qconfig_mapping, example_inputs)
self.assertTrue(isinstance(old_global_qconfig.activation(), HistogramObserver))
self.assertTrue(isinstance(my_qconfig.activation(), MinMaxObserver))
self.assertTrue(hasattr(m, "activation_post_process_0"))
self.assertTrue(hasattr(m, "activation_post_process_1"))
self.assertTrue(isinstance(m.activation_post_process_0, MinMaxObserver))
self.assertTrue(isinstance(m.activation_post_process_1, MinMaxObserver))
# Dummy classes for PrepareCustomConfig testing
class _DummyStandaloneModule:
pass
class _DummyFloatModule:
pass
class _DummyObservedModule:
pass
class _DummyQuantizedModule:
pass
class _DummyNonTraceableModule1:
pass
class _DummyNonTraceableModule2:
pass
def test_prepare_custom_config_set_standalone_module_name(self):
qconfig_mapping = QConfigMapping()
example_inputs = (torch.randn(3),)
child_prepare_custom_config = PrepareCustomConfig()
backend_config = BackendConfig("my_backend")
config_entry = StandaloneModuleConfigEntry(
qconfig_mapping, example_inputs, child_prepare_custom_config, backend_config)
prepare_custom_config = PrepareCustomConfig()
self.assertEqual(len(prepare_custom_config.standalone_module_names), 0)
prepare_custom_config.set_standalone_module_name(
"module1", qconfig_mapping, example_inputs, child_prepare_custom_config, backend_config)
self.assertEqual(list(prepare_custom_config.standalone_module_names.keys()), ["module1"])
self.assertEqual(prepare_custom_config.standalone_module_names["module1"], config_entry)
def test_prepare_custom_config_set_standalone_module_class(self):
qconfig_mapping = QConfigMapping()
example_inputs = (torch.randn(3),)
child_prepare_custom_config = PrepareCustomConfig()
backend_config = BackendConfig("my_backend")
config_entry = StandaloneModuleConfigEntry(
qconfig_mapping, example_inputs, child_prepare_custom_config, backend_config)
prepare_custom_config = PrepareCustomConfig()
self.assertEqual(len(prepare_custom_config.standalone_module_classes), 0)
prepare_custom_config.set_standalone_module_class(
self._DummyStandaloneModule, qconfig_mapping, example_inputs, child_prepare_custom_config, backend_config)
self.assertEqual(len(prepare_custom_config.standalone_module_classes), 1)
self.assertTrue(self._DummyStandaloneModule in prepare_custom_config.standalone_module_classes)
self.assertEqual(prepare_custom_config.standalone_module_classes[self._DummyStandaloneModule], config_entry)
def test_prepare_custom_config_set_float_to_observed_mapping(self):
prepare_custom_config = PrepareCustomConfig()
self.assertEqual(len(prepare_custom_config.float_to_observed_mapping), 0)
prepare_custom_config.set_float_to_observed_mapping(self._DummyFloatModule, self._DummyObservedModule, QuantType.STATIC)
self.assertEqual(len(prepare_custom_config.float_to_observed_mapping), 1)
self.assertEqual(list(prepare_custom_config.float_to_observed_mapping.keys()), [QuantType.STATIC])
self.assertEqual(len(prepare_custom_config.float_to_observed_mapping[QuantType.STATIC]), 1)
self.assertTrue(self._DummyFloatModule in prepare_custom_config.float_to_observed_mapping[QuantType.STATIC])
self.assertEqual(prepare_custom_config.float_to_observed_mapping[QuantType.STATIC][self._DummyFloatModule],
self._DummyObservedModule)
def test_prepare_custom_config_set_non_traceable_module_names(self):
prepare_custom_config = PrepareCustomConfig()
self.assertEqual(len(prepare_custom_config.non_traceable_module_names), 0)
prepare_custom_config.set_non_traceable_module_names(["module1", "module2"])
self.assertEqual(prepare_custom_config.non_traceable_module_names, ["module1", "module2"])
def test_prepare_custom_config_set_non_traceable_module_classes(self):
prepare_custom_config = PrepareCustomConfig()
self.assertEqual(len(prepare_custom_config.non_traceable_module_classes), 0)
prepare_custom_config.set_non_traceable_module_classes([self._DummyNonTraceableModule1, self._DummyNonTraceableModule2])
self.assertEqual(prepare_custom_config.non_traceable_module_classes,
[self._DummyNonTraceableModule1, self._DummyNonTraceableModule2])
def test_prepare_custom_config_set_input_quantized_indexes(self):
prepare_custom_config = PrepareCustomConfig()
self.assertEqual(len(prepare_custom_config.input_quantized_indexes), 0)
prepare_custom_config.set_input_quantized_indexes([0, 1])
self.assertEqual(prepare_custom_config.input_quantized_indexes, [0, 1])
def test_prepare_custom_config_set_output_quantized_indexes(self):
prepare_custom_config = PrepareCustomConfig()
self.assertEqual(len(prepare_custom_config.output_quantized_indexes), 0)
prepare_custom_config.set_output_quantized_indexes([0, 1])
self.assertEqual(prepare_custom_config.output_quantized_indexes, [0, 1])
def test_prepare_custom_config_set_preserved_attributes(self):
prepare_custom_config = PrepareCustomConfig()
self.assertEqual(len(prepare_custom_config.preserved_attributes), 0)
prepare_custom_config.set_preserved_attributes(["attr1", "attr2"])
self.assertEqual(prepare_custom_config.preserved_attributes, ["attr1", "attr2"])
def _get_dummy_prepare_custom_config_dict(self):
"""
Return a dummy prepare_custom_config_dict to test PrepareCustomConfig's to_dict and from_dict methods.
"""
return {
STANDALONE_MODULE_NAME_DICT_KEY: [(
"module1",
QConfigMapping(),
(torch.randn(3),),
PrepareCustomConfig(),
BackendConfig("my_backend"),
)],
STANDALONE_MODULE_CLASS_DICT_KEY: [(
self._DummyStandaloneModule,
QConfigMapping(),
(torch.randn(10),),
PrepareCustomConfig(),
BackendConfig("my_backend"),
)],
FLOAT_TO_OBSERVED_DICT_KEY: {
"static": {
self._DummyFloatModule: self._DummyObservedModule
},
},
NON_TRACEABLE_MODULE_NAME_DICT_KEY: ["module2", "module3"],
NON_TRACEABLE_MODULE_CLASS_DICT_KEY: [self._DummyNonTraceableModule1, self._DummyNonTraceableModule2],
INPUT_QUANTIZED_INDEXES_DICT_KEY: [0, 1],
OUTPUT_QUANTIZED_INDEXES_DICT_KEY: [0, 1],
PRESERVED_ATTRIBUTES_DICT_KEY: ["attr1", "attr2"]
}
def test_prepare_custom_config_from_dict(self):
prepare_custom_config_dict = self._get_dummy_prepare_custom_config_dict()
(sm_name, qm1, ei1, pcc1, bcd1) = prepare_custom_config_dict[STANDALONE_MODULE_NAME_DICT_KEY][0]
(sm_class, qm2, ei2, pcc2, bcd2) = prepare_custom_config_dict[STANDALONE_MODULE_CLASS_DICT_KEY][0]
sm_config_entry1 = StandaloneModuleConfigEntry(qm1, ei1, pcc1, bcd1)
sm_config_entry2 = StandaloneModuleConfigEntry(qm2, ei2, pcc2, bcd2)
prepare_custom_config = PrepareCustomConfig.from_dict(prepare_custom_config_dict)
# Standalone modules
self.assertEqual(len(prepare_custom_config.standalone_module_names), 1)
self.assertTrue(sm_name in prepare_custom_config.standalone_module_names)
self.assertEqual(prepare_custom_config.standalone_module_names[sm_name], sm_config_entry1)
self.assertEqual(len(prepare_custom_config.standalone_module_classes), 1)
self.assertTrue(sm_class in prepare_custom_config.standalone_module_classes)
self.assertEqual(prepare_custom_config.standalone_module_classes[sm_class], sm_config_entry2)
# Float to observed mapping
self.assertEqual(len(prepare_custom_config.float_to_observed_mapping), 1)
self.assertEqual(list(prepare_custom_config.float_to_observed_mapping.keys()), [QuantType.STATIC])
self.assertEqual(len(prepare_custom_config.float_to_observed_mapping[QuantType.STATIC]), 1)
self.assertTrue(self._DummyFloatModule in prepare_custom_config.float_to_observed_mapping[QuantType.STATIC])
self.assertEqual(prepare_custom_config.float_to_observed_mapping[QuantType.STATIC][self._DummyFloatModule],
self._DummyObservedModule)
# Other
self.assertEqual(prepare_custom_config.non_traceable_module_names, ["module2", "module3"])
self.assertEqual(prepare_custom_config.non_traceable_module_classes,
[self._DummyNonTraceableModule1, self._DummyNonTraceableModule2])
self.assertEqual(prepare_custom_config.input_quantized_indexes, [0, 1])
self.assertEqual(prepare_custom_config.output_quantized_indexes, [0, 1])
self.assertEqual(prepare_custom_config.preserved_attributes, ["attr1", "attr2"])
def test_prepare_custom_config_to_dict(self):
prepare_custom_config_dict = self._get_dummy_prepare_custom_config_dict()
(sm_name, qm1, ei1, pcc1, bcd1) = prepare_custom_config_dict[STANDALONE_MODULE_NAME_DICT_KEY][0]
(sm_class, qm2, ei2, pcc2, bcd2) = prepare_custom_config_dict[STANDALONE_MODULE_CLASS_DICT_KEY][0]
prepare_custom_config = PrepareCustomConfig() \
.set_standalone_module_name(sm_name, qm1, ei1, pcc1, bcd1) \
.set_standalone_module_class(sm_class, qm2, ei2, pcc2, bcd2) \
.set_float_to_observed_mapping(self._DummyFloatModule, self._DummyObservedModule) \
.set_non_traceable_module_names(["module2", "module3"]) \
.set_non_traceable_module_classes([self._DummyNonTraceableModule1, self._DummyNonTraceableModule2]) \
.set_input_quantized_indexes([0, 1]) \
.set_output_quantized_indexes([0, 1]) \
.set_preserved_attributes(["attr1", "attr2"])
# PrepareCustomConfig.to_dict also converts internal QConfigMappings and PrepareCustomConfigs to dicts
prepare_custom_config_dict[STANDALONE_MODULE_NAME_DICT_KEY][0] = (sm_name, qm1.to_dict(), ei1, pcc1.to_dict(), bcd1)
prepare_custom_config_dict[STANDALONE_MODULE_CLASS_DICT_KEY][0] = (sm_class, qm2.to_dict(), ei2, pcc2.to_dict(), bcd2)
self.assertEqual(prepare_custom_config.to_dict(), prepare_custom_config_dict)
def test_convert_custom_config_set_observed_to_quantized_mapping(self):
convert_custom_config = ConvertCustomConfig()
self.assertEqual(len(convert_custom_config.observed_to_quantized_mapping), 0)
convert_custom_config.set_observed_to_quantized_mapping(
self._DummyObservedModule, self._DummyQuantizedModule, QuantType.STATIC)
self.assertEqual(len(convert_custom_config.observed_to_quantized_mapping), 1)
self.assertEqual(list(convert_custom_config.observed_to_quantized_mapping.keys()), [QuantType.STATIC])
self.assertTrue(self._DummyObservedModule in convert_custom_config.observed_to_quantized_mapping[QuantType.STATIC])
self.assertEqual(convert_custom_config.observed_to_quantized_mapping[QuantType.STATIC][self._DummyObservedModule],
self._DummyQuantizedModule)
def test_convert_custom_config_set_preserved_attributes(self):
convert_custom_config = ConvertCustomConfig()
self.assertEqual(len(convert_custom_config.preserved_attributes), 0)
convert_custom_config.set_preserved_attributes(["attr1", "attr2"])
self.assertEqual(convert_custom_config.preserved_attributes, ["attr1", "attr2"])
def _get_dummy_convert_custom_config_dict(self):
"""
Return a dummy convert_custom_config_dict to test ConvertCustomConfig's to_dict and from_dict methods.
"""
return {
OBSERVED_TO_QUANTIZED_DICT_KEY: {
"static": {
self._DummyObservedModule: self._DummyQuantizedModule
},
},
PRESERVED_ATTRIBUTES_DICT_KEY: ["attr1", "attr2"]
}
def test_convert_custom_config_from_dict(self):
convert_custom_config_dict = self._get_dummy_convert_custom_config_dict()
convert_custom_config = ConvertCustomConfig.from_dict(convert_custom_config_dict)
self.assertEqual(len(convert_custom_config.observed_to_quantized_mapping), 1)
self.assertEqual(list(convert_custom_config.observed_to_quantized_mapping.keys()), [QuantType.STATIC])
self.assertEqual(len(convert_custom_config.observed_to_quantized_mapping[QuantType.STATIC]), 1)
self.assertTrue(self._DummyObservedModule in convert_custom_config.observed_to_quantized_mapping[QuantType.STATIC])
self.assertEqual(convert_custom_config.observed_to_quantized_mapping[QuantType.STATIC][self._DummyObservedModule],
self._DummyQuantizedModule)
self.assertEqual(convert_custom_config.preserved_attributes, ["attr1", "attr2"])
def test_convert_custom_config_to_dict(self):
convert_custom_config = ConvertCustomConfig() \
.set_observed_to_quantized_mapping(self._DummyObservedModule, self._DummyQuantizedModule) \
.set_preserved_attributes(["attr1", "attr2"])
self.assertEqual(convert_custom_config.to_dict(), self._get_dummy_convert_custom_config_dict())
def test_fuse_custom_config_set_preserved_attributes(self):
fuse_custom_config = FuseCustomConfig()
self.assertEqual(len(fuse_custom_config.preserved_attributes), 0)
fuse_custom_config.set_preserved_attributes(["attr1", "attr2"])
self.assertEqual(fuse_custom_config.preserved_attributes, ["attr1", "attr2"])
def test_fuse_custom_config_from_dict(self):
fuse_custom_config_dict = {PRESERVED_ATTRIBUTES_DICT_KEY: ["attr1", "attr2"]}
fuse_custom_config = FuseCustomConfig.from_dict(fuse_custom_config_dict)
self.assertEqual(fuse_custom_config.preserved_attributes, ["attr1", "attr2"])
def test_fuse_custom_config_to_dict(self):
fuse_custom_config_dict = {PRESERVED_ATTRIBUTES_DICT_KEY: ["attr1", "attr2"]}
fuse_custom_config = FuseCustomConfig().set_preserved_attributes(["attr1", "attr2"])
self.assertEqual(fuse_custom_config.to_dict(), fuse_custom_config_dict)
def test_remove_qconfig(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.avg_pool = torch.nn.AvgPool2d(1)
def forward(self, x):
return self.avg_pool(x)
m = M().eval()
qconfig_dict = {'': default_qconfig}
example_inputs = (torch.randn(1, 1, 1, 1),)
m = prepare_fx(m, qconfig_dict, example_inputs=example_inputs)
m(*example_inputs)
m = convert_fx(m)
m(*example_inputs)
for name, module in m.named_modules():
self.assertFalse(hasattr(module, 'qconfig'),
'qconfig is not removed for ' + name)
def test_return_none(self):
class M(torch.nn.Module):
def forward(self, x):
pass
m = M().eval()
qconfig_dict = {'': torch.ao.quantization.default_qconfig}
m = prepare_fx(m, qconfig_dict, example_inputs=(torch.randn(1),))
m = convert_fx(m)
def test_default_quant_after_none_qconfig(self):
""" Make sure default quant is inserted properly"""
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1 = torch.nn.Conv2d(1, 1, 1)
self.conv2 = torch.nn.Conv2d(1, 1, 1)
def forward(self, x):
x = self.conv1(x)
x = x.transpose(1, 2)
x = self.conv2(x)
m = M().eval()
qconfig_dict = {
"": default_qconfig,
"module_name": [
("conv1", None)
]
}
m = prepare_fx(m, qconfig_dict, example_inputs=(torch.randn(1, 1, 1, 1),))
m = convert_fx(m)
def test_qconfig_for_call_method(self):
class Sub(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv2d(1, 1, 1)
def forward(self, x):
x = x.transpose(2, 3)
x = self.conv(x)
return x.transpose(2, 3)
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.sub = Sub()
self.conv1 = torch.nn.Conv2d(1, 1, 1)
self.conv2 = torch.nn.Conv2d(1, 1, 1)
def forward(self, x):
x = self.conv1(x)
x = self.sub(x)
x = self.conv2(x)
return x.transpose(2, 3)
qconfig_dict1 = {"": default_qconfig, "module_name": [("sub", None)]}
# since sub is configured to have qconfig None, we should dequantize the output
# of self.conv1 and quantize the input of self.conv2
# dequantize after conv2 should happen after transpose since
# it is configured with default_qconfig
# nodes in Sub module instance is not quantized
node_list1 = [
ns.call_function(torch.quantize_per_tensor),
ns.call_module(nnq.Conv2d),
ns.call_method("dequantize"),
ns.call_method("transpose"),
ns.call_module(nn.Conv2d),
ns.call_method("transpose"),
ns.call_function(torch.quantize_per_tensor),
ns.call_module(nnq.Conv2d),
ns.call_method("transpose"),
ns.call_method("dequantize")
]
qconfig_dict2 = {"": None, "module_name": [("sub", default_qconfig)]}
# Only nodes in Sub module instance are quantized
# the first transpose is not quantized because the input is not quantized
node_list2 = [
ns.call_module(nn.Conv2d),
ns.call_function(torch.quantize_per_tensor),
ns.call_method("transpose"),
ns.call_module(nnq.Conv2d),
ns.call_method("transpose"),
ns.call_method("dequantize"),
ns.call_module(nn.Conv2d),
ns.call_method("transpose"),
]
for qconfig_dict, node_list in [
(qconfig_dict1, node_list1),
(qconfig_dict2, node_list2)
]:
example_inputs = (torch.randn(2, 1, 3, 3),)
m = M().eval()
m = prepare_fx(m, qconfig_dict, example_inputs=example_inputs)
m(torch.randn(2, 1, 3, 3))
m = convert_fx(m)
self.checkGraphModuleNodes(m, expected_node_list=node_list)
# make sure it runs
m(*example_inputs)
def test_qconfig_for_call_func(self):
class Linear(torch.nn.Module):
def __init__(self):
super().__init__()
self.w = torch.ones(5, 5)
self.b = torch.zeros(5)
def forward(self, x):
return torch.nn.functional.linear(x, self.w, self.b)
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.mods1 = torch.nn.Sequential(
Linear(),
Linear()
)
self.mods2 = Linear()
def forward(self, x):
x = self.mods1(x)
x = self.mods2(x)
return x
model = M().eval()
example_inputs = (torch.rand(5, 5),)
qconfig_dict = {"": default_qconfig, "module_name": [("mods2", None)]}
m = prepare_fx(model, qconfig_dict, example_inputs=example_inputs)
m(*example_inputs)
m = convert_fx(m)
node_list = [
ns.call_function(torch.quantize_per_tensor),
ns.call_function(torch.ops.quantized.linear),
ns.call_function(torch.ops.quantized.linear),
ns.call_method('dequantize'),
ns.call_function(torch.nn.functional.linear)
]
self.checkGraphModuleNodes(m, expected_node_list=node_list)
m(torch.rand(5, 5))
def test_preserve_attributes(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv2d(1, 1, 1)
def forward(self, x):
return self.conv(x)
m = M()
m.eval()
m.preserved_attr = 3
prepare_custom_config_dict = {
"preserved_attributes": ["preserved_attr"]
}
example_inputs = (torch.randn(1, 1, 1, 1),)
m = prepare_fx(
m,
{"": default_qconfig},
example_inputs=example_inputs,
prepare_custom_config=prepare_custom_config_dict)
def assertAttrPreserved(m):
self.assertTrue(hasattr(m, "preserved_attr"))
self.assertEqual(m.preserved_attr, 3)
assertAttrPreserved(m)
convert_custom_config_dict = {
"preserved_attributes": ["preserved_attr"]
}
m = convert_fx(m, convert_custom_config=convert_custom_config_dict)
assertAttrPreserved(m)
@skipIfNoFBGEMM
def test_qat_and_script(self):
model = LinearModelWithSubmodule().train()
qengine = torch.backends.quantized.engine
qconfig_dict = {'': torch.ao.quantization.get_default_qat_qconfig(qengine)}
x = torch.randn(5, 5)
example_inputs = (x,)
model = prepare_qat_fx(model, qconfig_dict, example_inputs=example_inputs)
# ensure scripting works
scripted = torch.jit.script(model)
# run one round to make sure model runs
scripted(x)
FileCheck().check_count('FakeQuantize = prim::GetAttr[name="', 4, exactly=True) \
.run(scripted.graph)
# disable fake_quant and observer
for epoch in range(3):
if epoch == 1:
scripted.apply(torch.ao.quantization.disable_observer)
if epoch == 2:
scripted.apply(torch.ao.quantization.disable_fake_quant)
# ensure the fake_quant and observer have been disabled.
matches = ['.fake_quant_enabled', '.observer_enabled']
for key, v in scripted.state_dict().items():
if any(x in key for x in matches):
self.assertEqual(v, torch.tensor([0], dtype=torch.int64))
# enable them back
scripted.apply(torch.ao.quantization.enable_fake_quant)
scripted.apply(torch.ao.quantization.enable_observer)
for key, v in scripted.state_dict().items():
if any(x in key for x in matches):
self.assertEqual(v, torch.tensor([1], dtype=torch.int64))
@skipIfNoFBGEMM
def test_save_observer_state_dict(self):
orig = LinearModelWithSubmodule().eval()
model = orig
qconfig_dict = {'': torch.ao.quantization.get_default_qconfig('fbgemm')}
x = torch.randn(5, 5)
model = prepare_fx(model, qconfig_dict, example_inputs=(x,))
# run it through input
model(x)
# save state_dict of model
obs_dict = torch.ao.quantization.get_observer_state_dict(model)
quant = convert_fx(model)
b = io.BytesIO()
torch.save(obs_dict, b)
b.seek(0)
# Load the stats into new model
model_2 = orig
model_2 = prepare_fx(model_2, qconfig_dict, example_inputs=(x,))
loaded_dict = torch.load(b)
torch.ao.quantization.load_observer_state_dict(model_2, loaded_dict)
quant_2 = convert_fx(model_2)
# Verify that loaded state dict produces same results.
self.assertEqual(quant(x), quant_2(x))
@skipIfNoFBGEMM
def test_custom_module_class(self):
class CustomModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(3, 3)
def forward(self, x):
return self.linear(x)
class ObservedCustomModule(torch.nn.Module):
def __init__(self, linear):
super().__init__()
self.linear = linear
def forward(self, x):
return self.linear(x)
@classmethod
def from_float(cls, float_module):
assert hasattr(float_module, 'qconfig')
observed = cls(float_module.linear)
observed.qconfig = float_module.qconfig
return observed
class StaticQuantCustomModule(torch.nn.Module):
def __init__(self, linear):
super().__init__()
self.linear = linear
def forward(self, x):
return self.linear(x)
@classmethod
def from_observed(cls, observed_module):
assert hasattr(observed_module, 'qconfig')
assert hasattr(observed_module, 'activation_post_process')
observed_module.linear.activation_post_process = \
observed_module.activation_post_process
quantized = cls(nnq.Linear.from_float(observed_module.linear))
return quantized
class DynamicQuantCustomModule(torch.nn.Module):
def __init__(self, linear):
super().__init__()
self.linear = linear
def forward(self, x):
return self.linear(x)
@classmethod
def from_observed(cls, observed_module):
assert hasattr(observed_module, 'qconfig')
observed_module.linear.qconfig = observed_module.qconfig
quantized = cls(nnqd.Linear.from_float(observed_module.linear))
return quantized
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(3, 3)
self.custom = CustomModule()
def forward(self, x):
x = self.linear(x)
x = self.custom(x)
return x
class RefM(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear1 = torch.nn.Linear(3, 3)
self.linear2 = torch.nn.Linear(3, 3)
def forward(self, x):
x = self.linear1(x)
x = self.linear2(x)
return x
# instantiate M and RefM and align the parameters
original_m = M().eval()
original_ref_m = RefM().eval()
original_ref_m.linear1.weight = torch.nn.Parameter(original_m.linear.weight.detach())
original_ref_m.linear1.bias = torch.nn.Parameter(original_m.linear.bias.detach())
original_ref_m.linear2.weight = torch.nn.Parameter(original_m.custom.linear.weight.detach())
original_ref_m.linear2.bias = torch.nn.Parameter(original_m.custom.linear.bias.detach())
a16_qconfig = QConfig(
activation=MinMaxObserver.with_args(dtype=torch.qint32, quant_min=0, quant_max=65536),
weight=default_weight_observer,
)
test_configs = {
"static": (default_qconfig, StaticQuantCustomModule, 3),
"static_a16": (a16_qconfig, StaticQuantCustomModule, 3),
"dynamic": (default_dynamic_qconfig, DynamicQuantCustomModule, 0)
}
for quant_type in [QuantType.STATIC, QuantType.DYNAMIC]:
key = _get_quant_type_to_str(quant_type)
qconfig, quantized_module_class, num_observers = test_configs[key]
qconfig_dict = {"": qconfig}
if key == "static":
prepare_custom_config_dict = {
"float_to_observed_custom_module_class": {
"static": {
CustomModule: ObservedCustomModule
}
}
}
convert_custom_config_dict = {
"observed_to_quantized_custom_module_class": {
"static": {
ObservedCustomModule: quantized_module_class
}
}
}
else:
prepare_custom_config_dict = {
"non_traceable_module_class": [
CustomModule
]
}
convert_custom_config_dict = {
"observed_to_quantized_custom_module_class": {
"dynamic": {
CustomModule: quantized_module_class
}
}
}
example_inputs = (torch.randn(3, 3),)
# check prepared model
m = prepare_fx(
copy.deepcopy(original_m),
qconfig_dict,
example_inputs=example_inputs,
prepare_custom_config=prepare_custom_config_dict)
# calibration
m(*example_inputs)
# all activation observers are inserted in the top level module
count_check = {
ns.call_module(torch.ao.quantization.MinMaxObserver): num_observers
}
self.checkGraphModuleNodes(m, expected_node_occurrence=count_check)
# check converted/quantized model
m = convert_fx(
m,
convert_custom_config=convert_custom_config_dict)
if quant_type == QuantType.STATIC:
count_check = {
ns.call_function(torch.quantize_per_tensor) : 1,
ns.call_module(nnq.Linear) : 1,
ns.call_method('dequantize') : 1,
}
self.checkGraphModuleNodes(m, expected_node_occurrence=count_check)
self.assertEqual(type(m.custom), quantized_module_class)
res = m(*example_inputs)
# quantize the reference model
ref_m = prepare_fx(
copy.deepcopy(original_ref_m), qconfig_dict, example_inputs=example_inputs)
ref_m(*example_inputs)
ref_m = convert_fx(ref_m)
ref_res = ref_m(*example_inputs)
self.assertEqual(res, ref_res)
@skipIfNoFBGEMM
def test_custom_module_class_input_has_multiple_users(self):
""" Tests that the flow still works when the input of custom module
has multiple users
"""
class CustomModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(3, 3)
def forward(self, x):
return self.linear(x)
class ObservedCustomModule(torch.nn.Module):
def __init__(self, linear):
super().__init__()
self.linear = linear
def forward(self, x):
return self.linear(x)
@classmethod
def from_float(cls, float_module):
assert hasattr(float_module, 'qconfig')
observed = cls(float_module.linear)
observed.qconfig = float_module.qconfig
return observed
class StaticQuantCustomModule(torch.nn.Module):
def __init__(self, linear):
super().__init__()
self.linear = linear
def forward(self, x):
return self.linear(x)
@classmethod
def from_observed(cls, observed_module):
assert hasattr(observed_module, 'qconfig')
assert hasattr(observed_module, 'activation_post_process')
observed_module.linear.activation_post_process = \
observed_module.activation_post_process
quantized = cls(nnq.Linear.from_float(observed_module.linear))
return quantized
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(3, 3)
self.custom = CustomModule()
def forward(self, x0):
x1 = self.custom(x0)
x2 = self.linear(x0)
return x1 + x2
prepare_custom_config_dict = {
"float_to_observed_custom_module_class": {
"static": {
CustomModule: ObservedCustomModule
}
}
}
convert_custom_config_dict = {
"observed_to_quantized_custom_module_class": {
"static": {
ObservedCustomModule: StaticQuantCustomModule
}
}
}
m = M().eval()
example_inputs = (torch.randn(3, 3),)
m = prepare_fx(
m,
{"": default_qconfig},
example_inputs=example_inputs,
prepare_custom_config=prepare_custom_config_dict)
# make sure it works
m = convert_fx(
m,
convert_custom_config=convert_custom_config_dict)
# make sure it runs
m(*example_inputs)
@skipIfNoFBGEMM
def test_custom_module_class_input_has_duplicate_nodes(self):
""" Tests that the flow still works when the graph has
multiple nodes with the same custom module target.
"""
class CustomModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(3, 3)
def forward(self, x):
return self.linear(x)
class ObservedCustomModule(torch.nn.Module):
def __init__(self, linear):
super().__init__()
self.linear = linear
def forward(self, x):
return self.linear(x)
@classmethod
def from_float(cls, float_module):
assert hasattr(float_module, 'qconfig')
observed = cls(float_module.linear)
observed.qconfig = float_module.qconfig
return observed
class StaticQuantCustomModule(torch.nn.Module):
def __init__(self, linear):
super().__init__()
self.linear = linear
def forward(self, x):
return self.linear(x)
@classmethod
def from_observed(cls, observed_module):
assert hasattr(observed_module, 'qconfig')
assert hasattr(observed_module, 'activation_post_process')
observed_module.linear.activation_post_process = \
observed_module.activation_post_process
quantized = cls(nnq.Linear.from_float(observed_module.linear))
return quantized
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.custom = CustomModule()
def forward(self, x0):
x1 = self.custom(x0)
x2 = self.custom(x0)
return x1 + x2
prepare_custom_config_dict = {
"float_to_observed_custom_module_class": {
"static": {
CustomModule: ObservedCustomModule
}
}
}
convert_custom_config_dict = {
"observed_to_quantized_custom_module_class": {
"static": {
ObservedCustomModule: StaticQuantCustomModule
}
}
}
m = M().eval()
example_inputs = (torch.randn(3, 3),)
m = prepare_fx(
m,
{"": default_qconfig},
example_inputs=example_inputs,
prepare_custom_config=prepare_custom_config_dict)
# make sure it works
m = convert_fx(
m,
convert_custom_config=convert_custom_config_dict)
# make sure it runs
m(*example_inputs)
@skipIfNoFBGEMM
def test_non_traceable_module(self):
class NonTraceable(torch.nn.Module):
def forward(self, x):
for k in x.keys():
print(x[k])
return x
class NonTraceable2(torch.nn.Module):
def forward(self, x):
# data dependent control flow is not traceable
for i in x:
print(i)
return x
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.m1 = NonTraceable()
self.m2 = NonTraceable2()
def forward(self, x):
x = self.m1(x)
x = self.m2(x)
return x
m = M().eval()
qconfig_dict = {"": default_qconfig}
prepare_custom_config_dict = {
"non_traceable_module_name": [
"m1"
],
"non_traceable_module_class": [
NonTraceable2
]
}
m = prepare_fx(
m, qconfig_dict,
example_inputs=({"key": torch.randn(1)},),
prepare_custom_config=prepare_custom_config_dict)
node_occurrence = {
ns.call_module(NonTraceable) : 1,
ns.call_module(NonTraceable2) : 1,
}
# make sure these modules are not traced
self.checkGraphModuleNodes(m, expected_node_occurrence=node_occurrence)
def test_prepared_model_deepcopy(self):
"""Ensures that copy.deepcopy works correctly on a prepared model.
"""
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv2d(1, 1, 1)
self._foobar = 'foobar'
self.foobar2 = 'foobar2'
def forward(self, x):
x = self.conv(x)
return x
m = M()
m.eval()
qconfig_dict = {'': torch.ao.quantization.default_qconfig}
example_inputs = (torch.randn(4, 1, 4, 4),)
prepared = prepare_fx(m, qconfig_dict, example_inputs=example_inputs)
# calibrate
prepared(*example_inputs)
# copy
prepared_copy = copy.deepcopy(prepared)
# quantize, should run with no errors
quantized = convert_fx(prepared_copy)
def test_quantized_model_type(self):
""" Test state_dict and deepcopy works properly in the quantized model
"""
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(5, 5)
def forward(self, x):
return self.linear(x)
example_inputs = (torch.rand(8, 5),)
m = M().eval()
m = prepare_fx(m, {"": default_qconfig}, example_inputs=example_inputs)
m = convert_fx(m)
# test deepcopy
m_copy = copy.deepcopy(m)
self.assertEqual(m_copy(*example_inputs), m(*example_inputs))
# test state_dict
state_dict = m.state_dict()
m_new = M().eval()
m_new = prepare_fx(m_new, {"": default_qconfig}, example_inputs=example_inputs)
m_new = convert_fx(m_new)
m_new.load_state_dict(state_dict)
self.assertEqual(m_new(*example_inputs), m(*example_inputs))
def test_dequantize(self):
r""" Test to make sure dequantize node are placed before
non-quantizable node
"""
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv2d(1, 1, 1)
self.act = torch.nn.GELU()
def forward(self, x):
x = self.conv(x)
return self.act(x)
data = torch.rand(5, 1, 3, 3, dtype=torch.float)
for quant_type in self.static_quant_types:
node_list = [
ns.call_module(nnq.Conv2d),
ns.call_method("dequantize"),
ns.call_module(nn.GELU),
]
self.checkGraphModeFxOp(
M().eval(), (data,), quant_type, expected_node_list=node_list)
def test_sequential(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.convs = torch.nn.Sequential(
torch.nn.Conv2d(1, 1, 1),
torch.nn.Conv2d(1, 1, 1)
)
def forward(self, x):
x = self.convs(x)
return x
data = torch.rand(5, 1, 3, 3, dtype=torch.float)
for quant_type in self.static_quant_types:
node_list = [
ns.call_module(nnq.Conv2d),
ns.call_module(nnq.Conv2d),
]
self.checkGraphModeFxOp(
M().eval(), (data,), quant_type, expected_node_list=node_list)
def _test_quantized_inputs_outputs(
self, prepare_custom_config_dict, prepare_count_check,
convert_count_check):
"""
Test the option to have inputs and outputs of the graph quantized
"""
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1 = torch.nn.Conv2d(1, 1, 1)
self.conv2 = torch.nn.Conv2d(1, 1, 1)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
return x
# quantized input, quantized output
m = M()
qconfig_dict = {'': torch.ao.quantization.default_qconfig}
example_inputs = (torch.randn(1, 1, 4, 4),)
m.eval()
mp = torch.ao.quantization.quantize_fx.prepare_fx(
m, qconfig_dict,
example_inputs=example_inputs,
prepare_custom_config=prepare_custom_config_dict)
self.checkGraphModuleNodes(mp, expected_node_occurrence=prepare_count_check)
mp(*example_inputs)
mq = torch.ao.quantization.quantize_fx.convert_fx(mp)
self.checkGraphModuleNodes(mq, expected_node_occurrence=convert_count_check)
def test_quantized_input_quantized_output(self):
prepare_custom_config_dict = {
'input_quantized_idxs': [0], 'output_quantized_idxs': [0]}
prepare_count_check = {
ns.call_module(torch.ao.quantization.MinMaxObserver): 2,
}
convert_count_check = {
ns.call_function(torch.quantize_per_tensor): 0,
ns.call_method('dequantize'): 0,
}
self._test_quantized_inputs_outputs(
prepare_custom_config_dict, prepare_count_check, convert_count_check)
def test_fp32_input_quantized_output(self):
prepare_custom_config_dict = {
'output_quantized_idxs': [0]}
prepare_count_check = {
ns.call_module(torch.ao.quantization.MinMaxObserver): 3,
}
convert_count_check = {
ns.call_function(torch.quantize_per_tensor): 1,
ns.call_method('dequantize'): 0,
}
self._test_quantized_inputs_outputs(
prepare_custom_config_dict, prepare_count_check, convert_count_check)
def test_quantized_input_fp32_output(self):
prepare_custom_config_dict = {
'input_quantized_idxs': [0]}
prepare_count_check = {
ns.call_module(torch.ao.quantization.MinMaxObserver): 2,
}
convert_count_check = {
ns.call_function(torch.quantize_per_tensor): 0,
ns.call_method('dequantize'): 1,
}
self._test_quantized_inputs_outputs(
prepare_custom_config_dict, prepare_count_check, convert_count_check)
def test_fp32_input_fp32_output(self):
prepare_custom_config_dict = {}
prepare_count_check = {
ns.call_module(torch.ao.quantization.MinMaxObserver): 3,
}
convert_count_check = {
ns.call_function(torch.quantize_per_tensor): 1,
ns.call_method('dequantize'): 1,
}
self._test_quantized_inputs_outputs(
prepare_custom_config_dict, prepare_count_check, convert_count_check)
@skipIfNoFBGEMM
def test_convtranspose_per_channel_fails_early(self):
r"""
Verifies that attempting to quantize a ConvTranspose module with per-Channel
weight observers fails in the prepare step, as opposed to the convert step.
"""
m = torch.nn.Sequential(torch.nn.ConvTranspose2d(1, 1, 1))
m.eval()
qconfig_dict = {'': torch.ao.quantization.get_default_qconfig('fbgemm')}
with self.assertRaises(AssertionError) as context:
mp = prepare_fx(m, qconfig_dict, example_inputs=(torch.randn(1, 1, 1, 1),))
self.assertTrue(
str(context.exception) ==
'Per channel weight observer is not supported yet for ConvTranspose{n}d.')
@skipIfNoFBGEMM
def test_qparams_buffers(self):
class Linear(torch.nn.Module):
def __init__(self):
super().__init__()
self.w = torch.ones(5, 5)
self.b = torch.zeros(5)
def forward(self, x):
return torch.nn.functional.linear(x, self.w, self.b)
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.mods1 = torch.nn.Sequential(
Linear(),
Linear()
)
self.mods2 = Linear()
def forward(self, x):
x = self.mods1(x)
x = self.mods2(x)
return x
model = M().eval()
qconfig_dict = {"": default_qconfig}
example_inputs = (torch.rand(5, 5),)
m = prepare_fx(model, qconfig_dict, example_inputs=example_inputs)
m(*example_inputs)
m = convert_fx(m)
keys = m.state_dict().keys()
quant_scale_count = quant_zero_point = scale_count = zero_point_count = 0
for k in keys:
if 'input_scale' in k:
quant_scale_count = quant_scale_count + 1
elif 'input_zero_point' in k:
quant_zero_point = quant_zero_point + 1
elif 'scale' in k:
scale_count = scale_count + 1
elif 'zero_point' in k:
zero_point_count = zero_point_count + 1
# Expect each quantized linear op to have a scale and zero point
self.assertTrue(scale_count == 3, "Expect each quantized linear op to have a scale in state_dict")
self.assertTrue(zero_point_count == 3, "Expect each quantized linear op to have a zero_point in state_dict")
m(*example_inputs)
# ensure it is scriptable
scripted = torch.jit.script(m)
scripted_keys = scripted.state_dict().keys()
scripted.mods1_0_packed_weight_0 = m.state_dict()["mods1_0_packed_weight_0"]
non_packed_weight_keys = [key for key in keys if "_packed_weight" not in key]
self.assertTrue(
set(scripted_keys) == set(non_packed_weight_keys),
"Expected the scripted model to preserve the state_dict for non-packed weight attributes")
# TODO: probably don't want to hardcode the attribute names, since they are generated
for attr_name in [
"mods1_0_input_scale_0", "mods1_0_input_zero_point_0",
"mods1_0_scale_1", "mods1_0_zero_point_1",
"mods1_1_scale_1", "mods1_1_zero_point_1",
"mods2_scale_1", "mods2_zero_point_1"]:
self.assertTrue(hasattr(m, attr_name), attr_name + " not found.")
@skipIfNoFBGEMM
def test_packed_weight_fused_op(self):
class Linear(torch.nn.Module):
def __init__(self):
super().__init__()
self.w = torch.ones(5, 5)
self.b = torch.zeros(5)
def forward(self, x):
return F.linear(x, self.w, self.b)
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.mods1 = torch.nn.Sequential(
Linear(),
Linear()
)
self.mods2 = Linear()
self.relu = F.relu
def forward(self, x):
x = self.mods1(x)
x = self.mods2(x)
x = self.relu(x)
return x
model = M().eval()
example_inputs = (torch.rand(5, 5),)
qconfig_dict = {"": default_qconfig}
m = prepare_fx(model, qconfig_dict, example_inputs=example_inputs)
m(*example_inputs)
m = convert_fx(m)
assert hasattr(m, "mods1_0_packed_weight_0")
assert hasattr(m, "mods1_1_packed_weight_0")
assert hasattr(m, "mods2_packed_weight_0")
@skipIfNoFBGEMM
def test_mul_add_fp16_config(self):
with override_quantized_engine('fbgemm'):
class Linear(torch.nn.Module):
def __init__(self):
super().__init__()
self.w = torch.ones(5, 5)
self.b = torch.zeros(5)
def forward(self, x):
return torch.nn.functional.linear(x, self.w, self.b)
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.mods1 = torch.nn.Sequential(
Linear(),
Linear()
)
self.mods2 = Linear()
def forward(self, x):
x = x * 5
x = x + 5
x = self.mods1(x)
x = self.mods2(x)
return x
model = M().eval()
qconfig_dict = {"": float16_dynamic_qconfig}
example_inputs = (torch.rand(5, 5),)
m = prepare_fx(model, qconfig_dict, example_inputs=example_inputs)
m = convert_fx(m)
# make sure it runs
m(*example_inputs)
def test_getattr_with_nontensor_result(self):
"""
Verifies that binary ops get quantized correctly if some
of the args are nodes but not Tensors, such as an `x.ndim`
pattern.
"""
class M1(torch.nn.Module):
def forward(self, x):
dims = x.ndim
dims_sub = dims - 1
dims_sub2 = dims_sub - 1
x = torch.add(x, dims_sub2)
return x
class M2(torch.nn.Module):
def forward(self, x):
dims = x.ndim
dims_sub = dims - 2
mul = [1] * dims_sub
dims_list = [-1, x.size(1)] + mul
x = x.view(dims_list)
return x
class M3(torch.nn.Module):
def forward(self, x):
shape = x.shape
x = x.view(shape)
return x
for cls in (M1, M2, M3):
m = cls().eval()
example_inputs = (torch.rand(4, 4, 4, 4),)
m(*example_inputs)
qconfig_dict = {'': torch.ao.quantization.default_qconfig}
mp = prepare_fx(m, qconfig_dict, example_inputs=example_inputs)
mp(torch.rand(4, 4, 4, 4))
mc = convert_fx(mp)
class _NonReferenceTestModel(nn.Module):
def __init__(self, func, lin_in, lin_out):
super().__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.lin = nn.Linear(lin_in, lin_out)
self.func = func
def forward(self, x, y, z):
x = self.pool(F.relu(self.conv1(x)))
x = torch.flatten(x, 1)
x = self.func(x, y, z)
x = self.lin(x)
return x
# This function looks at the node specified by the NodeInfo in the key of
# node_info_to_non_tensor_args and checks that the args at specified indices
# are not observed (since they are non tensors). If the args at those indices
# are a tuple/list (which do not show up as nodes) the function checks the
# individual elements of the tuple/list recursively.
def _check_not_observed(self, model, node_info_to_non_tensor_args):
# this is a helper function (for easier recursion) that checks whether
# arg_node is observed
def _check_node_not_observed(model, arg_node, node):
if isinstance(arg_node, (tuple, list)):
for new_node in arg_node:
_check_node_not_observed(model, new_node, node)
elif arg_node.op == "call_module":
self.assertTrue(
not _is_activation_post_process(getattr(model, arg_node.target)),
f"Arg: {arg_node} of node: {node} is observed but is not a float tensor",
)
for node in model.graph.nodes:
indices = node_info_to_non_tensor_args.get(
NodeInfo(node.op, node.target), []
)
for index in indices:
if index < len(node.args):
arg_node = node.args[index]
_check_node_not_observed(model, arg_node, node)
# This test checks that the model gets prepared correct, doesn't have observers
# on specific ops (see _check_not_observed) and that the prepared model runs
def _test_dtype_propagation(self, model, node_info_to_non_tensor_args, *args):
model.eval()
qconfig_dict = {"": torch.ao.quantization.get_default_qconfig("fbgemm")}
prepared_model = prepare_fx(model, qconfig_dict, example_inputs=tuple(args))
self._check_not_observed(prepared_model, node_info_to_non_tensor_args)
prepared_model(*args)
def test_masked_fill_nontensor_args_not_observed(self):
def func(x, y, z):
return x.masked_fill(y, z)
model = self._NonReferenceTestModel(func, 1176, 1)
args = [torch.randn(5, 3, 32, 32), torch.randn(1176) > 0, 0.1]
node_info_to_non_tensor_args = {NodeInfo("call_method", "masked_fill"): [1, 2]}
self._test_dtype_propagation(model, node_info_to_non_tensor_args, *args)
def test_permute_nontensor_args_not_observed(self):
def func(x, y, z):
return x.permute(y, z)
model = self._NonReferenceTestModel(func, 1176, 1)
args = [torch.randn(5, 3, 32, 32), 0, 1]
node_info_to_non_tensor_args = {NodeInfo("call_method", "permute"): [1, 2]}
self._test_dtype_propagation(model, node_info_to_non_tensor_args, *args)
def test_repeat_nontensor_args_not_observed(self):
def func(x, y, z):
return x.repeat(y, z)
model = self._NonReferenceTestModel(func, 1176, 1)
args = [torch.randn(5, 3, 32, 32), 2, 1]
node_info_to_non_tensor_args = {NodeInfo("call_method", "repeat"): [1, 2]}
self._test_dtype_propagation(model, node_info_to_non_tensor_args, *args)
def test_reshape_nontensor_args_not_observed(self):
def func(x, y, z):
return x.reshape(-1, y)
model = self._NonReferenceTestModel(func, 5, 1)
args = [torch.randn(5, 3, 32, 32), 5, None]
node_info_to_non_tensor_args = {NodeInfo("call_method", "reshape"): [2]}
self._test_dtype_propagation(model, node_info_to_non_tensor_args, *args)
def test_size_nontensor_args_not_observed(self):
def func(x, y, z):
return x.reshape((-1, x.size(y)))
model = self._NonReferenceTestModel(func, 5, 1)
args = [torch.randn(5, 3, 32, 32), 0, None]
node_info_to_non_tensor_args = {NodeInfo("call_method", "size"): [1]}
self._test_dtype_propagation(model, node_info_to_non_tensor_args, *args)
def test_transpose_nontensor_args_not_observed(self):
def func(x, y, z):
return x.transpose(y, z)
model = self._NonReferenceTestModel(func, 5, 1)
args = [torch.randn(5, 3, 32, 32), 0, 1]
node_info_to_non_tensor_args = {NodeInfo("call_method", "transpose"): [1, 2]}
self._test_dtype_propagation(model, node_info_to_non_tensor_args, *args)
def test_torch_transpose_nontensor_args_not_observed(self):
# TODO: make torch.transpose traceable by fx when using
# variable nontensor arguments
# func = lambda x, y, z: torch.transpose(x, y, z) # error
def func(x, y, z):
return torch.transpose(x, 0, 1)
model = self._NonReferenceTestModel(func, 5, 1)
node_info_to_non_tensor_args = {
NodeInfo("call_method", torch.transpose): [1, 2]
}
args = [torch.randn(5, 3, 32, 32), 0, 1]
self._test_dtype_propagation(model, node_info_to_non_tensor_args, *args)
def test_unsqueeze_nontensor_args_not_observed(self):
def func(x, y, z):
return x.unsqueeze(y)
model = self._NonReferenceTestModel(func, 1176, 1)
args = [torch.randn(5, 3, 32, 32), 1, None]
node_info_to_non_tensor_args = {NodeInfo("call_method", "unsqueeze"): [1]}
self._test_dtype_propagation(model, node_info_to_non_tensor_args, *args)
def test_unsqueeze__nontensor_args_not_observed(self):
def func(x, y, z):
return x.unsqueeze_(y)
model = self._NonReferenceTestModel(func, 1176, 1)
args = [torch.randn(5, 3, 32, 32), 1, None]
node_info_to_non_tensor_args = {NodeInfo("call_method", "unsqueeze_"): [1]}
self._test_dtype_propagation(model, node_info_to_non_tensor_args, *args)
def test_torch_unsqueeze_nontensor_args_not_observed(self):
# TODO: make torch.unsqueeze scriptable by fx when using
# variable nontensor arguments
# func = lambda x, y, z: torch.unsqueeze(x, y) # error
def func(x, y, z):
return torch.unsqueeze(x, 1)
model = self._NonReferenceTestModel(func, 1176, 1)
args = [torch.randn(5, 3, 32, 32), 1, None]
node_info_to_non_tensor_args = {NodeInfo("call_method", torch.unsqueeze): [1]}
self._test_dtype_propagation(model, node_info_to_non_tensor_args, *args)
def test_view_nontensor_args_not_observed(self):
def func(x, y, z):
return x.view(-1, y)
model = self._NonReferenceTestModel(func, 5, 1)
args = [torch.randn(5, 3, 32, 32), 5, None]
node_info_to_non_tensor_args = {NodeInfo("call_method", "view"): [2]}
self._test_dtype_propagation(model, node_info_to_non_tensor_args, *args)
def test_propagate_dtypes_for_known_nodes_list_args(self):
def func(x, y, z):
return x.reshape(y)
model = self._NonReferenceTestModel(func, 5, 1)
args = [torch.randn(5, 3, 32, 32), [-1, 5], None]
node_info_to_non_tensor_args = {NodeInfo("call_method", "reshape"): [1]}
self._test_dtype_propagation(model, node_info_to_non_tensor_args, *args)
def test_propagate_dtypes_for_known_nodes_split_list_args(self):
def func(x, y, z):
return x.reshape([y, z])
model = self._NonReferenceTestModel(func, 5, 1)
args = [torch.randn(5, 3, 32, 32), -1, 5]
node_info_to_non_tensor_args = {NodeInfo("call_method", "reshape"): [1]}
self._test_dtype_propagation(model, node_info_to_non_tensor_args, *args)
def test_propagate_dtypes_for_known_nodes_tuple_args(self):
def func(x, y, z):
return x.reshape(y)
model = self._NonReferenceTestModel(func, 5, 1)
args = [torch.randn(5, 3, 32, 32), (-1, 5), None]
node_info_to_non_tensor_args = {NodeInfo("call_method", "reshape"): [1]}
self._test_dtype_propagation(model, node_info_to_non_tensor_args, *args)
def test_propagate_dtypes_for_known_nodes_split_tuple_args(self):
def func(x, y, z):
return x.reshape((y, z))
model = self._NonReferenceTestModel(func, 5, 1)
args = [torch.randn(5, 3, 32, 32), -1, 5]
node_info_to_non_tensor_args = {NodeInfo("call_method", "reshape"): [1]}
self._test_dtype_propagation(model, node_info_to_non_tensor_args, *args)
def test_propagate_dtypes_for_known_nodes_dict_args(self):
def func(x, y, z):
return x.transpose(y["first"], y["second"])
model = self._NonReferenceTestModel(func, 5, 1)
args = [torch.randn(5, 3, 32, 32), {"first": 0, "second": 1}, None]
node_info_to_non_tensor_args = {NodeInfo("call_method", "transpose"): [1, 2]}
self._test_dtype_propagation(model, node_info_to_non_tensor_args, *args)
def test_propagate_dtypes_for_known_nodes_dict_tuple_args(self):
class reshape_module(nn.Module):
def forward(self, x, y, z):
return x.reshape(y["shape"])
model = self._NonReferenceTestModel(reshape_module(), 5, 1)
args = [torch.randn(5, 3, 32, 32), {"shape": (-1, 5)}, None]
node_info_to_non_tensor_args = {NodeInfo("call_method", "reshape"): [1]}
self._test_dtype_propagation(model, node_info_to_non_tensor_args, *args)
def test_propagate_dtypes_for_known_nodes_dict_split_tuple_args(self):
def func(x, y, z):
return x.reshape((y["first"], y["second"]))
model = self._NonReferenceTestModel(func, 5, 1)
args = [torch.randn(5, 3, 32, 32), {"first": -1, "second": 5}, None]
node_info_to_non_tensor_args = {NodeInfo("call_method", "transpose"): [1]}
self._test_dtype_propagation(model, node_info_to_non_tensor_args, *args)
def test_assert_on_size_after_quant_layer(self):
"""
Verifies that calculating a size of a quantized tensor works
correctly in quantization passes.
"""
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 1, 1)
def forward(self, x):
x = self.conv1(x)
torch._assert(x.size(1) == 1, 'foobar')
return x
m = M().eval()
example_inputs = (torch.rand(4, 1, 4, 4),)
m(*example_inputs)
qconfig_dict = {'': torch.ao.quantization.default_qconfig}
mp = prepare_fx(m, qconfig_dict, example_inputs=example_inputs)
mp(*example_inputs)
mc = convert_fx(mp)
mc(*example_inputs)
def test_fp32_sum(self):
"""
Verifies that fp32 sum works correctly if it's before or after
quantized layers.
"""
class M1(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 1, 1)
def forward(self, x):
x = self.conv1(x)
x = torch.stack([x])
x = torch.sum(x)
return x
class M2(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 1, 1)
self.conv2 = nn.Conv2d(1, 1, 1)
def forward(self, x):
x = self.conv1(x)
x1 = torch.stack([x])
x1 = torch.sum(x1, dim=0)
x2 = self.conv2(x1)
return x2
for cls in (M1, M2):
m = cls().eval()
example_inputs = (torch.rand(4, 1, 4, 4),)
m(*example_inputs)
qconfig_dict = {'': torch.ao.quantization.default_qconfig}
mp = prepare_fx(m, qconfig_dict, example_inputs=example_inputs)
mp(*example_inputs)
mc = convert_fx(mp)
mc(*example_inputs)
def test_fusion_pattern_unquantized(self):
"""
Ensure that leaving a possible fusion pattern of multiple nodes
unquantized runs through the APIs without errors.
"""
class Child(torch.nn.Module):
def __init__(self):
super().__init__()
self.relu = nn.ReLU()
def forward(self, x):
x = torch.add(x, 1.0)
x = torch.nn.functional.relu(x)
return x
class Parent(torch.nn.Module):
def __init__(self):
super().__init__()
self.child = Child()
self.conv = nn.Conv2d(1, 1, 1)
def forward(self, x):
x = self.child(x)
x = self.conv(x)
return x
m = Parent().eval()
qconfig_dict = {
'': torch.ao.quantization.default_qconfig,
'module_name': [
('child', None),
],
}
example_inputs = (torch.rand(1, 1, 1, 1),)
mp = prepare_fx(m, qconfig_dict, example_inputs=example_inputs)
mp(*example_inputs)
mc = convert_fx(mp)
def test_state_dict(self):
""" Make sure packed params appear in state_dict
"""
# test linear packed weight
class M1(torch.nn.Module):
def __init__(self):
super().__init__()
self.w = torch.rand(4, 30)
self.b = torch.rand(4)
def forward(self, x):
return F.linear(x, self.w, self.b)
m = M1().eval()
qconfig_dict = {"": default_qconfig}
m = prepare_fx(m, qconfig_dict, example_inputs=(torch.randn(1, 30),))
m = convert_fx(m)
state_dict = m.state_dict()
self.assertTrue("_packed_weight_0" in state_dict)
# test conv packed weight
class M2(torch.nn.Module):
def __init__(self):
super().__init__()
self.w = torch.rand(3, 3, 3, 3)
self.b = torch.rand(3)
self.stride = (1, 1)
self.padding = (0, 0)
self.dilation = (1, 1)
self.groups = 1
def forward(self, x):
return F.conv2d(x, self.w, self.b, self.stride, self.padding, self.dilation, self.groups)
m = M2().eval()
qconfig_dict = {"": default_qconfig}
m = prepare_fx(m, qconfig_dict, example_inputs=(torch.randn(1, 3, 3, 3),))
m = convert_fx(m)
state_dict = m.state_dict()
self.assertTrue("_packed_weight_0" in state_dict)
# test load
ref_weight, ref_bias = torch.ops.quantized.conv2d_unpack(state_dict["_packed_weight_0"])
data = torch.rand(1, 3, 5, 5)
ref_res = m(data)
m = M2().eval()
m = prepare_fx(m, qconfig_dict, (data,))
m = convert_fx(m)
res = m(data)
weight, bias = m._packed_weight_0.unpack()
# check that random model weight/bias does not match ref weight/bias
self.assertNotEqual(weight, ref_weight)
self.assertNotEqual(bias, ref_bias)
self.assertNotEqual(res, ref_res)
m.load_state_dict(state_dict)
def checkModel(m, data, ref_weight, ref_bias, ref_res):
res = m(data)
weight, bias = m._packed_weight_0.unpack()
# check that weight/bias matches after load the state_dict
self.assertEqual(weight, ref_weight)
self.assertEqual(bias, ref_bias)
self.assertEqual(res, ref_res)
checkModel(m, data, ref_weight, ref_bias, ref_res)
# Test save to disk and load back
m = M2().eval()
m = prepare_fx(m, qconfig_dict, example_inputs=(data,))
m = convert_fx(m)
m.load_state_dict(state_dict)
with TemporaryFileName() as fname:
torch.save(m.state_dict(), fname)
m.load_state_dict(torch.load(fname))
checkModel(m, data, ref_weight, ref_bias, ref_res)
@skipIfNoFBGEMM
def test_preserve_qconfig(self):
"""
Test to make sure the temporary config option to preserve qconfig attributes
in the model works
"""
with override_quantized_engine('fbgemm'):
class Linear(torch.nn.Module):
def __init__(self):
super().__init__()
self.w = torch.ones(5, 5)
self.b = torch.zeros(5)
def forward(self, x):
return torch.nn.functional.linear(x, self.w, self.b)
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.mods1 = torch.nn.Sequential(
Linear(),
Linear()
)
self.mods2 = torch.nn.Sigmoid()
def forward(self, x):
x = self.mods1(x)
x = self.mods2(x)
return x
model = M().eval()
qconfig_dict = {
"object_type": [
(torch.nn.functional.linear, float16_dynamic_qconfig),
],
}
example_inputs = (torch.rand(5, 5),)
m = prepare_fx(model, qconfig_dict, example_inputs=example_inputs)
m(*example_inputs)
m = convert_fx(m, _remove_qconfig=False)
self.assertTrue(hasattr(m.mods2, 'qconfig'))
def test_not_used(self):
""" Test quantizing a not used value"""
class M(torch.nn.Module):
def forward(self, x):
x = x + x
x.sigmoid_()
return x
m = M().eval()
qconfig_mapping = get_default_qconfig_mapping().set_global(float16_static_qconfig)
# make sure quantization runs
m = prepare_fx(m, qconfig_mapping, example_inputs=(torch.randn(1),))
m = convert_fx(m)
def test_qparams_fqn(self):
""" Test that the FQN of input_scale/zero_point is set
to that of first linear use. """
class Linear(torch.nn.Module):
def __init__(self):
super().__init__()
self.w = torch.ones(5, 5)
self.b = torch.zeros(5)
def forward(self, x):
return torch.nn.functional.linear(x, self.w, self.b)
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.mods1 = torch.nn.Sequential(
Linear(),
Linear()
)
def forward(self, x):
x = torch.cat((x,), 1)
tmp = x.size()
x = self.mods1(x)
y = x * tmp[0]
return y
model = M().eval()
qconfig_dict = {
"": None,
"object_type": [
(torch.nn.functional.linear, default_qconfig),
(torch.nn.functional.relu, default_qconfig),
],
}
example_inputs = (torch.rand(5, 5),)
m = prepare_fx(model, qconfig_dict, example_inputs=example_inputs)
m(*example_inputs)
m = convert_fx(m)
keys = m.state_dict().keys()
m(torch.randn(5, 5))
# TODO: probably don't want to hardcode the attribute names, since they are generated
for attr_name in [
"mods1_0_input_scale_0", "mods1_0_input_zero_point_0",
"mods1_0_scale_0", "mods1_0_zero_point_0",
"mods1_1_scale_0", "mods1_1_zero_point_0"]:
self.assertTrue(hasattr(m, attr_name), attr_name + " not found.")
def test_no_obs_between_unmatched_node_and_copy_node(self):
"""
Verifies that an observer is not inserted between an unmatched
node and a node matched to CopyNodeQuantizeHandler. This is done
because observers require activations to be Tensors, and there is
no guarantee that an output of an unmatched node is a Tensor.
"""
class M(nn.Module):
def __init__(self):
super().__init__()
self.relu = nn.ReLU()
def forward(self, x):
x = _user_func_with_complex_return_type(x)
x1 = x[0] + 1
return x1, x[1]
m = M().eval()
qconfig_dict = {'': torch.ao.quantization.default_qconfig}
example_inputs = (torch.randn(4, 4, 4, 4),)
mp = prepare_fx(m, qconfig_dict, example_inputs=example_inputs)
# if an observer is inserted after _user_func_with_complex_return_type,
# the following call will fail
mp(*example_inputs)
mc = convert_fx(mp)
mc(*example_inputs)
def test_fold_quant_dequant(self):
""" Test that the sequence of quant-dequant nodes in the
graph, get folded and we erase the extra dequant nodes.
"""
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.w = torch.ones(5, 5)
self.b = torch.zeros(5)
def forward(self, x):
x = torch.cat((x,), 1)
tmp = x.size()
x = torch.nn.functional.linear(x, self.w, self.b)
y = x * tmp[0]
return y
model = M().eval()
qconfig_dict = {
"": None,
"object_type": [
(torch.nn.functional.linear, default_qconfig),
],
}
example_inputs = (torch.rand(5, 5),)
m = prepare_fx(model, qconfig_dict, example_inputs=example_inputs)
m(*example_inputs)
m = convert_fx(m)
keys = m.state_dict().keys()
m(*example_inputs)
dequant = 0
quant = 0
for n in m.graph.nodes:
if n.op == "call_method" and n.target == "dequantize":
dequant = dequant + 1
if n.op == "call_function" and n.target == torch.quantize_per_tensor:
quant = quant + 1
self.assertEqual(dequant, 1)
self.assertEqual(quant, 1)
def test_quant_output_always_observed(self):
"""
If the output is hardcoded to be quantized, ensure that
there is always an observer, even if the last non-output node is not
quantizeable.
"""
qconfig_dict = {'': torch.ao.quantization.get_default_qat_qconfig('fbgemm')}
prepare_custom_config_dict = {'output_quantized_idxs': [0]}
example_inputs = (torch.randn(4, 1, 4, 4),)
# non-quantizeable node, quantized output
class M1(torch.nn.Module):
def __init__(self):
super().__init__()
self.identity = torch.nn.Identity()
def forward(self, x):
x = self.identity(x)
return x
m1 = M1()
self.checkGraphModeFxOp(
m1, example_inputs, QuantType.QAT,
prepare_expected_node_occurrence={
ns.call_module(torch.ao.quantization.FusedMovingAvgObsFakeQuantize): 2,
},
expected_node_occurrence={
ns.call_function(torch.quantize_per_tensor): 1,
},
prepare_custom_config=prepare_custom_config_dict)
# quantizeable node, quantized output
class M2(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv2d(1, 1, 1)
def forward(self, x):
x = self.conv(x)
return x
m2 = M2()
self.checkGraphModeFxOp(
m2, example_inputs, QuantType.QAT,
prepare_expected_node_occurrence={
# one for weights, one for activations
ns.call_module(torch.ao.quantization.FusedMovingAvgObsFakeQuantize): 2,
},
expected_node_occurrence={
ns.call_function(torch.quantize_per_tensor): 1,
},
prepare_custom_config=prepare_custom_config_dict)
# quantizeable node, quantized dictionary output
class M3(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv2d(1, 1, 1)
def forward(self, x):
x = self.conv(x)
return {"output": x}
m3 = M3()
self.checkGraphModeFxOp(
m3, example_inputs, QuantType.QAT,
prepare_expected_node_occurrence={
# one for weights, one for activations
ns.call_module(torch.ao.quantization.FusedMovingAvgObsFakeQuantize): 2,
},
expected_node_occurrence={
ns.call_function(torch.quantize_per_tensor): 1,
},
prepare_custom_config=prepare_custom_config_dict)
def test_deepcopy_preserve_attributes(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.attr = 3
def forward(self, x):
return x
m = M().eval()
m = prepare_fx(
m,
{"": default_qconfig},
example_inputs=(torch.randn(1),),
prepare_custom_config={"preserved_attributes": ["attr"]})
# preserved attributes are also stored in meta so that it doesn't get lost
# during deepcopy
self.assertTrue(hasattr(m, "attr"))
self.assertTrue("attr" in m.meta[_USER_PRESERVED_ATTRIBUTES_KEY])
m2 = copy.deepcopy(m)
self.assertTrue(hasattr(m2, "attr"))
self.assertTrue("attr" in m2.meta[_USER_PRESERVED_ATTRIBUTES_KEY])
m = convert_fx(m, convert_custom_config={"preserved_attributes": ["attr"]})
self.assertTrue(hasattr(m, "attr"))
self.assertTrue("attr" in m.meta[_USER_PRESERVED_ATTRIBUTES_KEY])
m2 = copy.deepcopy(m)
self.assertTrue(hasattr(m2, "attr"))
self.assertTrue("attr" in m2.meta[_USER_PRESERVED_ATTRIBUTES_KEY])
def test_output_lists_and_dicts(self):
"""Verify that specifying complicated output types does not crash.
"""
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(1, 1, 1)
def forward(self, x):
x = self.conv(x)
return {'foo': [x]}, [{'foo': [[x]]}]
m = M().eval()
qconfig_dict = {'': torch.ao.quantization.default_qconfig}
mp = prepare_fx(m, qconfig_dict, example_inputs=(torch.randn(1, 1, 1, 1),))
mc = convert_fx(mp)
def test_shape_followed_by_quantized_op(self):
""" Make sure that shape does not dequantize
the Tensor before the next operator
"""
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1 = torch.nn.Conv2d(2, 2, 2)
self.conv2 = torch.nn.Conv2d(2, 2, 2)
def forward(self, x):
x = self.conv1(x)
s = x.shape
torch._assert(s == x.shape, "")
x = self.conv2(x)
return x
# make sure quantization runs
m = M().eval()
example_inputs = (torch.randn(2, 2, 4, 4),)
m = prepare_fx(m, {"": default_qconfig}, example_inputs=example_inputs)
m = convert_fx(m)
m(*example_inputs)
node_occurrence = {
ns.call_function(torch.quantize_per_tensor): 1,
ns.call_method("dequantize"): 1
}
self.checkGraphModuleNodes(m, expected_node_occurrence=node_occurrence)
def test_trace_quantize_per_tensor(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv2d(1, 1, 1)
def forward(self, x):
x = self.conv(x)
return x
m = M().eval()
m = prepare_fx(m, {"": default_qconfig}, example_inputs=(torch.randn(1, 1, 3, 3),))
m = convert_fx(m)
# Make sure this runs without error
m = torch.fx.Transformer(m).transform()
def test_copy_node_has_shared_actpp_instance(self):
""" Test the output of CopyNode to have the same
observer/fake_quant instance as the input
"""
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.avgpool2d = torch.nn.AvgPool2d(kernel_size=3)
def forward(self, x):
x = self.avgpool2d(x)
return x
for quant_type in self.static_quant_types:
m = M()
# Checks that we have an observer for both input and output
occurrence_map = {
QuantType.STATIC: {
ns.call_module(torch.ao.quantization.MinMaxObserver): 2
},
QuantType.QAT: {
ns.call_module(torch.ao.quantization.FakeQuantize): 2
}
}
if quant_type == QuantType.QAT:
m.train()
prepare = prepare_qat_fx
qconfig = default_qat_qconfig
actpp_module_class = torch.ao.quantization.FakeQuantize
else:
m.eval()
prepare = prepare_fx
qconfig = default_qconfig
actpp_module_class = torch.ao.quantization.MinMaxObserver
example_inputs = (torch.randn(1, 3, 3, 3),)
m = prepare(m, {"": qconfig}, example_inputs=example_inputs)
# check that there is a duplicated observer instance
actpp_module_count = 0
for name, module in m.named_modules(remove_duplicate=False):
if isinstance(module, actpp_module_class):
actpp_module_count += 1
self.assertEqual(actpp_module_count, 2)
actpp_module_count = 0
for name, module in m.named_modules():
if isinstance(module, actpp_module_class):
actpp_module_count += 1
self.assertEqual(actpp_module_count, 1)
m_copy = copy.deepcopy(m)
m = convert_fx(m)
m_reference = convert_to_reference_fx(m_copy)
# checks for non-reference quantized model
node_occurrence = {
ns.call_function(torch.quantize_per_tensor): 1,
ns.call_method("dequantize"): 1
}
node_list = [
ns.call_function(torch.quantize_per_tensor),
ns.call_module(torch.nn.AvgPool2d),
ns.call_method("dequantize"),
]
self.checkGraphModuleNodes(m, expected_node_occurrence=node_occurrence, expected_node_list=node_list)
# checks for reference quantized model, for copy nodes we'll have
# dequant - copy_node - quant patterns which will be fused later
# in the backend lowering step
node_occurrence = {
ns.call_function(torch.quantize_per_tensor): 2,
ns.call_method("dequantize"): 2
}
node_list = [
ns.call_function(torch.quantize_per_tensor),
ns.call_method("dequantize"),
ns.call_module(torch.nn.AvgPool2d),
ns.call_function(torch.quantize_per_tensor),
ns.call_method("dequantize"),
]
self.checkGraphModuleNodes(m_reference, expected_node_occurrence=node_occurrence, expected_node_list=node_list)
def test_linear_qint8_activation(self):
"""Test support for qint8 activation in reference pattern
"""
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv2d(1, 2, 2, 2)
self.linear = torch.nn.Linear(8, 5)
def forward(self, x):
x = self.conv(x)
x = torch.flatten(x, 1)
x = self.linear(x)
return x
m = M().eval()
example_inputs = (torch.rand(2, 1, 5, 5),)
m = prepare_fx(
m,
{"": torch.ao.quantization.QConfig(
activation=torch.ao.quantization.HistogramObserver.with_args(
qscheme=torch.per_tensor_symmetric, dtype=torch.qint8
), weight=torch.ao.quantization.default_per_channel_weight_observer)},
example_inputs=example_inputs)
m = convert_to_reference_fx(m)
m(*example_inputs)
def test_preserve_tuple(self):
""" Test tuple input type is preserved
"""
class LSTM(nn.Module):
def __init__(self):
super().__init__()
self.lstm = nn.LSTM(50, 50, 1)
def forward(self, inputs: torch.Tensor, state: List[torch.Tensor]):
h = state[0]
c = state[1]
return self.lstm(inputs, (h, c))
m = LSTM().eval()
example_inputs = (torch.randn(5, 3, 50), torch.randn(2, 3, 50), torch.randn(2, 3, 50))
m = prepare_fx(m, {"": default_qconfig}, example_inputs=example_inputs)
# make sure the arg[1] of lstm module is a tuple
for n in m.graph.nodes:
if n.target == "lstm":
self.assertEqual(type(n.args[1]), tuple)
def _test_static_lstm_helper(self, model, prepare_node_occurrence, convert_node_occurrence):
"""
Helper method to validate the graph of a model with static LSTM.
"""
qconfig_mapping = get_default_qconfig_mapping()
prepare_custom_config = PrepareCustomConfig() \
.set_float_to_observed_mapping(torch.nn.LSTM, torch.ao.nn.quantizable.LSTM)
convert_custom_config = ConvertCustomConfig() \
.set_observed_to_quantized_mapping(torch.ao.nn.quantizable.LSTM, torch.ao.nn.quantized.LSTM)
example_inputs = (torch.rand(5, 3, 50), torch.rand(1, 3, 50), torch.randn(1, 3, 50))
model = prepare_fx(model, qconfig_mapping, example_inputs, prepare_custom_config=prepare_custom_config)
self.checkGraphModuleNodes(model, expected_node_occurrence=prepare_node_occurrence)
model(*example_inputs)
model = convert_fx(model, convert_custom_config=convert_custom_config)
self.checkGraphModuleNodes(model, expected_node_occurrence=convert_node_occurrence)
model(*example_inputs)
def test_static_lstm(self):
"""
Test statically quantized custom module LSTM followed by ops that consume individual
tensors of the output tuple.
"""
class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.lstm = nn.LSTM(50, 50, 1)
self.linear1 = nn.Linear(50, 10)
self.linear2 = nn.Linear(50, 10)
self.linear3 = nn.Linear(50, 10)
def forward(self, inputs: torch.Tensor, h0: torch.Tensor, c0: torch.Tensor):
(out, (h0_out, c0_out)) = self.lstm(inputs, (h0, c0))
out = self.linear1(out)
h0_out = self.linear2(h0_out)
c0_out = self.linear3(c0_out)
return (out, (h0_out, c0_out))
m = MyModel()
prepare_node_occurrence = {
ns.call_module(torch.ao.nn.quantizable.LSTM): 1,
}
convert_node_occurrence = {
ns.call_module(torch.ao.nn.quantized.LSTM): 1,
ns.call_function(torch.quantize_per_tensor): 3,
# lstm[0].dequantize()
# lstm[1][0].dequantize()
# lstm[1][1].dequantize()
ns.call_method("dequantize"): 3,
# lstm[0], lstm[1], lstm[1][0], lstm[1][1]
ns.call_function(operator.getitem): 4,
# No tuples are consumed
ns.call_function(tuple): 0,
}
self._test_static_lstm_helper(m, prepare_node_occurrence, convert_node_occurrence)
def test_static_lstm_consume_tuple(self):
"""
Test statically quantized custom module LSTM followed by a module that consumes the
output tuple, either as a whole or part of it.
"""
class ModuleAfterLSTM(nn.Module):
def __init__(self):
super().__init__()
self.identity = torch.nn.Identity()
def forward(self, x):
return self.identity(x)
class ConsumeWholeTuple(nn.Module):
def __init__(self):
super().__init__()
self.lstm = nn.LSTM(50, 50, 1)
self.module_after_lstm = ModuleAfterLSTM()
def forward(self, inputs: torch.Tensor, h0: torch.Tensor, c0: torch.Tensor):
x = self.lstm(inputs, (h0, c0))
x = self.module_after_lstm(x) # consume tuple (output, (hidden0, hidden1))
return x
class ConsumeHiddenTuple(ConsumeWholeTuple):
def forward(self, inputs: torch.Tensor, h0: torch.Tensor, c0: torch.Tensor):
x = self.lstm(inputs, (h0, c0))
x = self.module_after_lstm(x[1]) # consume tuple (hidden0, hidden1)
return x
# Test consuming the whole tuple (output, (hidden0, hidden1))
m1 = ConsumeWholeTuple()
prepare_node_occurrence = {
ns.call_module(torch.ao.nn.quantizable.LSTM): 1,
}
convert_node_occurrence1 = {
ns.call_module(torch.ao.nn.quantized.LSTM): 1,
ns.call_function(torch.quantize_per_tensor): 3,
# lstm[0].dequantize()
# lstm[1][0].dequantize()
# lstm[1][1].dequantize()
ns.call_method("dequantize"): 3,
# lstm[0], lstm[1], lstm[1][0], lstm[1][1]
ns.call_function(operator.getitem): 4,
# tuple(output_dq, tuple(hidden0_dq, hidden1_dq))
ns.call_function(tuple): 2,
}
self._test_static_lstm_helper(m1, prepare_node_occurrence, convert_node_occurrence1)
# Test consuming just the hidden tuple (hidden0, hidden1)
m2 = ConsumeHiddenTuple()
convert_node_occurrence2 = {
ns.call_module(torch.ao.nn.quantized.LSTM): 1,
ns.call_function(torch.quantize_per_tensor): 3,
# lstm[1][0].dequantize()
# lstm[1][1].dequantize()
ns.call_method("dequantize"): 2,
# lstm[1], lstm[1][0], lstm[1][1]
ns.call_function(operator.getitem): 3,
# tuple(hidden0_dq, hidden1_dq)
ns.call_function(tuple): 1,
}
self._test_static_lstm_helper(m2, prepare_node_occurrence, convert_node_occurrence2)
def test_static_lstm_with_custom_fixed_qparams(self):
"""
Test statically quantized LSTM with custom fixed qparams assigned to each of the
inner submodules. This flow requires users to extend `torch.ao.nn.quantizable.LSTM`
and use the child class in the custom module mapping.
"""
class MyModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.my_lstm = torch.nn.LSTM(50, 50, 1)
def forward(self, inputs: torch.Tensor, h0: torch.Tensor, c0: torch.Tensor):
x = self.my_lstm(inputs, (h0, c0))
return x
# Construct a BackendConfig that supports qint32 for certain ops
# TODO: build a BackendConfig from scratch instead of modifying an existing one
qint32_dtype_config = DTypeConfig(input_dtype=torch.qint32, output_dtype=torch.qint32)
my_backend_config = get_qnnpack_backend_config()
for config in my_backend_config.configs:
if config.pattern in [torch.nn.Sigmoid, torch.nn.Tanh, torch.add, torch.mul]:
config.add_dtype_config(qint32_dtype_config)
class UserObservedLSTM(torch.ao.nn.quantizable.LSTM):
"""
Example of user provided LSTM implementation that assigns fixed qparams
to the inner ops.
"""
@classmethod
def from_float(cls, float_lstm):
assert isinstance(float_lstm, cls._FLOAT_MODULE)
# uint16, [-16, 16)
linear_output_obs_ctr = FixedQParamsObserver.with_args(scale=2 ** -11, zero_point=2 ** 15, dtype=torch.qint32)
# uint16, [0, 1)
sigmoid_obs_ctr = FixedQParamsObserver.with_args(scale=2 ** -16, zero_point=0, dtype=torch.qint32)
# uint16, [-1, 1)
tanh_obs_ctr = FixedQParamsObserver.with_args(scale=2 ** -15, zero_point=2 ** 15, dtype=torch.qint32)
# int16, [-16, 16)
cell_state_obs_ctr = FixedQParamsObserver.with_args(scale=2 ** -11, zero_point=0, dtype=torch.qint32)
# uint8, [-1, 1)
hidden_state_obs_ctr = FixedQParamsObserver.with_args(scale=2 ** -7, zero_point=2 ** 7, dtype=torch.quint8)
example_inputs = (torch.rand(5, 3, 50), (torch.rand(1, 3, 50), torch.randn(1, 3, 50)))
return torch.ao.quantization.fx.lstm_utils._get_lstm_with_individually_observed_parts(
float_lstm=float_lstm,
example_inputs=example_inputs,
backend_config=my_backend_config,
linear_output_obs_ctr=linear_output_obs_ctr,
sigmoid_obs_ctr=sigmoid_obs_ctr,
tanh_obs_ctr=tanh_obs_ctr,
cell_state_obs_ctr=cell_state_obs_ctr,
hidden_state_obs_ctr=hidden_state_obs_ctr,
)
class UserQuantizedLSTM(torch.ao.nn.quantized.LSTM):
"""
Example of user provided LSTM implementation that produces a reference
quantized module from a `UserObservedLSTM`.
"""
@classmethod
def from_observed(cls, observed_lstm):
assert isinstance(observed_lstm, cls._FLOAT_MODULE)
return torch.ao.quantization.fx.lstm_utils._get_reference_quantized_lstm_module(
observed_lstm=observed_lstm,
backend_config=my_backend_config,
)
# FX graph mode quantization
m = MyModel()
qconfig_mapping = get_default_qconfig_mapping("qnnpack")
example_inputs = (torch.rand(5, 3, 50), torch.rand(1, 3, 50), torch.randn(1, 3, 50))
prepare_custom_config = PrepareCustomConfig() \
.set_float_to_observed_mapping(torch.nn.LSTM, UserObservedLSTM)
convert_custom_config = ConvertCustomConfig() \
.set_observed_to_quantized_mapping(torch.ao.nn.quantizable.LSTM, UserQuantizedLSTM)
prepared = prepare_fx(
m,
qconfig_mapping,
example_inputs,
prepare_custom_config,
backend_config=my_backend_config,
)
prepared(*example_inputs)
converted = convert_fx(
prepared,
convert_custom_config,
backend_config=my_backend_config,
)
converted(*example_inputs)
# Find the patterns [dq - op - q_to_specific_dtype] in the graph and
# verify that qparams and dtypes are set correctly in the quantize ops
node_name_to_expected_quantize_args = {
"igates": (None, None, torch.quint8),
"hgates": (None, None, torch.quint8),
"add": (2 ** -11, 2 ** 15, torch.qint32), # gates.add
"input_gate": (2 ** -16, 0, torch.qint32),
"forget_gate": (2 ** -16, 0, torch.qint32),
"cell_gate": (2 ** -15, 2 ** 15, torch.qint32),
"output_gate": (2 ** -16, 0, torch.qint32),
"mul": (2 ** -11, 0, torch.qint32), # fgate_cx.mul
"mul_1": (2 ** -11, 0, torch.qint32), # igate_cgate.mul
"add_1": (2 ** -11, 0, torch.qint32), # fgate_cx_igate_cgate.add
"mul_2": (2 ** -7, 2 ** 7, torch.quint8), # ogate_cy.mul
}
cell = converted.my_lstm.layers.get_submodule("0").layer_fw.cell
matched_names = set()
for node in cell.graph.nodes:
if node.name not in node_name_to_expected_quantize_args:
continue
matched_names.add(node.name)
# Match preceding dequantize
self.assertTrue(all(arg.target == "dequantize" for arg in node.args))
# Match following quantize with the specific qparams and dtypes
expected_scale, expected_zp, expected_dtype = node_name_to_expected_quantize_args[node.name]
for user in node.users.keys():
self.assertEqual(user.target, torch.quantize_per_tensor)
if expected_scale is not None:
self.assertEqual(getattr(cell, user.args[1].target), expected_scale)
if expected_zp is not None:
self.assertEqual(getattr(cell, user.args[2].target), expected_zp)
self.assertEqual(user.args[-1], expected_dtype)
# Ensure all patterns were matched
self.assertEqual(matched_names, set(node_name_to_expected_quantize_args.keys()))
def test_reroute_tuple_getitem_patterns(self):
"""
The following graph should redirect the output to `b`. After the transformation,
all other nodes, including the inputs `a` and `c`, are no longer needed.
a b c
| \\ /
\\ tuple
\\ /
tuple
/ \\
/ \\
| \\
| \\
| \\
getitem0 getitem1
| / \\
| getitem0 getitem1
| \\ /
\\ tuple
\\ /
\\ /
tuple
|
getitem1
|
getitem0
|
output
"""
# Construct graph manually because symbolic_trace does not insert tuple and getitem nodes
graph = torch.fx.Graph()
a = graph.create_node("placeholder", "a")
b = graph.create_node("placeholder", "b")
c = graph.create_node("placeholder", "c")
bc = graph.call_function(tuple, args=([b, c],))
abc = graph.call_function(tuple, args=([a, bc],))
# Break down tuple and reconstruct it again
a2 = graph.call_function(operator.getitem, args=(abc, 0))
bc2 = graph.call_function(operator.getitem, args=(abc, 1))
b2 = graph.call_function(operator.getitem, args=(bc2, 0))
c2 = graph.call_function(operator.getitem, args=(bc2, 1))
bc3 = graph.call_function(tuple, args=([b2, c2],))
abc2 = graph.call_function(tuple, args=([a2, bc3],))
# Output tuple[1][0]
bc4 = graph.call_function(operator.getitem, args=(abc2, 1))
b3 = graph.call_function(operator.getitem, args=(bc4, 0))
output = graph.output(b3)
# Do reroute
_reroute_tuple_getitem_pattern(graph)
# Assert that output reroutes to `b` directly, and all other nodes can be removed
output_ancestors = []
def gather_ancestors(current_node): # noqa: E306
for arg in current_node.args:
output_ancestors.append(arg)
gather_ancestors(arg)
gather_ancestors(output)
self.assertEqual(output_ancestors, [b])
self.assertEqual(output.args[0], b)
def test_relu_lowering(self):
class M(torch.nn.Module):
def forward(self, x):
return torch.nn.functional.relu(x)
m = M().eval()
m = prepare_fx(m, {"": default_qconfig}, example_inputs=(torch.randn(1),))
m_copy = copy.deepcopy(m)
m = convert_fx(m)
m_ref = convert_to_reference_fx(m_copy)
node_occurrence = {
ns.call_function(torch.quantize_per_tensor): 1,
ns.call_method("dequantize"): 1
}
node_occurrence_ref = {
ns.call_function(torch.quantize_per_tensor): 2,
ns.call_method("dequantize"): 2
}
self.checkGraphModuleNodes(m, expected_node_occurrence=node_occurrence)
self.checkGraphModuleNodes(m_ref, expected_node_occurrence=node_occurrence_ref)
@skipIfNoFBGEMM
def test_dynamic_with_fusion(self):
"""
Tests that dynamic quantization APIs work with Linear + Relu fusion
"""
with override_quantized_engine('fbgemm'):
class LinearRelu(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(5, 5)
self.relu = torch.nn.ReLU()
def forward(self, x):
x = self.linear(x)
return self.relu(x)
class Linear(torch.nn.Module):
def __init__(self):
super().__init__()
self.w = torch.ones(5, 5)
self.b = torch.zeros(5)
def forward(self, x):
return torch.nn.functional.linear(x, self.w, self.b)
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.mods1 = torch.nn.Sequential(LinearRelu(), LinearRelu())
self.mods2 = Linear()
self.relu = F.relu
def forward(self, x):
x = self.mods1(x)
x = self.mods2(x)
x = self.relu(x)
return x
dynamic_quantized_ops = {
float16_dynamic_qconfig: torch.ops.quantized.linear_relu_dynamic_fp16,
default_dynamic_qconfig: torch.ops.quantized.linear_relu_dynamic
}
for qconfig in [float16_dynamic_qconfig, default_dynamic_qconfig]:
model = M().eval()
qconfig_dict = {
"": qconfig
}
example_inputs = (torch.rand(5, 5),)
m = prepare_fx(model, qconfig_dict, example_inputs=example_inputs)
m = convert_fx(m)
m(*example_inputs)
node_list = [
ns.call_module(nniqd.LinearReLU),
ns.call_module(nniqd.LinearReLU),
ns.call_function(dynamic_quantized_ops[qconfig]),
]
self.checkGraphModuleNodes(m, expected_node_list=node_list)
@skipIfNoFBGEMM
def test_dynamic_with_fusion_multiple_uses(self):
"""
Tests that dynamic quantization APIs work with Linear + Relu fusion
"""
with override_quantized_engine('fbgemm'):
class LinearRelu(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(5, 5)
self.relu = torch.nn.ReLU()
def forward(self, x):
x = self.linear(x)
return self.relu(x)
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear_relu = LinearRelu()
def forward(self, x):
x = self.linear_relu(x)
x = self.linear_relu(x)
return x
for qconfig in [float16_dynamic_qconfig, default_dynamic_qconfig]:
model = M().eval()
qconfig_dict = {
"": qconfig
}
example_inputs = (torch.randn(5, 5),)
m = prepare_fx(model, qconfig_dict, example_inputs=example_inputs)
m = convert_fx(m)
m(*example_inputs)
node_list = [
ns.call_module(nniqd.LinearReLU),
ns.call_module(nniqd.LinearReLU),
]
self.checkGraphModuleNodes(m, expected_node_list=node_list)
@skipIfNoFBGEMM
def test_dynamic_linear_input_multiple_use(self):
"""
Tests input for dynamic linear being used by multiple ops
"""
with override_quantized_engine('fbgemm'):
class LinearRelu(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(5, 5)
self.relu = torch.nn.ReLU()
def forward(self, x):
x = self.linear(x)
return self.relu(x)
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.mod1 = LinearRelu()
self.mod2 = LinearRelu()
def forward(self, x):
y1 = self.mod1(x)
y2 = self.mod2(x)
return y1 + y2
for qconfig in [float16_dynamic_qconfig, default_dynamic_qconfig]:
model = M().eval()
qconfig_dict = {
"": qconfig
}
example_inputs = (torch.rand(5, 5, 5),)
m = prepare_fx(model, qconfig_dict, example_inputs=example_inputs)
m = convert_fx(m)
m(*example_inputs)
node_list = [
ns.call_module(nniqd.LinearReLU),
ns.call_module(nniqd.LinearReLU),
]
self.checkGraphModuleNodes(m, expected_node_list=node_list)
def test_ref_linear_module(self):
""" Make sure the numerics for models with ref linear module
matches models with fbgemm/qnnpack module
"""
class M1(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(10, 5)
def forward(self, x):
return self.linear(x)
class M2(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(10, 5)
self.relu = torch.nn.ReLU()
def forward(self, x):
return self.relu(self.linear(x))
for M in [M1, M2]:
m = M().eval()
example_inputs = (torch.randn(5, 10),)
m = prepare_fx(m, {"": default_qconfig}, example_inputs=example_inputs)
m_copy = copy.deepcopy(m)
m = convert_fx(m)
m_ref = convert_to_reference_fx(m_copy)
result = m(*example_inputs)
result_ref = m_ref(*example_inputs)
self.assertTrue(torch.equal(result, result_ref))
def test_ref_conv_module(self):
""" Make sure the numerics for models with ref conv module
matches models with fbgemm/qnnpack module
"""
convs = {
1: nn.Conv1d,
2: nn.Conv2d,
3: nn.Conv3d,
}
class M1(torch.nn.Module):
def __init__(self, dim):
super().__init__()
self.conv = convs[dim](3, 3, 3)
def forward(self, x):
return self.conv(x)
class M2(torch.nn.Module):
def __init__(self, dim):
super().__init__()
self.conv = convs[dim](3, 3, 3)
self.relu = torch.nn.ReLU()
def forward(self, x):
return self.relu(self.conv(x))
for dim, M in itertools.product([1, 2, 3], [M1, M2]):
m = M(dim).eval()
data = self.img_data_dict[dim][0][0]
m = prepare_fx(m, {"": default_qconfig}, example_inputs=(data,))
m_copy = copy.deepcopy(m)
m = convert_fx(m)
m_ref = convert_to_reference_fx(m_copy)
result = m(data)
result_ref = m_ref(data)
self.assertTrue(torch.equal(result, result_ref))
def test_sub_scalar(self):
class M(torch.nn.Module):
def forward(self, x):
x = x + 1
x = x - 1
x = x + 3
x = x - 4
return x
m = M().eval()
m = prepare_fx(m, {"": default_qconfig}, example_inputs=(torch.rand(3),))
m = convert_fx(m)
occurrence = {
ns.call_function(torch.quantize_per_tensor): 2,
ns.call_method("dequantize"): 2
}
self.checkGraphModuleNodes(m, expected_node_occurrence=occurrence)
def test_observer_fqn(self):
"""
Test to make sure the observer FQN is based on the quantizable op/module that it is observing
and uses the modules FQN to determine the observer name.
"""
class Linear(torch.nn.Module):
def __init__(self):
super().__init__()
self.w = torch.ones(5, 5)
self.b = torch.zeros(5)
def forward(self, x):
return torch.nn.functional.linear(x, self.w, self.b)
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.mods1 = torch.nn.Sequential(
Linear(),
Linear()
)
self.mods2 = Linear()
self.mods3 = torch.nn.Linear(5, 5)
def forward(self, x):
x = self.mods1(x)
x = torch.add(x, 4)
x = self.mods2(x)
y = torch.add(x, 2)
z = torch.mul(x, 5)
a = self.mods3(y)
return a, z
model = M().eval()
prepared = prepare_fx(model, {"": default_qconfig}, example_inputs=(torch.randn(1, 5)))
name_list = []
for name, mod in prepared.named_modules():
if isinstance(mod, torch.ao.quantization.observer.MinMaxObserver):
name_list.append(name)
expected_name_list = ['activation_post_process_0',
'activation_post_process_1',
'activation_post_process_2',
'activation_post_process_3',
'activation_post_process_4',
'activation_post_process_6',
'activation_post_process_7',
'activation_post_process_10']
assert name_list == expected_name_list
def test_conv_lowering(self):
convs = {1: nn.Conv1d, 2: nn.Conv2d, 3: nn.Conv3d}
qconvs = {1: nn.quantized.Conv1d, 2: nn.quantized.Conv2d, 3: nn.quantized.Conv3d}
class M(torch.nn.Module):
def __init__(self, dim):
super().__init__()
self.conv = convs[dim](3, 3, 3)
def forward(self, x):
return self.conv(x)
for dim in range(1, len(convs) + 1):
m = M(dim).eval()
data = self.img_data_dict[dim][0][0]
m = prepare_fx(m, {"": default_qconfig}, example_inputs=(data,))
m_ref = copy.deepcopy(m)
m_ref = convert_to_reference_fx(m_ref)
m = convert_fx(m)
out_ref = m_ref(data)
out = m(data)
# check that reference pattern for quantized conv module is fused
expected_node_occurrence = {
ns.call_function(torch.quantize_per_tensor): 1,
ns.call_module(qconvs[dim]): 1,
ns.call_method("dequantize"): 1
}
self.checkGraphModuleNodes(m, expected_node_occurrence=expected_node_occurrence)
# checking result match
self.assertTrue(torch.equal(out_ref, out))
def test_convert_qconfig_mapping(self):
class Linear(torch.nn.Module):
def __init__(self):
super().__init__()
self.w = torch.ones(5, 5)
self.b = torch.zeros(5)
def forward(self, x):
return torch.nn.functional.linear(x, self.w, self.b)
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.mods1 = torch.nn.Sequential(
Linear(),
Linear()
)
self.mods3 = torch.nn.Linear(5, 5)
def forward(self, x):
x = self.mods1(x)
x = torch.add(x, 4)
z = torch.mul(x, 5)
x = self.mods3(z)
return x
model = M().train()
for check in ["module_name", "object_type"]:
qconfig_dict = {"": None,
"object_type": [
(nn.functional.linear, get_default_qat_qconfig("fbgemm")),
(torch.add, get_default_qat_qconfig("fbgemm")),
(nn.Linear, get_default_qat_qconfig("fbgemm")),
],
}
example_inputs = (torch.rand(5, 5),)
prepared = prepare_qat_fx(model, qconfig_dict, example_inputs=example_inputs)
prepared(*example_inputs)
if check == "module_name":
convert_qconfig_dict = {"": None,
"object_type": [
(nn.functional.linear, get_default_qat_qconfig("fbgemm")),
(torch.add, get_default_qat_qconfig("fbgemm")),
(nn.Linear, get_default_qat_qconfig("fbgemm")),
],
"module_name": [("mods1.0", None)]}
node_occurrence = {
ns.call_function(torch.quantize_per_tensor): 2,
ns.call_function(torch.nn.functional.linear): 1,
ns.call_function(torch.ops.quantized.linear): 1,
ns.call_function(torch.ops.quantized.add): 1,
ns.call_method("dequantize"): 2
}
order_check = [
ns.call_function(torch.nn.functional.linear),
ns.call_function(torch.quantize_per_tensor),
ns.call_function(torch.ops.quantized.linear),
ns.call_function(torch.ops.quantized.add),
ns.call_method("dequantize"),
ns.call_function(torch.quantize_per_tensor),
ns.call_module(nnq.Linear),
ns.call_method("dequantize"),
]
elif check == "object_type":
convert_qconfig_dict = {"": None,
"object_type": [
(nn.functional.linear, get_default_qat_qconfig("fbgemm")),
(torch.add, get_default_qat_qconfig("fbgemm")),
(nn.Linear, None),
]}
node_occurrence = {
ns.call_function(torch.quantize_per_tensor): 1,
ns.call_function(torch.ops.quantized.linear): 2,
ns.call_function(torch.ops.quantized.add): 1,
ns.call_function(torch.mul): 1,
ns.call_method("dequantize"): 1
}
order_check = [
ns.call_function(torch.quantize_per_tensor),
ns.call_function(torch.ops.quantized.linear),
ns.call_function(torch.ops.quantized.linear),
ns.call_function(torch.ops.quantized.add),
ns.call_method("dequantize"),
ns.call_function(torch.mul),
ns.call_module(nn.Linear),
]
converted = convert_fx(prepared, qconfig_mapping=convert_qconfig_dict)
converted(torch.rand(5, 5))
self.checkGraphModuleNodes(
converted,
expected_node_occurrence=node_occurrence,
expected_node_list=order_check)
def _assertFixedQParamsFakeQuantizeEqual(self, fq1, fq2):
self.assertEqual(fq1()._observer_ctr, fq2()._observer_ctr)
def test_register_patterns(self):
def cleanUp():
del _DEFAULT_FUSION_PATTERNS["dummy_fusion"]
del _DEFAULT_QUANTIZATION_PATTERNS["dummy_quant"]
del _DEFAULT_QUANTIZATION_PATTERNS["dummy_quant2"]
del _DEFAULT_QUANTIZATION_PATTERNS["dummy_quant3"]
del _DEFAULT_OUTPUT_OBSERVER_MAP["dummy_quant2"]
del _DEFAULT_OUTPUT_OBSERVER_MAP["dummy_quant3"]
del _DEFAULT_OUTPUT_FAKE_QUANTIZE_MAP["dummy_quant2"]
del _DEFAULT_OUTPUT_FAKE_QUANTIZE_MAP["dummy_quant3"]
self.addCleanup(cleanUp)
@_register_fusion_pattern("dummy_fusion")
class DummyFusion:
pass
@_register_quant_pattern("dummy_quant")
class DummyQuant:
pass
@_register_quant_pattern("dummy_quant2", default_fixed_qparams_range_0to1_observer)
class DummyQuant2:
pass
@_register_quant_pattern("dummy_quant3", default_fixed_qparams_range_neg1to1_observer)
class DummyQuant3:
pass
self.assertEqual(_DEFAULT_FUSION_PATTERNS["dummy_fusion"], DummyFusion)
self.assertEqual(_DEFAULT_QUANTIZATION_PATTERNS["dummy_quant"], DummyQuant)
self.assertEqual(_DEFAULT_QUANTIZATION_PATTERNS["dummy_quant2"], DummyQuant2)
self.assertEqual(_DEFAULT_QUANTIZATION_PATTERNS["dummy_quant3"], DummyQuant3)
self.assertEqual(_DEFAULT_OUTPUT_OBSERVER_MAP["dummy_quant2"], default_fixed_qparams_range_0to1_observer)
self.assertEqual(_DEFAULT_OUTPUT_OBSERVER_MAP["dummy_quant3"], default_fixed_qparams_range_neg1to1_observer)
self._assertFixedQParamsFakeQuantizeEqual(_DEFAULT_OUTPUT_FAKE_QUANTIZE_MAP["dummy_quant2"],
default_fixed_qparams_range_0to1_fake_quant)
self._assertFixedQParamsFakeQuantizeEqual(_DEFAULT_OUTPUT_FAKE_QUANTIZE_MAP["dummy_quant3"],
default_fixed_qparams_range_neg1to1_fake_quant)
output_fake_quantize_map = get_default_output_activation_post_process_map(is_training=True)
output_observer_map = get_default_output_activation_post_process_map(is_training=False)
self.assertEqual(output_observer_map.get("dummy_quant3"), default_fixed_qparams_range_neg1to1_observer)
self._assertFixedQParamsFakeQuantizeEqual(output_fake_quantize_map.get("dummy_quant3"),
default_fixed_qparams_range_neg1to1_fake_quant)
def test_reuse_input_qconfig(self):
class M1(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv2d(3, 3, 3)
def forward(self, x):
x = self.conv(x)
x = x.reshape()
return x
class M2(torch.nn.Module):
def forward(self, x):
x = x.reshape()
return x
options = itertools.product([M1, M2], [True, False])
for M, is_qat in options:
m = M1().eval()
example_inputs = (torch.randn(1, 3, 3, 3),)
m = prepare_fx(m, get_default_qconfig_mapping(), example_inputs=example_inputs)
m = convert_fx(m)
node_list = [
ns.call_function(torch.quantize_per_tensor),
ns.call_module(nnq.Conv2d),
ns.call_method("reshape"),
ns.call_method("dequantize"),
]
self.checkGraphModuleNodes(
m,
expected_node_list=node_list)
m = M2().eval()
m = prepare_fx(m, get_default_qconfig_mapping(), example_inputs=example_inputs)
m = convert_fx(m)
node_occurrence = {
ns.call_function(torch.quantize_per_tensor): 0,
ns.call_method("dequnatize"): 0,
}
node_list = [
ns.call_method("reshape"),
]
self.checkGraphModuleNodes(
m,
expected_node_occurrence=node_occurrence,
expected_node_list=node_list)
def test_stack_trace_preserved_linear(self):
class M(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(1, 1)
def forward(self, x):
x = self.linear(x)
return x
m = M().eval()
mp = prepare_fx(m, get_default_qconfig_mapping(), example_inputs=(torch.randn(1, 1),))
found_stack_trace = False
for n in mp.graph.nodes:
if n.op == 'call_module' and n.target == 'linear':
found_stack_trace = n.stack_trace is not None
break
self.assertTrue(found_stack_trace)
# test reference model
mq = convert_to_reference_fx(copy.deepcopy(mp))
found_stack_trace = False
for n in mq.graph.nodes:
if n.op == 'call_module' and n.target == 'linear':
found_stack_trace = n.stack_trace is not None
break
self.assertTrue(found_stack_trace, f"stack trace not found, node: {n.format_node()}, is_reference: True")
# test quantized model
mq = convert_fx(mp)
found_stack_trace = False
for n in mq.graph.nodes:
if n.op == 'call_module' and n.target == 'linear':
found_stack_trace = n.stack_trace is not None
break
self.assertTrue(found_stack_trace, f"stack trace not found, node: {n.format_node()}, is_reference: False")
def test_qat_skip_untraced(self):
class UnTraceableModuleClass(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(2, 2)
def forward(self, x):
return self.linear(x)
class UnTraceableModuleName(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(2, 2)
def forward(self, x):
return self.linear(x)
class M(nn.Module):
def __init__(self):
super().__init__()
self.untraceable_module_class = UnTraceableModuleClass()
self.untraceable_module_name = UnTraceableModuleClass()
def forward(self, x):
x = self.untraceable_module_class(x)
x = self.untraceable_module_name(x)
return x
mod = M()
qconfig_dict = {"": torch.ao.quantization.get_default_qat_qconfig()}
prepare_custom_config_dict = {
"non_traceable_module_class": [UnTraceableModuleClass],
"non_traceable_module_name": ["untraceable_module_name"],
}
example_inputs = (torch.randn(2, 2),)
mod_prep = torch.ao.quantization.quantize_fx.prepare_qat_fx(
mod.train(), qconfig_dict, example_inputs=example_inputs,
prepare_custom_config=prepare_custom_config_dict
)
mod_prep = torch.ao.quantization.quantize_fx.prepare_qat_fx(
mod.train(), qconfig_dict, example_inputs=example_inputs,
prepare_custom_config=prepare_custom_config_dict
)
self.assertTrue(
isinstance(mod_prep.untraceable_module_class.linear, torch.nn.Linear)
)
self.assertTrue(
isinstance(mod_prep.untraceable_module_name.linear, torch.nn.Linear)
)
self.assertTrue(
type(mod_prep.untraceable_module_class.linear)
is not torch.ao.nn.qat.modules.linear.Linear,
"prepare_qat_fx shold not convert anything inside untraced module classes",
)
self.assertTrue(
type(mod_prep.untraceable_module_name.linear)
is not torch.ao.nn.qat.modules.linear.Linear,
"prepare_qat_fx shold not convert anything inside modules named in untraced_module_names",
)
def test_qconfig_dict_setup(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.Conv1d = torch.nn.Conv1d(1, 1, 1)
self.Conv2d = torch.nn.Conv2d(1, 1, 1)
self.Conv3d = torch.nn.Conv3d(1, 1, 1)
self.ConvTranspose1d = torch.nn.ConvTranspose1d(1, 1, 1)
self.ConvTranspose2d = torch.nn.ConvTranspose2d(1, 1, 1)
self.ConvTranspose3d = torch.nn.ConvTranspose3d(1, 1, 1)
self.Linear = torch.nn.Linear(1, 1, 1)
def forward(self, x):
x = self.Conv1d(x)
x = self.Conv2d(x)
x = self.Conv3d(x)
x = self.ConvTranspose1d(x)
x = self.ConvTranspose2d(x)
x = self.ConvTranspose3d(x)
x = self.Linear(x)
x = torch.nn.functional.conv1d(x, torch.rand(2, 2))
x = torch.nn.functional.conv2d(x, torch.rand(2, 2))
x = torch.nn.functional.conv3d(x, torch.rand(2, 2))
x = torch.nn.functional.linear(x, torch.rand(2, 2))
return x
backends = ["qnnpack", "fbgemm"]
for func in [get_default_qconfig_mapping, get_default_qat_qconfig_mapping]:
for backend in backends:
m = M().eval()
qconfig_dict = func(backend)
m = prepare_fx(m, qconfig_dict, example_inputs=(torch.randn(1, 1, 1, 1)))
for name, mod in m.named_modules():
if _is_activation_post_process(mod) and mod.dtype == torch.quint8:
if backend == "fbgemm":
lower_bnd = 0
upper_bnd = 127
else:
lower_bnd = 0
upper_bnd = 255
if issubclass(type(mod), FakeQuantize):
self.assertEqual(mod.activation_post_process.quant_min, lower_bnd)
self.assertEqual(mod.activation_post_process.quant_max, upper_bnd)
else:
self.assertEqual(mod.quant_min, lower_bnd)
self.assertEqual(mod.quant_max, upper_bnd)
def test_prepare_mode(self):
class LinearModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(5, 10)
def forward(self, x):
return self.linear(x)
def _test(prepare_fn, qconfig_dict):
m = LinearModel()
m1 = copy.deepcopy(m)
m1.train()
example_inputs = (torch.randn(1, 5),)
prepare_fn(m1, qconfig_dict, example_inputs=example_inputs)
m2 = copy.deepcopy(m)
m2.eval()
prepare_fn(m2, qconfig_dict, example_inputs=example_inputs)
# Ensure prepare_fx and prepare_qat_fx work in both training and eval modes
_test(prepare_fx, get_default_qconfig_mapping())
_test(prepare_qat_fx, get_default_qat_qconfig_mapping())
def _validate_qconfig_against_backend_config_constraints(
self,
model: torch.nn.Module,
qconfig: QConfig,
backend_config: BackendConfig,
satisfies_constraints: bool,
qconfig_name: Optional[str] = None):
"""
Helper method to validate whether `qconfig` satisfies the constraints specified in `backend_config`.
"""
qconfig_mapping = QConfigMapping().set_object_type(torch.nn.Linear, qconfig)
example_inputs = (torch.rand((1, 30), dtype=torch.float),)
model = prepare_fx(model, qconfig_mapping, example_inputs, backend_config=backend_config)
model(*example_inputs)
model = convert_fx(model, backend_config=backend_config)
if satisfies_constraints:
expected_node_occurrence = {
ns.call_module(torch.ao.nn.quantized.Linear) : 1,
ns.call_module(torch.nn.Linear) : 0,
}
else:
expected_node_occurrence = {
ns.call_module(torch.ao.nn.quantized.Linear) : 0,
ns.call_module(torch.nn.Linear) : 1,
}
try:
self.checkGraphModuleNodes(model, expected_node_occurrence=expected_node_occurrence)
except AssertionError as e:
if qconfig_name is not None:
print("ERROR: Validation for QConfig '%s' failed" % qconfig_name)
raise e
def test_backend_config_quantization_range(self):
"""
Check that quantization ranges specified through the BackendConfig are reflected in
the observers inserted into the model.
"""
class MyModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(30, 4).float()
def forward(self, x):
return self.linear(x)
dtype_config = DTypeConfig(
input_dtype=DTypeWithConstraints(
dtype=torch.quint8,
quant_min_lower_bound=0,
quant_max_upper_bound=31,
),
output_dtype=DTypeWithConstraints(
dtype=torch.quint8,
quant_min_lower_bound=0,
quant_max_upper_bound=31,
),
weight_dtype=DTypeWithConstraints(
dtype=torch.qint8,
quant_min_lower_bound=-64,
quant_max_upper_bound=63,
),
bias_dtype=torch.float,
)
backend_config = BackendConfig() \
.set_backend_pattern_config(BackendPatternConfig(torch.nn.Linear)
.set_observation_type(ObservationType.OUTPUT_USE_DIFFERENT_OBSERVER_AS_INPUT) # noqa: E128
.add_dtype_config(dtype_config)
.set_root_module(torch.nn.Linear)
.set_reference_quantized_module(nnqr.Linear))
def validate_qconfig(qconfig: QConfig, satisfies_constraints: bool):
self._validate_qconfig_against_backend_config_constraints(
MyModel(), qconfig, backend_config, satisfies_constraints)
# Case 1: QConfig ranges fit within backend ranges, OK
qconfig1 = QConfig(
activation=MinMaxObserver.with_args(quant_min=0, quant_max=15, dtype=torch.quint8),
weight=MinMaxObserver.with_args(quant_min=-32, quant_max=31, dtype=torch.qint8, qscheme=torch.per_tensor_symmetric))
validate_qconfig(qconfig1, satisfies_constraints=True)
# Case 2: QConfig activation range falls outside backend range, should fail
qconfig2 = QConfig(
activation=MinMaxObserver.with_args(quant_min=0, quant_max=63, dtype=torch.quint8),
weight=MinMaxObserver.with_args(dtype=torch.qint8, qscheme=torch.per_tensor_symmetric))
validate_qconfig(qconfig2, satisfies_constraints=False)
# Case 3: QConfig weight range falls outside backend range, should fail
qconfig3 = QConfig(
activation=MinMaxObserver.with_args(dtype=torch.quint8),
weight=MinMaxObserver.with_args(quant_min=-128, quant_max=127, dtype=torch.qint8, qscheme=torch.per_tensor_symmetric))
validate_qconfig(qconfig3, satisfies_constraints=False)
# Case 4: QConfig doesn't specify range, should fail
qconfig4 = QConfig(activation=ReuseInputObserver, weight=ReuseInputObserver)
validate_qconfig(qconfig4, satisfies_constraints=False)
def test_backend_config_scale_min(self):
"""
Test QConfig eps validation against the BackendConfig's min scale value.
"""
class MyModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(30, 4).float()
def forward(self, x):
return self.linear(x)
dtype_config = DTypeConfig(
input_dtype=DTypeWithConstraints(dtype=torch.quint8, scale_min_lower_bound=2 ** -12),
output_dtype=DTypeWithConstraints(dtype=torch.quint8, scale_min_lower_bound=2 ** -12),
weight_dtype=DTypeWithConstraints(dtype=torch.qint8, scale_min_lower_bound=2 ** -12),
bias_dtype=torch.float,
)
backend_config = BackendConfig() \
.set_backend_pattern_config(BackendPatternConfig(torch.nn.Linear)
.set_observation_type(ObservationType.OUTPUT_USE_DIFFERENT_OBSERVER_AS_INPUT) # noqa: E128
.add_dtype_config(dtype_config)
.set_root_module(torch.nn.Linear)
.set_reference_quantized_module(nnqr.Linear))
def validate_qconfig(qconfig: QConfig, satisfies_constraints: bool):
self._validate_qconfig_against_backend_config_constraints(
MyModel(), qconfig, backend_config, satisfies_constraints)
# Case 1: QConfig min scale value == backend min scale value, OK
qconfig1 = QConfig(
activation=MinMaxObserver.with_args(dtype=torch.quint8, eps=2 ** -12),
weight=MinMaxObserver.with_args(dtype=torch.qint8, qscheme=torch.per_tensor_symmetric, eps=2 ** -12))
validate_qconfig(qconfig1, satisfies_constraints=True)
# Case 2: QConfig min scale value > backend min scale value, OK
qconfig2 = QConfig(
activation=MinMaxObserver.with_args(dtype=torch.quint8, eps=2 ** -10),
weight=MinMaxObserver.with_args(dtype=torch.qint8, qscheme=torch.per_tensor_symmetric, eps=2 ** -10))
validate_qconfig(qconfig2, satisfies_constraints=True)
# Case 3: QConfig activation min scale value < backend min scale value, should fail
qconfig3 = QConfig(
activation=MinMaxObserver.with_args(dtype=torch.quint8, eps=2 ** -14),
weight=MinMaxObserver.with_args(dtype=torch.qint8, qscheme=torch.per_tensor_symmetric))
validate_qconfig(qconfig3, satisfies_constraints=False)
# Case 3: QConfig weight min scale value < backend min scale value, should fail
qconfig4 = QConfig(
activation=MinMaxObserver.with_args(dtype=torch.quint8),
weight=MinMaxObserver.with_args(dtype=torch.qint8, qscheme=torch.per_tensor_symmetric, eps=2 ** -14))
validate_qconfig(qconfig4, satisfies_constraints=False)
# Case 5: QConfig doesn't specify eps, should fail
qconfig5 = QConfig(
activation=FixedQParamsObserver.with_args(scale=1.0, zero_point=0),
weight=FixedQParamsObserver.with_args(scale=1.0, zero_point=0))
validate_qconfig(qconfig5, satisfies_constraints=False)
def test_qnnpack_backend_config(self):
"""
Test whether default QNNPACK QConfigs are compatible with the QNNPACK BackendConfig.
"""
class MyModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(30, 4).float()
def forward(self, x):
return self.linear(x)
all_qconfigs: List[Tuple[QConfig, str]] = [
(get_default_qconfig("qnnpack", version=0), "default_qnnpack_qconfig_v0"),
(get_default_qat_qconfig("qnnpack", version=0), "default_qat_qnnpack_qconfig_v0"),
(get_default_qat_qconfig("qnnpack", version=1), "default_qat_qnnpack_qconfig_v1"),
(default_symmetric_qnnpack_qconfig, "default_symmetric_qnnpack_qconfig"),
(default_symmetric_qnnpack_qat_qconfig, "default_symmetric_qnnpack_qat_qconfig"),
# TODO: Test these QConfigs once they are fixed, see https://github.com/pytorch/pytorch/issues/85862
# (default_per_channel_symmetric_qnnpack_qconfig, "default_per_channel_symmetric_qnnpack_qconfig"),
# (default_per_channel_symmetric_qnnpack_qat_qconfig, "default_per_channel_symmetric_qnnpack_qat_qconfig"),
]
backend_config = get_qnnpack_backend_config()
for qconfig, qconfig_name in all_qconfigs:
self._validate_qconfig_against_backend_config_constraints(
MyModel(), qconfig, backend_config, satisfies_constraints=True, qconfig_name=qconfig_name)
def test_symmetric_qnnpack_qconfig_mapping(self):
"""
Test whether `torch.ao.quantization.qconfig_mapping._get_symmetric_qnnpack_qconfig_mapping`
works with the QNNPACK BackendConfig.
"""
if "qnnpack" not in supported_qengines:
return
class MyModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(30, 4).float()
def forward(self, x):
return self.linear(x)
with override_quantized_engine("qnnpack"):
qconfig_mapping = _get_symmetric_qnnpack_qconfig_mapping()
example_inputs = (torch.rand((1, 30), dtype=torch.float),)
backend_config = get_qnnpack_backend_config()
model = MyModel()
model = prepare_fx(model, qconfig_mapping, example_inputs, backend_config=backend_config)
model(*example_inputs)
model = convert_fx(model, backend_config=backend_config)
expected_node_occurrence = {
ns.call_module(torch.ao.nn.quantized.Linear) : 1,
ns.call_module(torch.nn.Linear) : 0,
}
self.checkGraphModuleNodes(model, expected_node_occurrence=expected_node_occurrence)
model(*example_inputs)
def test_symmetric_qnnpack_qat_qconfig_mapping(self):
"""
Test whether `torch.ao.quantization.qconfig_mapping._get_symmetric_qnnpack_qat_qconfig_mapping`
works with the QNNPACK BackendConfig.
"""
if "qnnpack" not in supported_qengines:
return
class MyModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(30, 4).float()
def forward(self, x):
return self.linear(x)
with override_quantized_engine("qnnpack"):
qconfig_mapping = _get_symmetric_qnnpack_qat_qconfig_mapping()
example_inputs = (torch.rand((1, 30), dtype=torch.float),)
backend_config = get_qnnpack_backend_config()
model = MyModel()
model = prepare_fx(model, qconfig_mapping, example_inputs, backend_config=backend_config)
model(*example_inputs)
model = convert_fx(model, backend_config=backend_config)
expected_node_occurrence = {
ns.call_module(torch.ao.nn.quantized.Linear) : 1,
ns.call_module(torch.nn.Linear) : 0,
}
self.checkGraphModuleNodes(model, expected_node_occurrence=expected_node_occurrence)
model(*example_inputs)
def test_get_executorch_backend_config(self):
from torch.ao.quantization.backend_config import get_executorch_backend_config
# make sure this runs
executorch_backend_config = get_executorch_backend_config()
def test_backend_config_check_for_weight_and_bias(self):
""" Test to make sure the backend_config check for weight and bias
runs when the qconfig is None for the ops with weight and bias
previously the error was not hit because we first check input, and
the check for weight and bias are skipped.
"""
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.weight = torch.tensor((5, 5))
self.bias = torch.tensor((5,))
def forward(self, x):
return torch.addmm(self.bias, x, self.weight)
m = M().eval()
qconfig_mapping = QConfigMapping()
observation_type = ObservationType.OUTPUT_USE_DIFFERENT_OBSERVER_AS_INPUT
weighted_op_quint8_dtype_config = DTypeConfig(
input_dtype=torch.quint8,
output_dtype=torch.quint8,
weight_dtype=torch.qint8,
bias_dtype=torch.float,
)
dtype_configs = [weighted_op_quint8_dtype_config]
backend_pattern_config = BackendPatternConfig(torch.addmm) \
.set_observation_type(observation_type) \
.set_dtype_configs(dtype_configs) \
._set_input_type_to_index({"weight": 2, "bias": 0})
backend_config = BackendConfig() \
.set_backend_pattern_config(backend_pattern_config)
example_inputs = (torch.rand(1, 5),)
# make sure this runs
m = prepare_fx(m, qconfig_mapping, example_inputs, backend_config=backend_config)
def test_get_default_qconfig_valid_backend(self):
""" Checks that AssertionError is raised when non expected backend input is specified
"""
invalid_backends = ["imaginary_backend", 3]
for invalid_backend in invalid_backends:
with self.assertRaisesRegex(AssertionError, "not supported"):
qconfig = get_default_qconfig(invalid_backend)
with self.assertRaisesRegex(AssertionError, "not supported"):
qconfig = get_default_qat_qconfig(invalid_backend)
with self.assertRaisesRegex(AssertionError, "not supported"):
qconfig_mapping = get_default_qconfig_mapping(invalid_backend)
with self.assertRaisesRegex(AssertionError, "not supported"):
qconfig_mapping = get_default_qat_qconfig_mapping(invalid_backend)
def test__convert_to_reference_decomposed_fx(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(5, 10)
def forward(self, x):
return self.linear(x)
m = M().eval()
qconfig_mapping = get_default_qconfig_mapping("fbgemm")
example_inputs = (torch.randn(1, 5),)
m = prepare_fx(m, qconfig_mapping, example_inputs)
m_ref = copy.deepcopy(m)
m_ref = convert_to_reference_fx(m_ref)
m = _convert_to_reference_decomposed_fx(m)
expected_occurrence = {
ns.call_function(torch.ops.quantized_decomposed.quantize_per_tensor.default): 2,
ns.call_function(torch.ops.quantized_decomposed.dequantize_per_tensor.default): 2,
}
self.checkGraphModuleNodes(
m,
expected_node_occurrence=expected_occurrence)
# make sure it runs
res_ref = m_ref(*example_inputs)
res = m(*example_inputs)
self.assertEqual(res, res_ref)
@skipIfNoQNNPACK
def test__convert_to_reference_decomposed_fx_dynamic_quant(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(5, 10)
def forward(self, x):
return self.linear(x)
# to avoid reduce_range
with override_quantized_engine("qnnpack"):
m = M().eval()
qconfig_mapping = get_default_qconfig_mapping("fbgemm") \
.set_object_type(torch.nn.Linear, default_dynamic_qconfig)
example_inputs = (torch.randn(1, 5),)
m = prepare_fx(m, qconfig_mapping, example_inputs)
m(*example_inputs)
m_ref = copy.deepcopy(m)
m_ref = convert_to_reference_fx(m_ref)
m = _convert_to_reference_decomposed_fx(m)
expected_occurrence = {
ns.call_function(torch.ops.quantized_decomposed.choose_qparams.tensor): 1,
ns.call_function(torch.ops.quantized_decomposed.quantize_per_tensor.tensor): 1,
ns.call_function(torch.ops.quantized_decomposed.dequantize_per_tensor.tensor): 1,
}
self.checkGraphModuleNodes(
m,
expected_node_occurrence=expected_occurrence)
# make sure it runs
res_ref = m_ref(*example_inputs)
res = m(*example_inputs)
self.assertEqual(res, res_ref)
def test__convert_to_reference_decomposed_fx_per_channel_quant(self):
class M(torch.nn.Module):
def forward(self, x, weight, bias):
return F.linear(x, weight, bias)
m = M().eval()
qconfig_mapping = get_default_qconfig_mapping("fbgemm") \
.set_object_type(F.linear, default_per_channel_qconfig)
example_inputs = (torch.randn(1, 5), torch.randn(10, 5), torch.randn(10,))
m = prepare_fx(m, qconfig_mapping, example_inputs)
m(*example_inputs)
m_ref = copy.deepcopy(m)
m_ref = convert_to_reference_fx(m_ref)
m = _convert_to_reference_decomposed_fx(m)
expected_occurrence = {
# for input and output activations
ns.call_function(torch.ops.quantized_decomposed.quantize_per_tensor.default): 2,
ns.call_function(torch.ops.quantized_decomposed.dequantize_per_tensor.default): 2,
# for weight
ns.call_function(torch.ops.quantized_decomposed.quantize_per_channel.default): 1,
ns.call_function(torch.ops.quantized_decomposed.dequantize_per_channel.default): 1,
}
self.checkGraphModuleNodes(
m,
expected_node_occurrence=expected_occurrence)
# make sure it runs
res_ref = m_ref(*example_inputs)
res = m(*example_inputs)
self.assertEqual(res, res_ref)
def test_change_backend_config_for_fixed_qparam_ops(self):
""" Making sure we can skip validation of qconfigs for fixedqparam ops based
on BackendConfig
"""
class M(nn.Module):
def __init__(self):
super().__init__()
self.tanh = torch.nn.Tanh()
def forward(self, x: torch.Tensor):
x = self.tanh(x)
return x
model = M().eval()
# we set a global default_qconfig, which will be ignored since the backend
# we defined doesn't support anything
# this is to make sure we don't validate the qconfig when BackendConfig does not
# have fixed qparam op related configurations
qconfig_mapping = QConfigMapping().set_global(default_qconfig)
backend_config = BackendConfig()
# make sure this runs
model = prepare_fx(
model,
qconfig_mapping=qconfig_mapping,
example_inputs=(torch.randn(1, 2, 3, 4),),
backend_config=backend_config
)
def test_channel_shuffle_lowering(self):
# Three versions of channel shuffle
class M1(torch.nn.Module):
def __init__(self):
super().__init__()
self.op = torch.nn.ChannelShuffle(2)
def forward(self, x):
return self.op(x + x) + x
class M2(torch.nn.Module):
def forward(self, x):
return torch.channel_shuffle(x + x, 2) + x
class M3(torch.nn.Module):
def forward(self, x):
return torch.nn.functional.channel_shuffle(x + x, 2) + x
x = torch.randn(4, 4, 4, 4)
# torch.channel_shuffle is equivalent to torch.nn.functional.channel_shuffle
model_node_pairs = [
(M1().eval(), ns.call_module(torch.nn.ChannelShuffle)),
(M2().eval(), ns.call_function(torch.channel_shuffle)),
(M3().eval(), ns.call_function(torch.channel_shuffle))
]
for m, node in model_node_pairs:
m = prepare_fx(m, {"": default_qconfig}, example_inputs=(x,))
m_copy = copy.deepcopy(m)
m = convert_fx(m)
m_ref = convert_to_reference_fx(m_copy)
node_occurrence = {
node: 1,
ns.call_function(torch.quantize_per_tensor): 1,
ns.call_method("dequantize"): 1
}
node_occurrence_ref = {
node: 1,
ns.call_function(torch.quantize_per_tensor): 4,
ns.call_method("dequantize"): 4
}
self.checkGraphModuleNodes(m, expected_node_occurrence=node_occurrence)
self.checkGraphModuleNodes(m_ref, expected_node_occurrence=node_occurrence_ref)
def test_match_pattern_with_multiple_args(self):
""" Test that we can match a pattern that has multiple arguments
Pattern:
shape \
transpose (observed) -> reshape -> output (observed) ->
where `reshape` has two arguments
"""
def _get_pattern_configs():
backend_pattern_configs = []
observation_type = ObservationType.OUTPUT_SHARE_OBSERVER_WITH_INPUT
weighted_op_quint8_dtype_config = DTypeConfig(
input_dtype=torch.quint8,
output_dtype=torch.quint8,
weight_dtype=torch.qint8,
bias_dtype=torch.float,
)
dtype_configs = [weighted_op_quint8_dtype_config]
def root_node_getter(node_pattern):
reshape, transpose, shape = node_pattern
return transpose
backend_pattern_configs.append(
BackendPatternConfig()
._set_pattern_complex_format((torch.reshape, torch.transpose, MatchAllNode)) # noqa: E131
.set_observation_type(observation_type)
.set_dtype_configs(dtype_configs)
._set_root_node_getter(root_node_getter)
)
return backend_pattern_configs
backend_config = BackendConfig().set_backend_pattern_configs(_get_pattern_configs())
class M(torch.nn.Module):
def forward(self, x):
x = torch.transpose(x, 0, 1)
x = torch.reshape(x, (-1,))
return x
m = M().eval()
qconfig_mapping = QConfigMapping().set_global(default_qconfig)
example_inputs = (torch.randn(1, 3, 3, 3),)
m = prepare_fx(m, qconfig_mapping, example_inputs, backend_config=backend_config)
node_occurrence = {
# one for input of the pattern and one for output of the pattern
ns.call_module(MinMaxObserver): 2
}
self.checkGraphModuleNodes(m, expected_node_occurrence=node_occurrence)
def _test_linear_activation_fusion_lowering_helper(
self, module, example_inputs, qconfig_mapping,
backend_config, fused_module, root_module, activation_module):
node_occurrence = {
ns.call_function(torch.quantize_per_tensor): 1,
ns.call_method("dequantize"): 1,
ns.call_module(fused_module): 1,
ns.call_module(root_module): 0,
ns.call_module(activation_module): 0,
}
node_occurrence_ref = {
ns.call_function(torch.quantize_per_tensor): 2,
ns.call_method("dequantize"): 2,
}
m = module.eval()
m = prepare_fx(m, qconfig_mapping,
example_inputs=example_inputs,
backend_config=backend_config)
m_copy = copy.deepcopy(m)
m = convert_fx(m, backend_config=backend_config)
m_ref = convert_to_reference_fx(m_copy)
self.checkGraphModuleNodes(m, expected_node_occurrence=node_occurrence)
self.checkGraphModuleNodes(m_ref, expected_node_occurrence=node_occurrence_ref)
m(*example_inputs)
@skipIfNoONEDNN
def test_linear_leaky_relu_lowering(self):
""" Test fusion and lowering of Linear - (bn -) LeakyReLU
by FX. For onednn backedn only.
"""
from torch.ao.quantization.backend_config import get_onednn_backend_config
qconfig_mapping = get_default_qconfig_mapping('onednn')
with override_quantized_engine('onednn'):
for with_bn in [True, False]:
m = LinearBnLeakyReluModel(with_bn)
self._test_linear_activation_fusion_lowering_helper(
m,
m.get_example_inputs(),
qconfig_mapping,
get_onednn_backend_config(),
nniq.LinearLeakyReLU,
nn.Linear,
nn.LeakyReLU)
@skipIfNoONEDNN
def test_linear_tanh_lowering(self):
""" Test fusion and lowering of Linear - Tanh
by FX. For onednn backedn only.
"""
from torch.ao.quantization.backend_config import get_onednn_backend_config
qconfig_mapping = get_default_qconfig_mapping('onednn')
# TODO Currently it's required that separate ops in a fused op/module have the same qconfig.
# Need to be able to support fusion of ops with different qconfigs
# Since tanh must have 'fixed_qparams_qconfig' while linear should use
# the global qconfig, we need to set qconfigs for them manually here for
# fusion and cannot put such configs in onednn's default qconfig_mapping.
# Known issue:
# Cannot fuse linear - tanh and quantize standalone tanh at the same time.
qconfig = get_default_qconfig('onednn')
qconfig_mapping.set_object_type(torch.nn.Linear, qconfig)
qconfig_mapping.set_object_type(torch.nn.Tanh, qconfig)
with override_quantized_engine('onednn'):
m = LinearTanhModel()
self._test_linear_activation_fusion_lowering_helper(
m,
m.get_example_inputs(),
qconfig_mapping,
get_onednn_backend_config(),
nniq.LinearTanh,
nn.Linear,
nn.Tanh)
@override_qengines
def test_linear_size_view(self):
class M(torch.nn.Module):
def __init__(self, use_relu=False):
super().__init__()
self.linear = torch.nn.Linear(16, 32)
self.relu = torch.nn.ReLU()
self.use_relu = use_relu
def forward(self, x):
x = self.linear(x)
if self.use_relu:
x = self.relu(x)
return x.view(x.size(0), 1, 4, 8)
for use_relu in [False, True]:
model_fp32 = M(use_relu).eval()
qengine = torch.backends.quantized.engine
qconfig_mapping = get_default_qconfig_mapping(qengine)
x = torch.randn((5, 16))
model_fp32(x)
prepared_model = prepare_fx(model_fp32, qconfig_mapping, x)
prepared_model(x)
quantized_model = convert_fx(prepared_model)
node_occurrence = {
ns.call_module(nnq.Linear): 0 if use_relu else 1,
ns.call_module(nniq.LinearReLU): 1 if use_relu else 0,
ns.call_function(torch.quantize_per_tensor): 1,
ns.call_method("dequantize"): 1
}
self.checkGraphModuleNodes(quantized_model, expected_node_occurrence=node_occurrence)
@override_qengines
def test_linear_shape_view(self):
class M(torch.nn.Module):
def __init__(self, use_relu=False):
super().__init__()
self.linear = torch.nn.Linear(16, 32)
self.relu = torch.nn.ReLU()
self.use_relu = use_relu
def forward(self, x):
x = self.linear(x)
if self.use_relu:
x = self.relu(x)
return x.view(x.shape[0], 1, 4, 8)
for use_relu in [False, True]:
model_fp32 = M(use_relu).eval()
qengine = torch.backends.quantized.engine
qconfig_mapping = get_default_qconfig_mapping(qengine)
x = torch.randn((5, 16))
model_fp32(x)
prepared_model = prepare_fx(model_fp32, qconfig_mapping, x)
prepared_model(x)
quantized_model = convert_fx(prepared_model)
node_occurrence = {
ns.call_module(nnq.Linear): 0 if use_relu else 1,
ns.call_module(nniq.LinearReLU): 1 if use_relu else 0,
ns.call_function(torch.quantize_per_tensor): 1,
ns.call_method("dequantize"): 1
}
self.checkGraphModuleNodes(quantized_model, expected_node_occurrence=node_occurrence)
def test_mixed_dtypes(self):
"""
Test that multiple dtypes can be used in the same model for different layers,
and the dtypes will be converted correctly between the layers.
"""
class MyModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear1 = torch.nn.Linear(5, 5)
self.linear2 = torch.nn.Linear(5, 5)
self.sigmoid = torch.nn.Sigmoid()
self.tanh = torch.nn.Tanh()
self.float_functional = torch.ao.nn.quantized.FloatFunctional()
def forward(self, x: torch.Tensor):
x = self.linear1(x) # qint32
x = self.linear2(x) # quint8
linear2 = x
x = self.sigmoid(x) # back to qint32
x = self.tanh(x) # back to quint8
x = self.float_functional.add(linear2, x) # adding two quint8's together
return x
def make_qconfig(scale, zp, dtype):
return QConfig(
activation=FixedQParamsObserver.with_args(scale=scale, zero_point=zp, dtype=dtype),
weight=torch.ao.quantization.default_weight_observer)
# Set up a QConfigMapping that specifies different qparams and dtypes for different layers
qconfig_mapping = QConfigMapping() \
.set_global(get_default_qconfig("qnnpack")) \
.set_module_name("linear1", make_qconfig(1234, 11, torch.qint32)) \
.set_module_name("linear2", make_qconfig(2345, 22, torch.quint8)) \
.set_object_type(torch.nn.Sigmoid, make_qconfig(3456, 33, torch.qint32)) \
.set_object_type(torch.nn.Tanh, make_qconfig(4567, 44, torch.quint8))
# Set up BackendConfig that supports the dtypes configured in the above QConfigMapping
weighted_op_qint32_dtype_config = DTypeConfig(
input_dtype=torch.qint32,
output_dtype=torch.qint32,
weight_dtype=torch.qint8,
bias_dtype=torch.float,
)
fixed_qparams_op_quint8_dtype_config = DTypeConfig(
input_dtype=torch.quint8,
output_dtype=torch.quint8,
)
fixed_qparams_op_qint32_dtype_config = DTypeConfig(
input_dtype=torch.qint32,
output_dtype=torch.qint32,
)
backend_config = get_qnnpack_backend_config()
for config in backend_config.configs:
if config.pattern == torch.nn.Linear:
config.add_dtype_config(weighted_op_qint32_dtype_config)
elif config.pattern in [torch.nn.Sigmoid, torch.nn.Tanh]:
config.add_dtype_config(fixed_qparams_op_quint8_dtype_config)
config.add_dtype_config(fixed_qparams_op_qint32_dtype_config)
# Produce the reference quantized model
m = MyModule()
example_inputs = (torch.rand(5, 5),)
prepared = prepare_fx(m, qconfig_mapping, example_inputs, backend_config=backend_config)
prepared(*example_inputs) # calibrate
converted = convert_to_reference_fx(prepared, backend_config=backend_config)
converted(*example_inputs)
# Verify that the reference model is correct
#
# Reference model until add should be:
# fp32_input -> q_to_int32 -> [dq -> linear1_fp32 -> q_to_int32] -> dq ->
# q_to_uint8 -> [dq -> linear2_fp32 -> q_to_uint8] -> dq (linear2_dq) ->
# q_to_int32 -> [dq -> sigmoid_fp32 -> q_to_int32] -> dq ->
# q_to_uint8 -> [dq -> tanh_fp32 -> q_to_uint8] -> dq (tanh_dq)
#
# Complete reference model with add should be:
# [(linear2_dq, tanh_dq) -> add_fp32 -> q_to_uint8] -> dq -> fp32_output
target_to_expected_dtypes = {
"linear1": torch.qint32,
"linear2": torch.quint8,
"sigmoid": torch.qint32,
"tanh": torch.quint8,
torch.add: torch.quint8,
}
# Find the patterns [dq - op_fp32 - q_to_specific_dtype] in the graph
linear2_node = tanh_node = None
for node in converted.graph.nodes:
if node.target not in target_to_expected_dtypes:
continue
# Match preceding dequantize
self.assertTrue(len(node.args) == 1 or len(node.args) == 2)
self.assertTrue(all(arg.target == "dequantize" for arg in node.args))
# Match following quantize with the specific dtypes
self.assertEqual(len(node.users), 1)
user = list(node.users.keys())[0]
self.assertEqual(user.target, torch.quantize_per_tensor)
self.assertEqual(user.args[-1], target_to_expected_dtypes[node.target])
# Match [dq - torch.add(linear2_dq, tanh_dq) - q]
if node.target == "linear2":
linear2_node = node
elif node.target == "tanh":
tanh_node = node
elif node.target == torch.add:
linear2_dq, tanh_dq = node.args
self.assertEqual(tanh_dq.args[0].args[0], tanh_node)
self.assertEqual(linear2_dq.args[0].args[0], linear2_node)
def test_lowering_functional_conv_with_kwargs(self):
dim_to_op = {
1: F.conv1d,
2: F.conv2d,
3: F.conv3d,
}
dim_to_qop = {
1: torch.ops.quantized.conv1d,
2: torch.ops.quantized.conv2d,
3: torch.ops.quantized.conv3d,
}
class Mod(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, dimension):
super().__init__()
self.dim = dimension
self.op = dim_to_op[dimension]
kernel_sizes = [kernel_size] * self.dim
self.weight = nn.Parameter(torch.randn(out_channels, in_channels, *kernel_sizes))
def forward(self, input):
return self.op(input, self.weight, bias=None, stride=[1] * self.dim,
padding=[0] * self.dim, dilation=[1] * self.dim, groups=1)
for dimension in [1, 2, 3]:
model = Mod(3, 16, 3, dimension)
model.eval()
qconfig_mapping = get_default_qconfig_mapping()
input_shape = (1, 3, *([8] * dimension))
example_inputs = torch.randn(input_shape)
prepared_model = prepare_fx(model, qconfig_mapping, example_inputs)
prepared_model(example_inputs)
quantized_model = convert_fx(prepared_model)
# This should pass
quantized_model(example_inputs)
# Ensure the quantized model has the expected op
node_occurrence = {
ns.call_function(dim_to_qop[dimension]): 1,
}
self.checkGraphModuleNodes(quantized_model, expected_node_occurrence=node_occurrence)
def test_lowering_functional_conv_transpose_with_kwargs(self):
dim_to_op = {
1: F.conv_transpose1d,
2: F.conv_transpose2d,
3: F.conv_transpose3d,
}
dim_to_qop = {
1: torch.ops.quantized.conv_transpose1d,
2: torch.ops.quantized.conv_transpose2d,
3: torch.ops.quantized.conv_transpose3d,
}
class Mod(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, dimension):
super().__init__()
self.dim = dimension
self.op = dim_to_op[dimension]
kernel_sizes = [kernel_size] * self.dim
self.weight = nn.Parameter(torch.randn(in_channels, out_channels, *kernel_sizes))
def forward(self, input):
return self.op(input, self.weight, bias=None, stride=[1] * self.dim,
padding=[0] * self.dim, output_padding=[0] * self.dim,
dilation=[1] * self.dim, groups=1)
for dimension in [1, 2, 3]:
model = Mod(3, 16, 3, dimension)
model.eval()
qconfig_mapping = get_default_qconfig_mapping()
input_shape = (1, 3, *([8] * dimension))
example_inputs = torch.randn(input_shape)
prepared_model = prepare_fx(model, qconfig_mapping, example_inputs)
prepared_model(example_inputs)
quantized_model = convert_fx(prepared_model)
# This should pass
quantized_model(example_inputs)
# Ensure the quantized model has the expected op
node_occurrence = {
ns.call_function(dim_to_qop[dimension]): 1,
}
self.checkGraphModuleNodes(quantized_model, expected_node_occurrence=node_occurrence)
def test_lowering_functional_linear_with_kwargs(self):
class Mod(nn.Module):
def __init__(self, in_channels, out_channels):
super().__init__()
self.weight = nn.Parameter(torch.randn(out_channels, in_channels))
def forward(self, input):
return F.linear(input, self.weight, bias=None)
model = Mod(8, 4)
model.eval()
qconfig_mapping = get_default_qconfig_mapping()
example_inputs = torch.randn(1, 8)
prepared_model = prepare_fx(model, qconfig_mapping, example_inputs)
prepared_model(example_inputs)
quantized_model = convert_fx(prepared_model)
# This should pass
quantized_model(example_inputs)
# Ensure the quantized model has the expected op
node_occurrence = {
ns.call_function(torch.ops.quantized.linear): 1,
}
self.checkGraphModuleNodes(quantized_model, expected_node_occurrence=node_occurrence)
@skipIfNoFBGEMM
class TestQuantizeFxOps(QuantizationTestCase):
def setUp(self):
super().setUp()
self.custom_qconfig = torch.ao.quantization.QConfig(
activation=torch.ao.quantization.observer.HistogramObserver.with_args(
qscheme=torch.per_tensor_symmetric, dtype=torch.qint8
),
weight=torch.ao.quantization.default_per_channel_weight_observer
)
self.common_quant_patterns = {
torch.nn.ConvTranspose1d: DefaultNodeQuantizeHandler,
torch.nn.ConvTranspose2d: DefaultNodeQuantizeHandler,
torch.nn.ELU: DefaultNodeQuantizeHandler,
torch.nn.LeakyReLU: DefaultNodeQuantizeHandler,
torch.nn.Hardswish: DefaultNodeQuantizeHandler,
torch.nn.InstanceNorm1d: DefaultNodeQuantizeHandler,
torch.nn.InstanceNorm2d: DefaultNodeQuantizeHandler,
torch.nn.InstanceNorm3d: DefaultNodeQuantizeHandler,
torch.nn.LayerNorm: DefaultNodeQuantizeHandler,
torch.nn.SiLU: DefaultNodeQuantizeHandler,
torch.nn.Mish: DefaultNodeQuantizeHandler,
torch.nn.GELU: DefaultNodeQuantizeHandler,
torch.nn.Softmax: DefaultNodeQuantizeHandler,
torch.nn.functional.elu: DefaultNodeQuantizeHandler,
torch.nn.functional.hardswish: DefaultNodeQuantizeHandler,
torch.nn.functional.instance_norm: DefaultNodeQuantizeHandler,
torch.nn.functional.layer_norm: DefaultNodeQuantizeHandler,
torch.nn.functional.leaky_relu: DefaultNodeQuantizeHandler,
torch.nn.functional.silu: DefaultNodeQuantizeHandler,
torch.nn.functional.mish: DefaultNodeQuantizeHandler,
torch.nn.functional.gelu: DefaultNodeQuantizeHandler,
torch.nn.functional.softmax: DefaultNodeQuantizeHandler,
torch.sum: DefaultNodeQuantizeHandler
}
"""Unit tests for individual ops
"""
@skipIfNoFBGEMM
def test_linear_module(self):
with override_quantized_engine('fbgemm'):
class LinearModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(30, 4).float()
def forward(self, x):
return self.linear(x)
class LinearReLUModel(torch.nn.Module):
def __init__(self, f_relu=False):
super().__init__()
self.linear = torch.nn.Linear(30, 4).float()
if f_relu:
self.relu = F.relu
else:
self.relu = torch.nn.ReLU()
def forward(self, x):
x = self.linear(x)
x = self.relu(x)
return x
class LinearBnModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(4, 4).float()
self.bn = torch.nn.BatchNorm1d(4)
def forward(self, x):
x = self.linear(x)
x = self.bn(x)
return x
# Test linear
data = (torch.rand((1, 30), dtype=torch.float),)
for quant_type in self.all_quant_types:
model = LinearModel()
quantized_module = nnqd.Linear if quant_type == QuantType.DYNAMIC else nnq.Linear
quantized_node = ns.call_module(quantized_module)
result_dict = self.checkGraphModeFxOp(model, data, quant_type, quantized_node)
if quant_type in self.static_quant_types:
self.assertEqual(result_dict["quantized_output"], result_dict["quantized_reference_output"])
# TODO: enable test for dynamic quant
# Test linear-relu
for f_relu, quant_type in itertools.product([True, False], [QuantType.STATIC, QuantType.QAT]):
model = LinearReLUModel(f_relu)
quantized_node = ns.call_module(nniq.LinearReLU)
result_dict = self.checkGraphModeFxOp(model, data, quant_type, quantized_node)
self.assertEqual(result_dict["quantized_output"], result_dict["quantized_reference_output"])
# Test linear-bn
data = (torch.rand((4, 4), dtype=torch.float),)
for quant_type in self.static_quant_types:
model = LinearBnModel()
quantized_node = ns.call_module(nnq.Linear)
result_dict = self.checkGraphModeFxOp(model, data, quant_type, quantized_node)
self.assertEqual(result_dict["quantized_output"], result_dict["quantized_reference_output"])
@skipIfNoFBGEMM
def test_functional_linear(self):
with override_quantized_engine('fbgemm'):
class FuncLinear(torch.nn.Module):
def __init__(self, use_bias, has_relu, f_relu):
super().__init__()
self.w = torch.randn(4, 30)
self.b = torch.randn(4)
self.use_bias = use_bias
if has_relu:
if f_relu:
self.relu_or_id = F.relu
else:
self.relu_or_id = torch.nn.ReLU()
else:
self.relu_or_id = torch.nn.Identity()
def forward(self, x):
if self.use_bias:
x = F.linear(x, self.w, self.b)
else:
x = F.linear(x, self.w)
x = self.relu_or_id(x)
return x
data = (torch.rand((1, 30), dtype=torch.float),)
quant_type_to_qlinear_fun = {
QuantType.DYNAMIC: ns.call_function(torch.ops.quantized.linear_dynamic),
QuantType.STATIC: ns.call_function(torch.ops.quantized.linear),
QuantType.QAT: ns.call_function(torch.ops.quantized.linear),
}
quant_type_to_qlinear_relu_fun = {
# we don't have linear_relu_dynamic
QuantType.DYNAMIC: ns.call_function(torch.ops.quantized.linear_relu_dynamic),
QuantType.STATIC: ns.call_function(torch.ops.quantized.linear_relu),
QuantType.QAT: ns.call_function(torch.ops.quantized.linear_relu),
}
options = itertools.product(
self.all_quant_types,
(True, False), # use_bias
(True, False), # has_relu
(True, False), # functional relu
)
for quant_type, use_bias, has_relu, f_relu in options:
# when has_relu is False, we are using an nn.Identity and
# we will insert observer/fake_quant for the output of nn.Identity since
# it is a copy node, that's why we have extra observer/fake_quant
# when has_relu is False
quant_type_to_prepare_expected_node_occurrence = {
QuantType.DYNAMIC: {
ns.call_module(torch.ao.quantization.PlaceholderObserver): 1,
ns.call_module(torch.ao.quantization.MinMaxObserver): 1,
},
# There should be 3 observers: after input, weight and activation.
# one more observer for torch.nn.Identity when there is no relu
QuantType.STATIC: {
ns.call_module(torch.ao.quantization.HistogramObserver): 2 if has_relu else 3,
ns.call_module(torch.ao.quantization.PerChannelMinMaxObserver): 1,
},
# There should be 3 observers: after input, weight and activation.
QuantType.QAT: {
ns.call_module(torch.ao.quantization.FusedMovingAvgObsFakeQuantize): 3 if has_relu else 4,
},
}
model = FuncLinear(use_bias, has_relu, f_relu)
if has_relu:
qlinear_fun = quant_type_to_qlinear_relu_fun[quant_type]
else:
qlinear_fun = quant_type_to_qlinear_fun[quant_type]
if quant_type != QuantType.DYNAMIC:
num_dequantize = 1
else:
# we will have an extra quantize_per_tensor_dynamic + dequantize for
# nn.Identity right now, but it will be fixed after we use
# backend_config to configure the default pt backend
num_dequantize = int(not has_relu)
convert_node_occurrence = {
ns.call_function(torch.quantize_per_tensor): 1 if quant_type != QuantType.DYNAMIC else 0,
qlinear_fun: 1,
ns.call_method("dequantize"): num_dequantize if quant_type != QuantType.DYNAMIC else 0,
}
prepare_expected_node_occurrence = \
quant_type_to_prepare_expected_node_occurrence[quant_type]
result_dict = self.checkGraphModeFxOp(
model, data, quant_type, qlinear_fun,
prepare_expected_node_occurrence=prepare_expected_node_occurrence,
expected_node_occurrence=convert_node_occurrence)
if quant_type != QuantType.DYNAMIC:
self.assertEqual(result_dict["quantized_output"], result_dict["quantized_reference_output"])
# Ensure packed weights in lowered models are folded
self.assertIn("_packed_weight_0", result_dict["quantized"].state_dict().keys())
@skipIfNoFBGEMM
def test_linear_dynamic_fp16(self):
with override_quantized_engine('fbgemm'):
class FuncLinear(torch.nn.Module):
def __init__(self, use_bias, has_relu, f_relu):
super().__init__()
self.w = torch.randn(4, 30)
self.b = torch.randn(4)
self.use_bias = use_bias
if has_relu:
if f_relu:
self.relu = F.relu
else:
self.relu = torch.nn.ReLU()
else:
self.relu = torch.nn.Identity()
def forward(self, x):
if self.use_bias:
x = F.linear(x, self.w, self.b)
else:
x = F.linear(x, self.w)
x = self.relu(x)
return x
data = (torch.rand((1, 30), dtype=torch.float),)
options = itertools.product(
(True, False), # use_bias
(True, False), # has_relu
(True, False), # functional relu
(True, False), # is_reference
)
for use_bias, has_relu, f_relu, is_reference in options:
model = FuncLinear(use_bias, has_relu, f_relu)
if is_reference:
qlinear_fun = ns.call_function(torch.nn.functional.linear)
else:
if has_relu:
qlinear_fun = ns.call_function(torch.ops.quantized.linear_relu_dynamic_fp16)
else:
qlinear_fun = ns.call_function(torch.ops.quantized.linear_dynamic_fp16)
prepare_node_occurrence = {
# activation and weight
ns.call_module(torch.ao.quantization.PlaceholderObserver): 2
}
convert_node_occurrence = {
qlinear_fun: 1,
# weight
ns.call_method("to"): 1 if is_reference else 0
}
self.checkGraphModeFxOp(
model, data, QuantType.DYNAMIC, qlinear_fun,
is_reference=is_reference,
custom_qconfig_dict={"": float16_dynamic_qconfig},
prepare_expected_node_occurrence=prepare_node_occurrence,
expected_node_occurrence=convert_node_occurrence)
def test_linear_static_fp16(self):
class FuncLinear(torch.nn.Module):
def __init__(self, use_bias, has_relu, f_relu):
super().__init__()
self.w = torch.randn(4, 30)
self.b = torch.randn(4)
self.use_bias = use_bias
if has_relu:
if f_relu:
self.relu = F.relu
else:
self.relu = torch.nn.ReLU()
else:
self.relu = torch.nn.Identity()
def forward(self, x):
if self.use_bias:
x = F.linear(x, self.w, self.b)
else:
x = F.linear(x, self.w)
x = self.relu(x)
return x
data = (torch.rand((1, 30), dtype=torch.float),)
options = itertools.product(
(True, False), # use_bias
(True, False), # has_relu
(True, False), # functional relu
(True, False), # is_reference
)
backend_config = get_test_only_legacy_native_backend_config()
for use_bias, has_relu, f_relu, is_reference in options:
model = FuncLinear(use_bias, has_relu, f_relu)
linear_fun = ns.call_function(torch.nn.functional.linear)
# when has_relu is False, we are using an nn.Identity and
# we will insert observer/fake_quant for the output of nn.Identity since
# it is a copy node, that's why we have extra observer/fake_quant
# when has_relu is False
prepare_node_occurrence = {
# activation, weight, bias and output
ns.call_module(torch.ao.quantization.PlaceholderObserver): 3 + int(use_bias) + int(not has_relu),
}
# We have extra to and dequantize when is_reference is True
# and has_relu is False since when has_relu is False, we
# have an nn.Identity in the model, which is a CopyNode
# and we would add extra quant - dequant for CopyNode in
# reference patterns
convert_node_occurrence = {
# we don't support static fp16 ops, so the linear function
# is unfused
linear_fun: 1,
# activation, weight, bias and output
ns.call_method("to"): 3 + int(use_bias) + int(not has_relu and is_reference),
ns.call_method("dequantize"): 3 + int(use_bias) + int(not has_relu and is_reference)
}
self.checkGraphModeFxOp(
model, data, QuantType.DYNAMIC, linear_fun,
is_reference=is_reference,
custom_qconfig_dict={"": float16_static_qconfig},
prepare_expected_node_occurrence=prepare_node_occurrence,
expected_node_occurrence=convert_node_occurrence,
backend_config=backend_config)
@skipIfNoFBGEMM
def test_conv_module(self):
conv_module = {1 : torch.nn.Conv1d, 2 : torch.nn.Conv2d, 3 : torch.nn.Conv3d}
class ConvWrapper(torch.nn.Module):
def __init__(self, dim):
super().__init__()
self.conv = conv_module[dim](3, 3, 3).float()
def forward(self, x):
return self.conv(x)
options = itertools.product([1, 2, 3], self.static_quant_types)
quantized_nodes = {
# dim
1: ns.call_module(nnq.Conv1d),
2: ns.call_module(nnq.Conv2d),
3: ns.call_module(nnq.Conv3d),
}
for dim, quant_type in options:
self.checkGraphModeFxOp(
ConvWrapper(dim), self.img_data_dict[dim], quant_type,
quantized_nodes[dim])
@skipIfNoFBGEMM
def test_functional_conv(self):
with override_quantized_engine('fbgemm'):
""" Test for function conv and functional conv + relu
"""
convs = {
1: torch.nn.functional.conv1d,
2: torch.nn.functional.conv2d,
3: torch.nn.functional.conv3d,
}
class FuncConv(torch.nn.Module):
def __init__(self, dim, use_bias, has_relu, f_relu):
super().__init__()
self.dim = dim
self.w = torch.randn(tuple([3] * (dim + 2)))
self.b = torch.randn(3) if use_bias else None
self.stride = tuple([1] * dim)
self.padding = tuple([0] * dim)
self.dilation = tuple([1] * dim)
self.groups = 1
self.use_bias = use_bias
if has_relu:
if f_relu:
self.relu = F.relu
else:
self.relu = torch.nn.ReLU()
else:
self.relu = torch.nn.Identity()
def forward(self, x):
x = convs[self.dim](x, self.w, self.b, self.stride, self.padding, self.dilation, self.groups)
x = self.relu(x)
return x
quant_type_to_qconv_fun = {
QuantType.STATIC: {
1: ns.call_function(torch.ops.quantized.conv1d),
2: ns.call_function(torch.ops.quantized.conv2d),
3: ns.call_function(torch.ops.quantized.conv3d)
},
QuantType.QAT: {
1: ns.call_function(torch.ops.quantized.conv1d),
2: ns.call_function(torch.ops.quantized.conv2d),
3: ns.call_function(torch.ops.quantized.conv3d)
},
}
quant_type_to_qconv_relu_fun = {
QuantType.STATIC: {
1: ns.call_function(torch.ops.quantized.conv1d_relu),
2: ns.call_function(torch.ops.quantized.conv2d_relu),
3: ns.call_function(torch.ops.quantized.conv3d_relu)
},
QuantType.QAT: {
1: ns.call_function(torch.ops.quantized.conv1d_relu),
2: ns.call_function(torch.ops.quantized.conv2d_relu),
3: ns.call_function(torch.ops.quantized.conv3d_relu)
},
}
options = itertools.product(
[1, 2, 3], # dims
self.static_quant_types,
(True, False), # use_bias
(True, False), # has_relu
(True, False), # functional relu
)
for dim, quant_type, use_bias, has_relu, f_relu in options:
# when has_relu is False, we are using an nn.Identity and
# we will insert observer/fake_quant for the output of nn.Identity since
# it is a copy node, that's why we have extra observer/fake_quant
# when has_relu is False
quant_type_to_prepare_expected_node_occurrence = {
QuantType.DYNAMIC: {},
# There should be 3 observers: after input, weight and activation.
QuantType.STATIC: {
ns.call_module(torch.ao.quantization.HistogramObserver): 2 if has_relu else 3,
ns.call_module(torch.ao.quantization.PerChannelMinMaxObserver): 1,
},
# There should be 3 observers: after input, weight and activation.
QuantType.QAT: {
ns.call_module(torch.ao.quantization.FusedMovingAvgObsFakeQuantize): 3 if has_relu else 4,
},
}
data_dims = [2, 3] + [4] * dim
data = (torch.randn(tuple(data_dims), dtype=torch.float),)
model = FuncConv(dim, use_bias, has_relu, f_relu)
if has_relu:
qconv_fun = quant_type_to_qconv_relu_fun[quant_type][dim]
else:
qconv_fun = quant_type_to_qconv_fun[quant_type][dim]
convert_node_occurrence = {
ns.call_function(torch.quantize_per_tensor): 1,
qconv_fun: 1,
ns.call_method("dequantize"): 1
}
prepare_expected_node_occurrence = \
quant_type_to_prepare_expected_node_occurrence[quant_type]
result_dict = self.checkGraphModeFxOp(
model, data, quant_type, qconv_fun,
prepare_expected_node_occurrence=prepare_expected_node_occurrence,
expected_node_occurrence=convert_node_occurrence)
if quant_type != QuantType.DYNAMIC:
self.assertEqual(result_dict["quantized_output"], result_dict["quantized_reference_output"])
# Ensure packed weights in lowered models are folded
self.assertIn("_packed_weight_0", result_dict["quantized"].state_dict().keys())
@skipIfNoFBGEMM
def test_quantized_conv_relu(self):
"""tests for conv1d_relu/conv2d_relu/conv3d_relu"""
conv_module = {1 : torch.nn.Conv1d, 2 : torch.nn.Conv2d, 3 : torch.nn.Conv3d}
class ConvNdRelu(torch.nn.Module):
def __init__(self, dim, inplace):
super().__init__()
self.conv = conv_module[dim](3, 3, 3).float()
self.relu = torch.nn.ReLU(inplace)
def forward(self, x):
return self.relu(self.conv(x))
class ConvNdFunctionalRelu(torch.nn.Module):
def __init__(self, dim):
super().__init__()
self.conv = conv_module[dim](3, 3, 3).float()
def forward(self, x):
return F.relu(self.conv(x))
class ConvNdInplaceFunctionalRelu(torch.nn.Module):
def __init__(self, dim):
super().__init__()
self.conv = conv_module[dim](3, 3, 3).float()
def forward(self, x):
return F.relu(self.conv(x), True)
options = itertools.product([1, 2, 3], self.static_quant_types)
quantized_nodes = {
# dim
1: ns.call_module(nniq.ConvReLU1d),
2: ns.call_module(nniq.ConvReLU2d),
3: ns.call_module(nniq.ConvReLU3d),
}
for dim, quant_type in options:
for m in [ConvNdRelu(dim, True),
ConvNdRelu(dim, False),
ConvNdFunctionalRelu(dim),
ConvNdInplaceFunctionalRelu(dim)]:
self.checkGraphModeFxOp(
m, self.img_data_dict[dim], quant_type,
quantized_nodes[dim])
def _test_binary_op_int8_impl(self, binary_op, ibinary_op, quantized_op):
data = (torch.randn(1, 1, 1, 1, dtype=torch.float),
torch.randn(1, 1, 1, 1, dtype=torch.float))
options = itertools.product([True, False], [True, False], [True, False])
quant_type = QuantType.STATIC
# testing for default int8 static quant
for is_inplace, is_scalar, is_reference in options:
if is_reference:
node_list = [
ns.call_method("dequantize"),
ns.call_function(binary_op),
ns.call_function(torch.quantize_per_tensor)
]
quantized_node = None
else:
node_list = None
quantized_node = ns.call_function(quantized_op)
self.checkGraphModeFxOp(
BinaryOp(binary_op, ibinary_op, is_inplace, is_scalar), data, quant_type,
quantized_node, expected_node_list=node_list, is_reference=is_reference)
# This tests the binary op should be quantized even when it is not feed with a
# quantized input
self.checkGraphModeFxOp(
BinaryOpNonQuantizedInput(binary_op, ibinary_op, is_inplace, is_scalar),
data, quant_type, quantized_node,
expected_node_list=node_list, is_reference=is_reference)
def _test_binary_op_float16_impl(self, binary_op, ibinary_op):
data = (torch.randn(1, 1, 1, 1, dtype=torch.float),
torch.randn(1, 1, 1, 1, dtype=torch.float))
quant_type = QuantType.STATIC
# testing for fp16 static quant
# we are producing fp16 patterns
options = itertools.product([True, False], [True, False])
custom_qconfig_dict = {
"object_type": [(binary_op, float16_static_qconfig)]
}
backend_config = get_test_only_legacy_native_backend_config()
for is_inplace, is_scalar in options:
node_occurrence = {
# output_conv1, output_add1, output_add2 for scalar
# output_conv1, output_conv2, output_add1, output_add2 for non-scalar
ns.call_method("to"): 3 if is_scalar else 4
}
self.checkGraphModeFxOp(
BinaryOp(binary_op, ibinary_op, is_inplace, is_scalar), data, quant_type,
expected_node_occurrence=node_occurrence,
custom_qconfig_dict=custom_qconfig_dict,
backend_config=backend_config)
node_occurrence = {
# input_add, output_add for scalar
# input_add1, input_add2, output_add for non-scalar
ns.call_method("to"): 2 if is_scalar else 3
}
self.checkGraphModeFxOp(
BinaryOpNonQuantizedInput(binary_op, ibinary_op, is_inplace, is_scalar), data, quant_type,
expected_node_occurrence=node_occurrence,
custom_qconfig_dict=custom_qconfig_dict,
backend_config=backend_config)
def _test_binary_op_relu_int8_impl(self, binary_op, ibinary_op, quantized_op):
data = (torch.rand((1, 1, 1, 1), dtype=torch.float),
torch.rand((1, 1, 1, 1), dtype=torch.float))
quant_type = QuantType.STATIC
quantized_node = ns.call_function(quantized_op)
options = itertools.product(
[True, False], [nn.ReLU, F.relu, torch.relu], [True, False])
for is_inplace_op, relu_callable, is_scalar in options:
model = BinaryOpRelu(
binary_op, ibinary_op, is_inplace_op, relu_callable, is_scalar)
self.checkGraphModeFxOp(
model, data, quant_type, quantized_node)
def _test_binary_op_relu_float16_impl(self, binary_op, ibinary_op):
data = (torch.rand((1, 1, 1, 1), dtype=torch.float),
torch.rand((1, 1, 1, 1), dtype=torch.float))
quant_type = QuantType.STATIC
options = itertools.product(
[True, False], [nn.ReLU, F.relu, torch.relu], [True, False])
custom_qconfig_dict = {
"": float16_static_qconfig,
"object_type": [(torch.nn.Conv2d, None)]
}
backend_config = get_test_only_legacy_native_backend_config()
for is_inplace_op, is_functional_relu, is_scalar in options:
node_occurrence = {
ns.call_method("to"): 3 if is_scalar else 4
}
model = BinaryOpRelu(
binary_op, ibinary_op, is_inplace_op, is_functional_relu, is_scalar)
self.checkGraphModeFxOp(
model, data, quant_type, custom_qconfig_dict=custom_qconfig_dict,
expected_node_occurrence=node_occurrence,
backend_config=backend_config)
@skipIfNoFBGEMM
def test_add(self):
self._test_binary_op_int8_impl(
operator.add, operator.iadd, torch.ops.quantized.add)
self._test_binary_op_float16_impl(
operator.add, operator.iadd)
@unittest.skip("This is no longer needed right now, can enable later with new api")
def test_sub(self):
self._test_binary_op_float16_impl(operator.sub, operator.isub)
self._test_binary_op_float16_impl(torch.sub, None)
@unittest.skip("This is no longer needed right now, can enable later with new api")
def test_div(self):
self._test_binary_op_float16_impl(operator.truediv, operator.itruediv)
self._test_binary_op_float16_impl(torch.div, None)
@skipIfNoFBGEMM
def test_mul(self):
self._test_binary_op_int8_impl(
operator.mul, operator.imul, torch.ops.quantized.mul)
self._test_binary_op_float16_impl(operator.mul, operator.imul)
@unittest.skip("This is no longer needed right now, can enable later with new api")
def test_sum(self):
class Sum(torch.nn.Module):
def forward(self, x):
x = torch.sum(x, [1], keepdim=True)
x = torch.sum(x, [1])
return x
data = torch.randn(1, 2, 3, 4, dtype=torch.float)
quant_type = QuantType.STATIC
# testing for fp16 static quant
# we are producing fp16 patterns
custom_qconfig_dict = {
"object_type": [(torch.sum, float16_static_qconfig)]
}
node_occurrence = {
# input_sum1, output_sum1, output_sum2
ns.call_method("to"): 3
}
self.checkGraphModeFxOp(
Sum(), data, quant_type,
expected_node_occurrence=node_occurrence,
custom_qconfig_dict=custom_qconfig_dict)
@unittest.skip("This is no longer needed right now, can enable later with new api")
def test_bmm(self):
class BMMMethod(torch.nn.Module):
def forward(self, x, y):
return x.bmm(y)
data = (torch.randn(1, 1, 1, dtype=torch.float),
torch.randn(1, 1, 1, dtype=torch.float))
quant_type = QuantType.STATIC
# testing for fp16 static quant
# we are producing fp16 patterns
custom_qconfig_dict = {
"object_type": [(torch.bmm, float16_static_qconfig),
("bmm", float16_static_qconfig)]
}
node_occurrence = {
# input_bmm1, input_bmm2, output_bmm
ns.call_method("to"): 3
}
self.checkGraphModeFxOp(
BinaryOpNonQuantizedInput(torch.bmm, None, False, False), data, quant_type,
expected_node_occurrence=node_occurrence,
custom_qconfig_dict=custom_qconfig_dict)
# TODO: support call_method("bmm")
# we can transform call_method("bmm") to call_function(torch.bmm)
# self.checkGraphModeFxOp(
# BMMMethod(), data, quant_type,
# expected_node_occurrence=node_occurrence,
# custom_qconfig_dict=custom_qconfig_dict,
# print_debug_info=True)
@skipIfNoFBGEMM
def test_add_relu(self):
self._test_binary_op_relu_int8_impl(
operator.add, operator.iadd, torch.ops.quantized.add_relu)
self._test_binary_op_relu_float16_impl(
operator.add, operator.iadd)
@skipIfNoFBGEMM
def test_add_relu_multiple_uses_of_relu(self):
class Sub(torch.nn.Module):
def __init__(self):
super().__init__()
self.relu = torch.nn.ReLU(inplace=True)
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.sub = Sub()
def forward(self, x, y):
x = x + y
x = self.sub.relu(x)
x = x + y
x = self.sub.relu(x)
return x
m = M().eval()
example_inputs = (torch.randn(3), torch.randn(3))
m = prepare_fx(m, {"": default_qconfig}, example_inputs=example_inputs)
m = convert_fx(m)
node_occurrence = {
ns.call_function(torch.quantize_per_tensor): 2,
ns.call_function(torch.ops.quantized.add_relu): 2,
ns.call_method("dequantize"): 1,
}
self.checkGraphModuleNodes(m, expected_node_occurrence=node_occurrence)
# check the model is scriptable
m = torch.jit.script(m)
# check the model is runnable
m(*example_inputs)
@skipIfNoFBGEMM
def test_mul_relu(self):
self._test_binary_op_relu_int8_impl(
operator.mul, operator.imul, torch.ops.quantized.mul_relu)
self._test_binary_op_relu_float16_impl(
operator.mul, operator.imul)
# TODO(future PR): make more generic
def _test_quantized_add_mul_qat(self, model, example_inputs, expected_node_occurrence):
qconfig_dict = {'': torch.ao.quantization.get_default_qat_qconfig('fbgemm')}
mp = prepare_qat_fx(model, qconfig_dict, example_inputs=example_inputs)
self.checkGraphModuleNodes(
mp, expected_node_occurrence=expected_node_occurrence)
@skipIfNoFBGEMM
def test_quantized_add_qat(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1 = torch.nn.Conv2d(1, 1, 1)
self.conv2 = torch.nn.Conv2d(1, 1, 1)
def forward(self, x):
x = torch.add(x, 1.0)
x = self.conv1(x)
x = torch.add(x, 1.0)
x = torch.relu(x)
x = self.conv2(x)
return x
m = M()
example_inputs = (torch.randn(1, 1, 1, 1),)
expected_node_occurrence = {
ns.call_module(torch.ao.quantization.FusedMovingAvgObsFakeQuantize): 5,
}
self._test_quantized_add_mul_qat(m, example_inputs, expected_node_occurrence)
@skipIfNoFBGEMM
def test_quantized_mul_qat(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1 = torch.nn.Conv2d(1, 1, 1)
self.conv2 = torch.nn.Conv2d(1, 1, 1)
def forward(self, x):
x = torch.mul(x, 1.0)
x = self.conv1(x)
x = torch.mul(x, 1.0)
x = torch.relu(x)
x = self.conv2(x)
return x
m = M()
example_inputs = (torch.randn(1, 1, 1, 1),)
expected_node_occurrence = {
ns.call_module(torch.ao.quantization.FusedMovingAvgObsFakeQuantize): 5,
}
self._test_quantized_add_mul_qat(m, example_inputs, expected_node_occurrence)
def test_int8_input_no_unnecessary_fq(self):
"""
If the inputs to the graph are quantized and the only node
does not need an activation observer, verifies that the
activation observer is not inserted.
"""
class M(nn.Module):
def __init__(self, scalar):
super().__init__()
self.scalar = scalar
self.add_func = torch.ao.nn.quantized.FloatFunctional()
def forward(self, x):
return self.add_func.add_scalar(x, self.scalar)
m = M(0.5)
mp = torch.ao.quantization.quantize_fx.prepare_qat_fx(
m, {'': torch.ao.quantization.get_default_qat_qconfig('fbgemm')},
example_inputs=(torch.randn(1),),
prepare_custom_config={"input_quantized_idxs": [0]})
expected_node_occurrence = {
ns.call_module(torch.ao.quantization.FusedMovingAvgObsFakeQuantize): 1,
}
self.checkGraphModuleNodes(
mp, expected_node_occurrence=expected_node_occurrence)
@skipIfNoFBGEMM
def test_cat(self):
""" quantization of the output of cat will depend on the
input of cat. we only quantize the output of cat when its inputs are quantized.
"""
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1 = torch.nn.Conv2d(2, 2, 2).float()
self.conv2 = torch.nn.Conv2d(2, 2, 2).float()
def forward(self, x, y):
x = self.conv1(x)
y = self.conv2(y)
return torch.cat([x, y], 1)
example_inputs = (torch.randn(1, 2, 5, 5, dtype=torch.float),
torch.randn(1, 2, 5, 5, dtype=torch.float))
quantized_node = ns.call_function(torch.cat)
options = itertools.product(self.static_quant_types, [True, False])
for quant_type, is_reference in options:
if is_reference:
converted_node_list = [
ns.call_method("dequantize"),
ns.call_function(torch.cat),
ns.call_function(torch.quantize_per_tensor)
]
converted_node_occurrence = {
# inputs and outputs of the two conv, and output of cat
ns.call_method("dequantize"): 5,
ns.call_function(torch.cat): 1,
# inputs and outputs of the two conv, and output of cat
ns.call_function(torch.quantize_per_tensor): 5,
}
else:
converted_node_list = None
converted_node_occurrence = {
# output of cat
ns.call_method("dequantize"): 1,
ns.call_function(torch.cat): 1,
# for two inputs
ns.call_function(torch.quantize_per_tensor): 2,
}
self.checkGraphModeFxOp(
M(),
example_inputs,
quant_type,
quantized_node,
expected_node_list=converted_node_list,
expected_node_occurrence=converted_node_occurrence,
is_reference=is_reference)
# check cat is using the same observer for input and output
m = M().eval()
m = prepare_fx(m, {"": default_qconfig}, example_inputs=example_inputs)
# two inputs and one output of torch.cat are using same observer, so we have
# 2 observers that's replicated
all_observers = len(dict(m.named_modules(remove_duplicate=False)))
distinct_observers = len(dict(m.named_modules()))
self.assertEqual(all_observers, distinct_observers + 2)
# make sure the converted model runs
m = convert_fx(m)
m(*example_inputs)
@skipIfNoFBGEMM
def test_qbatch_norm(self):
bn_module = {
# TODO: quantized batchnorm 1d module is missing
# 1 : torch.nn.BatchNorm1d,
2 : torch.nn.BatchNorm2d,
3 : torch.nn.BatchNorm3d,
}
class M(torch.nn.Module):
def __init__(self, dim):
super().__init__()
self.bn = bn_module[dim](3).to(torch.float)
def forward(self, x):
return self.bn(x)
options = itertools.product(self.static_quant_types, [2, 3], [True, False])
quantized_nodes = {
False: {
# 1: ns.call_module(nnq.BatchNorm1d),
2: ns.call_module(nnq.BatchNorm2d),
3: ns.call_module(nnq.BatchNorm3d),
},
True: {
# 1: ns.call_module(nn.BatchNorm1d),
2: ns.call_module(nn.BatchNorm2d),
3: ns.call_module(nn.BatchNorm3d),
}
}
for quant_type, dim, is_reference in options:
self.checkGraphModeFxOp(
M(dim), self.img_data_dict[dim], quant_type, quantized_nodes[is_reference][dim], is_reference=is_reference)
@skipIfNoFBGEMM
def test_qbatch_norm_relu(self):
bn_module = {2 : torch.nn.BatchNorm2d, 3 : torch.nn.BatchNorm3d}
class BNRelu(torch.nn.Module):
def __init__(self, dim, inplace):
super().__init__()
self.bn = bn_module[dim](3).to(torch.float)
self.relu = torch.nn.ReLU(inplace=inplace)
def forward(self, x):
return self.relu(self.bn(x))
class BNFuncRelu(torch.nn.Module):
def __init__(self, dim):
super().__init__()
self.bn = bn_module[dim](3).to(torch.float)
def forward(self, x):
return F.relu(self.bn(x), False)
class BNFuncInplaceRelu(torch.nn.Module):
def __init__(self, dim):
super().__init__()
self.bn = bn_module[dim](3).to(torch.float)
def forward(self, x):
return F.relu(self.bn(x), True)
options = itertools.product(self.static_quant_types, [2, 3], [True, False])
quantized_nodes = {
True: {
2: ns.call_module(nni.BNReLU2d),
3: ns.call_module(nni.BNReLU3d),
},
False: {
2: ns.call_module(nniq.BNReLU2d),
3: ns.call_module(nniq.BNReLU3d),
}
}
for quant_type, dim, is_reference in options:
for instance in [BNRelu(dim, True), BNRelu(dim, False),
BNFuncRelu(dim), BNFuncInplaceRelu(dim)]:
self.checkGraphModeFxOp(
instance, self.img_data_dict[dim], quant_type,
quantized_nodes[is_reference][dim], is_reference=is_reference)
def _test_activation_impl(
self, float_module, float_op, quantized_module, quantized_op):
''' Test for activation op(with inplace options), float_op can be
torch op or functional op
'''
class M(torch.nn.Module):
def __init__(self, is_module, inplace):
super().__init__()
self.is_module = is_module
self.inplace = inplace
if self.is_module:
self.op = float_module(self.inplace)
else:
self.op = float_op
def forward(self, input):
if self.is_module:
return self.op(input)
else:
return self.op(input, self.inplace)
options = itertools.product([True, False], [True, False], self.static_quant_types, [True, False])
quantized_nodes = {
# is_module
True: {
# is_reference
True: ns.call_module(float_module),
False: ns.call_module(quantized_module),
},
False: {
True: ns.call_function(float_op),
False: ns.call_function(quantized_op),
}
}
for is_module, is_inplace, quant_type, is_reference in options:
self.checkGraphModeFxOp(
M(is_module, is_inplace), self.img_data_2d,
quant_type, quantized_nodes[is_module][is_reference], is_reference=is_reference)
def test_hardswish(self):
self._test_activation_impl(nn.Hardswish, F.hardswish, nnq.Hardswish, torch.ops.quantized.hardswish)
def test_elu(self):
self._test_activation_impl(nn.ELU, F.elu, nnq.ELU, torch.ops.quantized.elu)
def test_leaky_relu(self):
self._test_activation_impl(nn.LeakyReLU, F.leaky_relu, nnq.LeakyReLU, torch.ops.quantized.leaky_relu)
def test_prelu(self):
class M(torch.nn.Module):
def __init__(self, num_param: int):
super().__init__()
self.op = torch.nn.PReLU(num_parameters=num_param)
def forward(self, input):
return self.op(input)
X = [[torch.randn(4, 4, 4, 4, dtype=torch.float)]]
options = itertools.product([1, 4], self.static_quant_types, [True, False])
quantized_nodes = {
# is_reference
True: ns.call_module(torch.nn.PReLU),
False: ns.call_module(torch.ao.nn.quantized.PReLU),
}
for num_parameter, quant_type, is_reference in options:
self.checkGraphModeFxOp(
M(num_parameter), X, quant_type, quantized_nodes[is_reference],
is_reference=is_reference)
def _test_norm_impl(
self, float_module, float_op, op_args, data, quantized_module, quantized_op,
skip_op_arg_for_functional=False):
''' Test for normalization op, float_op can be torch op or functional op,
op_args is a list of positional argument for the module/op
'''
class M(torch.nn.Module):
def __init__(self, is_module):
super().__init__()
self.is_module = is_module
if self.is_module:
self.op = float_module(*op_args)
else:
self.op = float_op
def forward(self, input):
if self.is_module:
return self.op(input)
else:
args = [input]
if not skip_op_arg_for_functional:
args += op_args
return self.op(*args)
options = itertools.product([True, False], self.static_quant_types)
quantized_nodes = {
# is_module
True: ns.call_module(quantized_module),
False: ns.call_function(quantized_op),
}
for is_module, quant_type in options:
self.checkGraphModeFxOp(
M(is_module), data, quant_type, quantized_nodes[is_module])
def _test_norm_float16_impl(
self, float_module, float_op, op_args, data,
skip_op_arg_for_functional=False):
''' Test for normalization op, float_op can be torch op or functional op,
op_args is a list of positional argument for the module/op
'''
class M(torch.nn.Module):
def __init__(self, is_module):
super().__init__()
self.is_module = is_module
if self.is_module:
self.op = float_module(*op_args)
else:
self.op = float_op
def forward(self, input):
if self.is_module:
return self.op(input)
else:
args = [input]
if not skip_op_arg_for_functional:
args += op_args
return self.op(*args)
options = itertools.product([True, False], self.static_quant_types)
qconfig_dict = {
"object_type": [
(float_module, float16_static_qconfig),
(float_op, float16_static_qconfig)
]
}
node_occurrence = {
ns.call_method("to"): 2
}
for is_module, quant_type in options:
self.checkGraphModeFxOp(
M(is_module), data, quant_type, custom_qconfig_dict=qconfig_dict, expected_node_occurrence=node_occurrence)
def test_layer_norm(self):
data = (torch.rand((1, 2, 5, 5), dtype=torch.float),)
self._test_norm_impl(
nn.LayerNorm, F.layer_norm, [[2, 5, 5]], data, nnq.LayerNorm, torch.ops.quantized.layer_norm)
def test_instance_norm(self):
data_1d = (torch.rand((1, 4, 5), dtype=torch.float),)
data_2d = (torch.rand((1, 4, 5, 1), dtype=torch.float),)
data_3d = (torch.rand((1, 4, 5, 1, 1), dtype=torch.float),)
data_dict = {1 : data_1d, 2 : data_2d, 3 : data_3d}
instance_norm_modules = {1 : nn.InstanceNorm1d,
2 : nn.InstanceNorm2d,
3 : nn.InstanceNorm3d}
quantized_instance_norm_modules = {
1 : nnq.InstanceNorm1d,
2 : nnq.InstanceNorm2d,
3 : nnq.InstanceNorm3d
}
for dim in [1, 2, 3]:
data = data_dict[dim]
module = instance_norm_modules[dim]
quantized_module = quantized_instance_norm_modules[dim]
self._test_norm_impl(
module, F.instance_norm, [4], data,
quantized_module, torch.ops.quantized.instance_norm,
skip_op_arg_for_functional=True)
def test_norm_weight_bias(self):
class Linear(torch.nn.Module):
def __init__(self):
super().__init__()
self.w = torch.ones(5, 5)
self.b = torch.zeros(5)
def forward(self, x):
return torch.nn.functional.linear(x, self.w, self.b)
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.mods1 = Linear()
self.scale = torch.randn(5, 5)
self.bias = torch.randn(5, 5)
def forward(self, x):
x1 = self.mods1(x)
y = F.layer_norm(x1, [5, 5], weight=self.scale, bias=self.bias)
return y
model = M()
expected_occurrence = {
ns.call_function(torch.quantize_per_tensor): 1,
ns.call_function(torch.ops.quantized.linear): 1,
ns.call_function(torch.ops.quantized.layer_norm): 1,
ns.call_method("dequantize"): 1,
}
self.checkGraphModeFxOp(
model,
(torch.rand(5, 5),),
QuantType.STATIC,
expected_node_occurrence=expected_occurrence,
custom_qconfig_dict=get_default_qconfig_mapping().to_dict()
)
def _test_default_node_quant_handler_ops(
self, module, functional, qconfig, is_reference=True, node_list=None, additional_quant_pattern_dict=None
):
class M(torch.nn.Module):
def __init__(self, mod, func):
super().__init__()
self.module = mod()
self.functional = func
def forward(self, x):
x = self.module(x)
x = self.functional(x)
return x
if node_list is None:
node_list = []
if additional_quant_pattern_dict is None:
additional_quant_pattern_dict = {}
data = torch.randn((2, 2, 2, 2))
quant_type = QuantType.STATIC
prepare_custom_qconfig_dict = {"additional_quant_pattern": additional_quant_pattern_dict}
qconfig_dict = {"": qconfig}
m = M(module, functional).eval()
m_prep = prepare_fx(m, qconfig_dict, prepare_custom_qconfig_dict)
m_prep(data)
convert_fn = convert_to_reference_fx if is_reference else convert_fx
m_quant = convert_fn(m_prep, is_reference=is_reference)
m_quant(data)
self.checkGraphModuleNodes(m_quant, expected_node_list=node_list)
@unittest.skip("TODO: reenable with backend_config api")
def test_gelu_normal(self):
module = torch.nn.GELU
functional = torch.nn.functional.gelu
qconfig = torch.ao.quantization.get_default_qconfig("fbgemm")
is_reference = False
node_list = [
ns.call_module(module),
ns.call_function(functional),
]
self._test_default_node_quant_handler_ops(
module, functional, qconfig, is_reference, node_list)
@unittest.skip("TODO: reenable with backend_config api")
def test_softmax_normal(self):
module = torch.nn.Softmax
functional = torch.nn.functional.softmax
qconfig = torch.ao.quantization.get_default_qconfig("fbgemm")
is_reference = False
node_list = [
ns.call_module(torch.ao.nn.quantized.Softmax),
ns.call_function(functional),
]
self._test_default_node_quant_handler_ops(
module, functional, qconfig, is_reference, node_list)
@unittest.skip("This is no longer needed right now, can enable later with new api")
def test_gelu_reference(self):
module = torch.nn.GELU
functional = torch.nn.functional.gelu
qconfig = torch.ao.quantization.get_default_qconfig("fbgemm")
is_reference = True
node_list = [
ns.call_function(torch.quantize_per_tensor),
ns.call_method("dequantize"),
ns.call_module(module),
ns.call_function(torch.quantize_per_tensor),
ns.call_method('dequantize'),
ns.call_function(functional),
ns.call_function(torch.quantize_per_tensor),
ns.call_method('dequantize')
]
# TODO: change these to use backend_config
additional_patterns = {torch.nn.GELU: DefaultNodeQuantizeHandler,
torch.nn.functional.gelu: DefaultNodeQuantizeHandler}
self._test_default_node_quant_handler_ops(
module, functional, qconfig, is_reference, node_list, additional_patterns)
self._test_default_node_quant_handler_ops(module, functional, self.custom_qconfig, is_reference, node_list,
additional_quant_pattern_dict=self.common_quant_patterns)
@unittest.skip("This is no longer needed right now, can enable later with new api")
def test_softmax_reference(self):
module = torch.nn.Softmax
functional = torch.nn.functional.softmax
qconfig = torch.ao.quantization.get_default_qconfig("fbgemm")
is_reference = True
node_list = [
ns.call_function(torch.quantize_per_tensor),
ns.call_method("dequantize"),
ns.call_module(module),
ns.call_function(torch.quantize_per_tensor),
ns.call_method('dequantize'),
ns.call_function(functional),
ns.call_function(torch.quantize_per_tensor),
ns.call_method('dequantize')
]
additional_patterns = {torch.nn.Softmax: DefaultNodeQuantizeHandler,
torch.nn.functional.softmax: DefaultNodeQuantizeHandler}
self._test_default_node_quant_handler_ops(
module, functional, qconfig, is_reference, node_list, additional_patterns)
self._test_default_node_quant_handler_ops(module, functional, self.custom_qconfig, is_reference, node_list,
additional_quant_pattern_dict=self.common_quant_patterns)
@unittest.skip("This is no longer needed right now, can enable later with new api")
def test_silu_reference(self):
module = torch.nn.SiLU
functional = torch.nn.functional.silu
qconfig = float16_static_qconfig
is_reference = True
node_list = [
ns.call_method("to"),
ns.call_method("dequantize"),
ns.call_module(module),
ns.call_method("to"),
ns.call_method('dequantize'),
ns.call_function(functional),
ns.call_method("to"),
ns.call_method('dequantize')
]
self._test_default_node_quant_handler_ops(
module, functional, qconfig, is_reference, node_list)
node_list = [
ns.call_function(torch.quantize_per_tensor),
ns.call_method("dequantize"),
ns.call_module(module),
ns.call_function(torch.quantize_per_tensor),
ns.call_method("dequantize"),
ns.call_function(functional),
ns.call_function(torch.quantize_per_tensor),
ns.call_method("dequantize")
]
self._test_default_node_quant_handler_ops(module, functional, self.custom_qconfig, is_reference, node_list,
additional_quant_pattern_dict=self.common_quant_patterns)
@unittest.skip("This is no longer needed right now, can enable later with new api")
def test_mish_reference(self):
module = torch.nn.Mish
functional = torch.nn.functional.mish
qconfig = float16_static_qconfig
is_reference = True
node_list = [
ns.call_method("to"),
ns.call_method("dequantize"),
ns.call_module(module),
ns.call_method("to"),
ns.call_method('dequantize'),
ns.call_function(functional),
ns.call_method("to"),
ns.call_method('dequantize')
]
self._test_default_node_quant_handler_ops(
module, functional, qconfig, is_reference, node_list)
node_list = [
ns.call_function(torch.quantize_per_tensor),
ns.call_method("dequantize"),
ns.call_module(module),
ns.call_function(torch.quantize_per_tensor),
ns.call_method("dequantize"),
ns.call_function(functional),
ns.call_function(torch.quantize_per_tensor),
ns.call_method("dequantize")
]
self._test_default_node_quant_handler_ops(module, functional, self.custom_qconfig, is_reference, node_list,
additional_quant_pattern_dict=self.common_quant_patterns)
def test_bmm_int_reference(self):
""" int8 is not supported for bmm so we won't produce reference
pattern for it
"""
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.bmm = torch.bmm
def forward(self, x, y):
out = self.bmm(x, y)
return out
data_x = torch.randn((2, 2, 2,))
data_y = torch.randn((2, 2, 2,))
example_inputs = (data_x, data_y)
qconfig_dict = {"": torch.ao.quantization.get_default_qconfig("fbgemm")}
is_reference = True
node_list = [
ns.call_function(torch.bmm),
]
m = M().eval()
m_prep = prepare_fx(m, qconfig_dict, example_inputs=example_inputs)
m_prep(*example_inputs)
convert_fn = convert_to_reference_fx if is_reference else convert_fx
m_quant = convert_fn(m_prep)
m_quant(*example_inputs)
self.checkGraphModuleNodes(m_quant, expected_node_list=node_list)
@skipIfNoFBGEMM
def test_clamp(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv2d(2, 2, 2).float()
self.relu6 = torch.nn.ReLU6()
self.relu6_ = torch.nn.ReLU6(True)
self.hardtanh = torch.nn.Hardtanh()
self.hardtanh_ = torch.nn.Hardtanh(inplace=True)
def forward(self, x):
x = self.conv(x)
x = self.relu6(x)
self.relu6_(x)
x = F.relu6(x)
x = torch.clamp(x, -3, 3)
x = x.clamp(-2.5, 2.5)
# x = x.clamp_(-2, 2) # Enable when quantized `clamp_` is ready
x = self.hardtanh(x)
self.hardtanh_(x)
x = F.hardtanh(x)
return x
data = (torch.rand((1, 2, 5, 5), dtype=torch.float),)
# list of node that should occur in order
node_list = [
ns.call_function(torch.quantize_per_tensor),
ns.call_module(nnq.Conv2d),
ns.call_method('dequantize')
]
for quant_type in self.static_quant_types:
self.checkGraphModeFxOp(
M(), data, quant_type, expected_node_list=node_list)
def test_fixed_qparams_ops_fp16(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.sigmoid = torch.nn.Sigmoid()
self.tanh = torch.nn.Tanh()
def forward(self, x):
x = self.sigmoid(x)
x = torch.sigmoid(x)
x = x.sigmoid()
x = self.tanh(x)
x = torch.tanh(x)
x = x.tanh()
return x
data = (torch.randn((2, 2, 2, 2), dtype=torch.float),)
quant_type = QuantType.STATIC
# TODO: use get_default_qconfig_mapping once it handles fp16
qconfig_mapping = QConfigMapping().set_global(float16_static_qconfig)
backend_config = get_test_only_legacy_native_backend_config()
node_occurrence = {
ns.call_method("to"): 7
}
self.checkGraphModeFxOp(
M(), data, quant_type, custom_qconfig_dict=qconfig_mapping,
expected_node_occurrence=node_occurrence,
backend_config=backend_config)
def test_fixed_qparams_ops_qint8(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.sigmoid = torch.nn.Sigmoid()
self.tanh = torch.nn.Tanh()
def forward(self, x):
x = self.sigmoid(x)
x = torch.sigmoid(x)
x = x.sigmoid()
x = self.tanh(x)
x = torch.tanh(x)
x = x.tanh()
return x
data = (torch.randn((2, 2, 2, 2), dtype=torch.float),)
quant_type = QuantType.STATIC
qconfig = torch.ao.quantization.QConfig(
activation=HistogramObserver.with_args(qscheme=torch.per_tensor_symmetric, dtype=torch.quint8),
weight=default_weight_observer)
qconfig_mapping = get_default_qconfig_mapping().set_global(qconfig)
node_occurrence = {
ns.call_function(torch.quantize_per_tensor): 7,
ns.call_method("dequantize"): 7
}
self.checkGraphModeFxOp(
M(), data, quant_type, custom_qconfig_dict=qconfig_mapping,
expected_node_occurrence=node_occurrence, is_reference=True)
def test_fixed_qparams_ops_wrong_qconfig(self):
""" Test that wrong qconfigs for fixed qparams ops results in the ops not being quantized.
"""
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.sigmoid = torch.nn.Sigmoid()
self.tanh = torch.nn.Tanh()
def forward(self, x):
x = self.sigmoid(x)
x = torch.sigmoid(x)
x = x.sigmoid()
x = self.tanh(x)
x = torch.tanh(x)
x = x.tanh()
return x
data = (torch.randn((2, 2, 2, 2), dtype=torch.float),)
qconfig_mapping = QConfigMapping().set_global(default_qconfig)
m = M().eval()
node_occurrence = {
ns.call_function(torch.quantize_per_tensor): 0,
ns.call_method("dequantize"): 0,
}
self.checkGraphModeFxOp(
m, data, QuantType.STATIC, custom_qconfig_dict=qconfig_mapping,
expected_node_occurrence=node_occurrence, is_reference=True)
self.assertTrue(isinstance(m.sigmoid, torch.nn.Sigmoid))
self.assertTrue(isinstance(m.tanh, torch.nn.Tanh))
@skipIfNoFBGEMM
def test_general_shape_ops(self):
""" A test that checks dequantize will be swapped for
all supported general shape ops like aten::flatten
without actually checking for execution of these ops
"""
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.maxpool1d = torch.nn.MaxPool1d(kernel_size=3)
self.maxpool2d = torch.nn.MaxPool2d(kernel_size=3)
self.maxpool3d = torch.nn.MaxPool3d(kernel_size=3)
self.dropout = torch.nn.Dropout()
self.conv1 = torch.nn.Conv2d(3, 3, 3)
self.conv2 = torch.nn.Conv2d(3, 3, 3)
self.relu = torch.nn.ReLU()
def forward(self, x):
x = self.conv1(x)
# add_scalar
x = x + 3
# mul_scalar
x = x * 3
# add_scalar_out
x += 3
# mul_scalar_out
x *= 3
# add_scalar_relu
x = x + 3
x = F.relu(x)
# add_scalar_relu_out
x += 3
x = F.relu(x)
# mul_scalar_relu
x = x * 3
x = F.relu(x)
# mul_scalar_relu_out
x *= 3
x = F.relu(x)
x = self.maxpool1d(x)
x = self.maxpool2d(x)
x = self.maxpool3d(x)
x = torch.flatten(x)
x = x.reshape([-1])
x = x.resize_(1, 1, x)
x = x.view(-1)
# prim::ListConstruct
xs = [x, x]
# prim::ListUnpack
x, y = xs
# prim::TupleConstruct
xs = (x, x)
# prim::TupleUnpack
x, y = xs
x = x.transpose(1, 2)
x = x.contiguous()
# chunk is not supported since observer only supports
# observing single Tensor currently
x, y = torch.chunk(x, 2)
x = F.dropout(x)
x = self.dropout(x)
x = x.permute(0, 2, 3, 1)
x = x.repeat_interleave(3, 1)
x = torch.repeat_interleave(x, 3, 1)
x = self.relu(x)
x = F.relu(x)
x = F.relu(x, inplace=True)
x = x.relu()
x.relu_()
x = x.squeeze(0)
x.squeeze_(0)
x = torch.squeeze(x, 0)
x = x.unsqueeze(0)
x.unsqueeze_(0)
x = torch.unsqueeze(x, 0)
x = x.detach()
x.detach_()
x = x.repeat(4, 2)
y = []
y.append(x)
z = torch.stack(y, 0)
z = [z, z]
x, _ = z
x = self.conv2(x)
return x
example_inputs = (torch.rand(1, 3, 10, 10),)
# This model is not executable since we just put all ops
# in the same forward
m = M().eval()
qconfig_dict = {'': default_qconfig}
prepared = prepare_fx(m, qconfig_dict, example_inputs=example_inputs)
# not runnable
quantized = convert_fx(prepared)
# This checks that the dequantize from the output of first conv
# is being propagated to the end, so that we don't insert extra
# observers and also successfully fused two quantized::conv2d
# patterns
# one quantize_per_tensor for input
# check exact counts of quantize and dequantize
count_check = {
# input of conv and two outputs of getitem
ns.call_function(torch.quantize_per_tensor) : 2,
# output of the model and two outputs of getitem
ns.call_method('dequantize') : 2
}
order_check = [
ns.call_function(torch.quantize_per_tensor),
ns.call_module(nnq.Conv2d),
ns.call_module(nnq.Conv2d),
ns.call_method('dequantize'),
]
self.checkGraphModuleNodes(
quantized,
expected_node_occurrence=count_check,
expected_node_list=order_check)
# Checking the is_reference output
m = M().eval()
qconfig_dict = {'': default_qconfig}
prepared = prepare_fx(m, qconfig_dict, example_inputs=example_inputs)
# not runnable
quantized = convert_to_reference_fx(prepared)
@skipIfNoFBGEMM
def test_ave_pool_with_custom_cfg(self):
""" A test that checks correct patterns are produced for
avg_pool2d with customized config
"""
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.avg_pool2d = torch.nn.AvgPool2d(3)
def forward(self, x):
x = self.avg_pool2d(x)
return x
# This model is not executable since we just put all ops
# in the same forward
m = M().eval()
# nothing to fuse so skipping the fuse step
qconfig_dict = {'': default_qconfig}
example_inputs = (torch.randn(1, 3, 3, 3),)
prepared = prepare_fx(
m, qconfig_dict, example_inputs=example_inputs,
prepare_custom_config={"input_quantized_idxs": [0]})
# not runnable
quantized = convert_fx(prepared)
# This checks that the dequantize from the output of first conv
# is being propagated to the end, so that we don't insert extra
# observers
# check exact counts of quantize and dequantize
count_check = {
ns.call_method('dequantize') : 1
}
order_check = [
ns.call_module(nn.AvgPool2d),
ns.call_method('dequantize'),
]
self.checkGraphModuleNodes(
quantized,
expected_node_occurrence=count_check,
expected_node_list=order_check)
@skipIfNoFBGEMM
def test_general_value_ops(self):
""" A test that checks correct patterns are produced for
all supported general value ops like aten::avg_pool2d \
without actually checking for execution of these ops
"""
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv2d(3, 3, 3)
self.avg_pool1d = torch.nn.AvgPool1d(3)
self.avg_pool2d = torch.nn.AvgPool2d(3)
self.avg_pool3d = torch.nn.AvgPool3d(3)
self.adaptive_avg_pool1d = torch.nn.AdaptiveAvgPool1d(1)
self.adaptive_avg_pool2d = torch.nn.AdaptiveAvgPool2d((1, 1))
self.adaptive_avg_pool3d = torch.nn.AdaptiveAvgPool3d((1, 1, 1))
def forward(self, x):
x = self.conv(x)
x = self.avg_pool1d(x)
x = self.avg_pool2d(x)
x = self.avg_pool3d(x)
x = self.adaptive_avg_pool1d(x)
x = self.adaptive_avg_pool2d(x)
x = self.adaptive_avg_pool3d(x)
x = F.avg_pool1d(x, 3)
x = F.avg_pool2d(x, 3)
x = F.avg_pool3d(x, 3)
x = F.adaptive_avg_pool1d(x, (1))
x = F.adaptive_avg_pool2d(x, (1, 1))
x = F.adaptive_avg_pool3d(x, (1, 1, 1))
x = torch.mean(x)
x = torch.mean(x, [2, 3], False)
x = x.mean()
x = x.mean([2, 3], True)
x = F.interpolate(x, 4, mode='nearest')
x = F.interpolate(x, 4, mode='linear')
x = self.conv(x)
return x
# This model is not executable since we just put all ops
# in the same forward
m = M().eval()
# nothing to fuse so skipping the fuse step
qconfig_dict = {'': default_qconfig}
example_inputs = (torch.randn(1, 3, 3, 3),)
prepared = prepare_fx(m, qconfig_dict, example_inputs=example_inputs)
# not runnable
quantized = convert_fx(prepared)
# This checks that the dequantize from the output of first conv
# is being propagated to the end, so that we don't insert extra
# observers
# check exact counts of quantize and dequantize
count_check = {
ns.call_function(torch.quantize_per_tensor) : 1,
ns.call_method('dequantize') : 1
}
order_check = [
ns.call_function(torch.quantize_per_tensor),
ns.call_module(nnq.Conv2d),
ns.call_module(nnq.Conv2d),
ns.call_method('dequantize'),
]
self.checkGraphModuleNodes(
quantized,
expected_node_occurrence=count_check,
expected_node_list=order_check)
def test_copy_node_fp32_input(self):
""" CopyNode works for both fp32 and int8 inputs, this is a test to make
sure that a CopyNode can be successfully quantized in both cases
"""
class M(torch.nn.Module):
def forward(self, x):
x = x.relu()
return x
m = M().eval()
m = prepare_fx(m, {"": default_reuse_input_qconfig}, example_inputs=(torch.randn(1),))
m = convert_fx(m)
# make sure it runs
m(torch.rand(1))
def test_getitem(self):
""" Make sure we only insert observer for getitem if the following node is matched
or needs to be quantized
"""
class M(torch.nn.Module):
def forward(self, xs):
x = xs[0]
return x
m = M().eval()
example_inputs = (torch.rand(1, 2),)
qconfig_mapping = get_default_qconfig_mapping()
m = prepare_fx(m, qconfig_mapping, example_inputs=example_inputs)
self.checkGraphModuleNodes(m, expected_node_occurrence={
ns.call_module(torch.ao.quantization.MinMaxObserver): 0
})
m = convert_fx(m)
m(*example_inputs)
class M2(torch.nn.Module):
def forward(self, xs):
x = xs[0]
x = torch.sigmoid(x)
return x
m2 = M2().eval()
example_inputs = ([torch.rand(1, 2)],)
qconfig_mapping = get_default_qconfig_mapping()
m2 = prepare_fx(m2, qconfig_mapping, example_inputs=example_inputs)
self.checkGraphModuleNodes(m2, expected_node_occurrence={
ns.call_module(torch.ao.quantization.FixedQParamsObserver): 2
})
m2 = convert_fx(m2)
self.checkGraphModuleNodes(m2, expected_node_list=[
ns.call_function(torch.quantize_per_tensor),
ns.call_method("dequantize")
])
m2(*example_inputs)
# testing prepare recognizes non-Tensor input for getitem
class M3(torch.nn.Module):
def forward(self, x):
s = x.shape
n, c = s[:2]
x = torch.sigmoid(x)
return x
m3 = M3().eval()
example_inputs = (torch.rand(1, 2, 3, 4),)
qconfig_mapping = get_default_qconfig_mapping()
m3 = prepare_fx(m3, qconfig_mapping, example_inputs=example_inputs)
self.checkGraphModuleNodes(m3, expected_node_occurrence={
ns.call_module(torch.ao.quantization.FixedQParamsObserver): 2
})
m3 = convert_fx(m3)
self.checkGraphModuleNodes(m3, expected_node_list=[
ns.call_function(torch.quantize_per_tensor),
ns.call_method("dequantize")
])
m3(*example_inputs)
@skipIfNoFBGEMM
def test_fixed_qparams_ops(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv2d(3, 3, 3)
self.sigmoid = torch.nn.Sigmoid()
self.hardsigmoid = torch.nn.Hardsigmoid()
self.tanh = torch.nn.Tanh()
self.softmax = torch.nn.Softmax(dim=0)
def forward(self, x):
x = self.conv(x)
# F.sigmoid is deprecated
x = self.sigmoid(x)
x = torch.sigmoid(x)
x = x.sigmoid()
x = self.hardsigmoid(x)
x = F.hardsigmoid(x)
x = F.hardsigmoid(x, inplace=True)
x = self.tanh(x)
# F.tanh is deprecated
x = torch.tanh(x)
x = x.tanh()
# TODO(future PR): handle F.softmax
x = self.softmax(x)
return x
for eval_mode in [True, False]:
# This model is not executable since we just put all ops
# in the same forward
m = M()
if eval_mode:
m.eval()
qconfig_mapping = get_default_qconfig_mapping()
prepare = prepare_fx
fq_count = 10
else:
m.train()
qconfig_mapping = get_default_qat_qconfig_mapping()
prepare = prepare_qat_fx
fq_count = 10
# nothing to fuse so skipping the fuse step
m_copy = copy.deepcopy(m)
example_inputs = (torch.rand(3, 3, 3, 3),)
prepared = prepare(m, qconfig_mapping, example_inputs=example_inputs)
prepared_copy = copy.deepcopy(prepared)
# check that prepare does not change model result
if eval_mode:
self.assertEqual(m_copy(*example_inputs), prepared_copy(*example_inputs))
# check the correct number of activation_post_process is inserted
expected_activation_post_process = FixedQParamsObserver if eval_mode else FixedQParamsFakeQuantize
count_check = {
ns.call_module(expected_activation_post_process) : fq_count,
}
self.checkGraphModuleNodes(
prepared,
expected_node_occurrence=count_check)
# not runnable
quantized = convert_fx(prepared)
quantized_reference = convert_to_reference_fx(prepared_copy)
# This checks that the dequantize from the output of first conv
# is being propagated to the end, so that we don't insert extra
# observers
# check exact counts of quantize and dequantize
count_check = {
ns.call_function(torch.quantize_per_tensor) : 1,
ns.call_method('dequantize') : 1
}
order_check = [
ns.call_function(torch.quantize_per_tensor),
ns.call_module(nnq.Conv2d),
ns.call_module(nn.Sigmoid),
ns.call_module(nnq.Softmax),
ns.call_method('dequantize'),
]
self.checkGraphModuleNodes(
quantized,
expected_node_occurrence=count_check,
expected_node_list=order_check)
reference_count_check = {
ns.call_function(torch.quantize_per_tensor) : 12,
ns.call_method('dequantize') : 12
}
reference_order_check = [
ns.call_function(torch.quantize_per_tensor),
ns.call_method('dequantize'),
ns.call_module(nnqr.Conv2d),
ns.call_function(torch.quantize_per_tensor),
ns.call_method('dequantize'),
ns.call_module(nn.Sigmoid),
ns.call_function(torch.quantize_per_tensor),
ns.call_method('dequantize'),
ns.call_module(nn.Softmax),
ns.call_function(torch.quantize_per_tensor),
ns.call_method('dequantize'),
]
self.checkGraphModuleNodes(
quantized_reference,
expected_node_occurrence=reference_count_check,
expected_node_list=reference_order_check)
# Verify that softmax scale and zero_point are correct
self.assertTrue(quantized.softmax.scale - (1.0 / 256) <= 1e-8)
self.assertTrue(quantized.softmax.zero_point == 0)
def test_float_functional(self):
class TorchAdd(nn.Module):
"""Wrapper around torch.add so that all ops can be found at build"""
def __init__(self):
super().__init__()
self.add_func = nnq.FloatFunctional()
def forward(self, x, y):
return self.add_func.add(x, y)
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.ff1 = TorchAdd()
self.ff2 = nnq.FloatFunctional()
self.ff3 = nnq.FloatFunctional()
self.ff4 = nnq.FloatFunctional()
self.ff5 = nnq.FloatFunctional()
self.ff6 = nnq.FloatFunctional()
def forward(self, x):
x = self.ff1(x, x)
x = self.ff2.add_scalar(x, 3)
x = self.ff3.mul(x, x)
x = self.ff4.mul_scalar(x, 3)
x = self.ff5.add_relu(x, x)
x = self.ff6.cat([x])
return x
example_inputs = (torch.rand(3, 3),)
# Note: QAT test succeeded by chance, to make it actually work
# we need to fix eager mode FloatFunctional by removing
# activation_post_process in add_scalar and mul_scalar
for quant_type in self.static_quant_types:
m = M()
ref_m = torch.ao.quantization.QuantWrapper(M())
is_qat = quant_type == QuantType.QAT
if is_qat:
m.train()
ref_m.train()
qconfig = default_qat_qconfig
expected_act_post_process = torch.ao.quantization.FakeQuantize
else:
m.eval()
ref_m.eval()
qconfig = default_qconfig
expected_act_post_process = torch.ao.quantization.MinMaxObserver
prepare_fx_function = prepare_qat_fx if is_qat else prepare_fx
qconfig_dict = {"": qconfig}
m = prepare_fx_function(m, qconfig_dict, example_inputs=example_inputs)
node_occurrence = {
ns.call_module(expected_act_post_process): 7,
ns.call_module(torch.ao.nn.quantized.FloatFunctional): 0
}
self.checkGraphModuleNodes(m, expected_node_occurrence=node_occurrence)
m(*example_inputs)
node_list = [
ns.call_function(torch.quantize_per_tensor),
ns.call_function(torch.ops.quantized.add),
ns.call_function(torch.ops.quantized.add),
ns.call_function(torch.ops.quantized.mul),
ns.call_function(torch.ops.quantized.mul),
ns.call_function(torch.ops.quantized.add_relu),
ns.call_function(torch.cat),
ns.call_method('dequantize')
]
m = convert_fx(m)
self.checkGraphModuleNodes(m, expected_node_list=node_list)
# make sure numerics match with eager mode
ref_m.qconfig = qconfig
prepare_function = prepare_qat if is_qat else prepare
ref_m = prepare_function(ref_m)
ref_m(*example_inputs)
ref_m = convert(ref_m)
# FX Graph Mode and Eager Mode now diverages in numerics of add_scalar and mul_scalar
# self.assertEqual(m(data), ref_m(data))
def test_embedding(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.emb = torch.nn.Embedding(num_embeddings=10, embedding_dim=12)
def forward(self, indices):
return self.emb(indices)
for qconfig_type in [float_qparams_weight_only_qconfig, float_qparams_weight_only_qconfig_4bit]:
model = M().eval()
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])
example_inputs = (indices,)
quantized_node = ns.call_module(nnq.Embedding)
configs = [
(qconfig_type, ns.call_module(nnq.Embedding)),
(None, ns.call_module(nn.Embedding)),
(default_qconfig, ns.call_module(nn.Embedding)),
]
for qconfig, node in configs:
qconfig_dict = {"": qconfig}
m = prepare_fx(model, qconfig_dict, example_inputs=example_inputs)
self.checkGraphModuleNodes(m, expected_node_occurrence={
ns.call_module(torch.ao.quantization.MinMaxObserver): 0
})
m = convert_fx(m)
self.checkGraphModuleNodes(m, expected_node=node)
# make sure it runs
m(*example_inputs)
def test_embedding_bag(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.emb = torch.nn.EmbeddingBag(num_embeddings=10, embedding_dim=12, include_last_offset=True)
def forward(self, indices, offsets):
return self.emb(indices, offsets)
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])
quantized_node = ns.call_module(nnq.EmbeddingBag)
example_inputs = (indices, offsets)
for dtype in [torch.quint8, torch.quint4x2]:
model = M().eval()
float_qparams_observer = PerChannelMinMaxObserver.with_args(dtype=dtype,
qscheme=torch.per_channel_affine_float_qparams,
ch_axis=0)
float_qparams_qconfig = QConfig(activation=default_placeholder_observer,
weight=float_qparams_observer)
self.checkGraphModeFxOp(
model,
example_inputs,
QuantType.DYNAMIC,
quantized_node,
custom_qconfig_dict={"": float_qparams_qconfig}
)
# check it works in None and static qconfig
for qconfig in [None, default_qconfig]:
qconfig_dict = {"": default_qconfig}
m = M().eval()
m = prepare_fx(model, qconfig_dict, example_inputs=example_inputs)
self.checkGraphModuleNodes(m, expected_node_occurrence={
ns.call_module(torch.ao.quantization.MinMaxObserver): 0
})
m = convert_fx(m)
self.checkGraphModuleNodes(m, expected_node=ns.call_module(nn.EmbeddingBag))
# make sure it runs
m(*example_inputs)
def _test_rnn_impl(self, qconfigs, M, module_type_strs, module_types, sample_input):
options = itertools.product(qconfigs, module_type_strs)
for qconfig, module_type_str in options:
model_eager = M(module_type_str).eval()
model_graph = copy.deepcopy(model_eager)
if torch.backends.quantized.engine == 'qnnpack' and \
qconfig is float16_dynamic_qconfig:
continue
# fp16 dynamic quant is not supported for qnnpack
eager_qconfig_dict = {x : qconfig for x in module_types}
model_eager = quantize_dynamic(model_eager, qconfig_spec=eager_qconfig_dict)
graph_qconfig_dict = {
"object_type": [
(x, qconfig) for x in module_types
]
}
model_graph = prepare_fx(model_graph, graph_qconfig_dict, example_inputs=(sample_input,))
model_graph = convert_fx(model_graph)
self.assertEqual(model_eager(sample_input), model_graph(sample_input))
self.checkScriptable(model_graph, [[sample_input]], True)
@override_qengines
def test_rnn_cell(self):
if torch.backends.quantized.engine not in ('fbgemm', 'qnnpack'):
return
qconfigs = [per_channel_dynamic_qconfig, default_dynamic_qconfig, float16_dynamic_qconfig]
module_type_strs = ['LSTMCell', 'GRUCell', 'RNNTanh', 'RNNReLU']
module_types = [torch.nn.LSTMCell, torch.nn.GRUCell, torch.nn.RNNCell]
sample_input = torch.tensor([[100, -155],
[-155, 100],
[100, -155]], dtype=torch.float)
self._test_rnn_impl(qconfigs, RNNCellDynamicModel, module_type_strs, module_types, sample_input)
@override_qengines
def test_rnn(self):
if torch.backends.quantized.engine not in ('fbgemm', 'qnnpack'):
return
qconfigs = [per_channel_dynamic_qconfig, default_dynamic_qconfig, float16_dynamic_qconfig]
module_type_strs = ['LSTM', 'GRU']
module_types = [torch.nn.LSTM, torch.nn.GRU]
niter = 10
sample_input = torch.tensor([[100, -155],
[-155, 100],
[100, -155]], dtype=torch.float).unsqueeze(0).repeat(niter, 1, 1)
self._test_rnn_impl(qconfigs, RNNDynamicModel, module_type_strs, module_types, sample_input)
def _test_conv_transpose_impl(
self, float_cls: Callable, q_cls: Callable, data: torch.Tensor):
with override_quantized_engine('qnnpack'):
# Create fp32 versions of FX and Eager models
m1 = torch.nn.Sequential(float_cls(1, 1, 1))
m2 = torch.nn.Sequential(float_cls(1, 1, 1))
m2.load_state_dict(m1.state_dict())
m2 = torch.ao.quantization.QuantWrapper(m2)
# FX graph
result_dict = self.checkGraphModeFxOp(
m1, (data,), QuantType.STATIC,
expected_node_occurrence={
ns.call_module(q_cls): 1,
})
q_result1 = result_dict["quantized_output"]
# Eager
m2.qconfig = get_default_qconfig(torch.backends.quantized.engine)
m2.eval()
m2p = torch.ao.quantization.prepare(m2)
m2p(data)
m2q = torch.ao.quantization.convert(m2p)
q_result2 = m2q(data)
# verify results match
self.assertEqual(q_result1, q_result2)
@unittest.skipUnless('qnnpack' in supported_qengines,
"This Pytorch Build has not been built with or does not support QNNPACK")
def test_conv_transpose_1d(self):
self._test_conv_transpose_impl(
torch.nn.ConvTranspose1d, nnq.ConvTranspose1d, torch.randn(4, 1, 4))
@unittest.skipUnless('qnnpack' in supported_qengines,
"This Pytorch Build has not been built with or does not support QNNPACK")
def test_conv_transpose_2d(self):
self._test_conv_transpose_impl(
torch.nn.ConvTranspose2d, nnq.ConvTranspose2d, torch.randn(4, 1, 4, 4))
def test_reshape_fp16(self):
class M(torch.nn.Module):
def __init__(self, w, b):
super().__init__()
self.w = w
self.b = b
def forward(self, x):
x = torch.nn.functional.linear(x, self.w)
x = x.reshape(-1, 4)
x = torch.nn.functional.linear(x, self.w)
return x
w = torch.randn(4, 4)
b = torch.randn(4)
m = M(w, b).eval()
qconfig_dict = {
# reshape will be quantized to fp16 as requested by this qconfig
"": float16_static_qconfig,
"object_type": [
(torch.nn.functional.linear, default_qconfig)
]
}
backend_config = get_test_only_legacy_native_backend_config()
example_inputs = (torch.randn(1, 4),)
m = prepare_fx(
m, qconfig_dict, example_inputs=example_inputs,
backend_config=backend_config)
expected_occurrence = {
# input and weight of first and second linear, output of first and second linear
ns.call_module(torch.ao.quantization.MinMaxObserver): 6,
# we insert placeholder observer for both input and output of reshape
ns.call_module(torch.ao.quantization.PlaceholderObserver): 2
}
self.checkGraphModuleNodes(
m,
expected_node_occurrence=expected_occurrence
)
m = convert_fx(m, backend_config=backend_config)
expected_occurrence = {
ns.call_function(torch.quantize_per_tensor): 2,
# dequantize after first linear, before reshape and before output
ns.call_method("dequantize"): 3,
# before reshape, to(fp16)
ns.call_method("to"): 1,
ns.call_function(torch.ops.quantized.linear): 2
}
self.checkGraphModuleNodes(
m,
expected_node_occurrence=expected_occurrence
)
# make sure it runs
m(torch.randn(2, 4))
def test_multiple_qconfigs_for_single_value(self):
""" Test multiple qconfigs for a single value"""
class M(torch.nn.Module):
def __init__(self, w, b):
super().__init__()
self.w = w
self.b = b
def forward(self, x):
x = torch.nn.functional.linear(x, self.w)
x = torch.sigmoid(x)
return x
w = torch.randn(4, 4)
b = torch.randn(4)
m = M(w, b).eval()
# TODO: use get_default_qconfig_mapping once it handles fp16
qconfig_mapping = QConfigMapping() \
.set_global(float16_static_qconfig) \
.set_object_type(torch.nn.functional.linear, default_qconfig)
example_inputs = (torch.randn(1, 4),)
backend_config = get_test_only_legacy_native_backend_config()
m = prepare_fx(
m, qconfig_mapping, example_inputs=example_inputs,
backend_config=backend_config)
expected_occurrence = {
# input and weight of linear, output of linear
ns.call_module(torch.ao.quantization.MinMaxObserver): 3,
# input and output of sigmoid
ns.call_module(torch.ao.quantization.PlaceholderObserver): 2,
}
self.checkGraphModuleNodes(
m,
expected_node_occurrence=expected_occurrence
)
# make sure it runs
m = convert_fx(m)
expected_occurrence = {
ns.call_function(torch.quantize_per_tensor): 1,
ns.call_method("dequantize"): 3,
ns.call_method("to"): 2
}
self.checkGraphModuleNodes(
m,
expected_node_occurrence=expected_occurrence
)
def test_boolean_tensor(self):
""" Make sure we don't insert observer for boolean Tensors """
class M(torch.nn.Module):
def forward(self, x, mask):
mask = mask.unsqueeze(0)
mask = mask.unsqueeze(1)
x = x.masked_fill(mask, 1)
return x
m = M().eval()
example_inputs = (torch.rand(1, 2, 3, 4), torch.rand(3, 4).bool())
m = prepare_fx(m, {"": default_qconfig}, example_inputs=example_inputs)
expected_occurrence = {
ns.call_module(torch.ao.quantization.MinMaxObserver): 0
}
self.checkGraphModuleNodes(
m,
expected_node_occurrence=expected_occurrence)
m = convert_fx(m)
m(*example_inputs)
def test_chunk(self):
class M(torch.nn.Module):
def forward(self, x):
x, y = torch.chunk(x, 2)
x = x + y
return x
m = M().eval()
example_inputs = (torch.rand(2, 2, 2, 2),)
m = prepare_fx(m, {"": default_qconfig}, example_inputs=example_inputs)
m(*example_inputs)
m = convert_fx(m)
m(*example_inputs)
# make sure everything runs
def test_ref_pattern_multi_use(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(5, 5)
self.linear1 = torch.nn.Linear(5, 5)
def forward(self, x):
y = self.linear(x)
z = self.linear1(x)
a = torch.mul(z, 5)
b = torch.add(z, 5)
return (y, a, b)
m = M().eval()
qconfig_dict = {
"": None,
"object_type": [
(torch.nn.Linear, get_default_qconfig("fbgemm")),
(torch.nn.ReLU, get_default_qconfig("fbgemm")),
],
}
example_inputs = (torch.randn(1, 5),)
m = prepare_fx(m, qconfig_dict, example_inputs=example_inputs)
m = convert_fx(m)
expected_occurrence = {
ns.call_function(torch.quantize_per_tensor): 1,
ns.call_module(nnq.Linear): 2,
ns.call_method("dequantize"): 2,
ns.call_function(torch.add): 1,
ns.call_function(torch.mul): 1,
}
self.checkGraphModuleNodes(
m,
expected_node_occurrence=expected_occurrence)
def test_qmatmul(self):
class M(torch.nn.Module):
def forward(self, x, y):
z = torch.matmul(x, y)
return z
m = M().eval()
example_inputs = (torch.randn(2, 2), torch.randn(2, 2))
qconfig_dict = get_default_qconfig_mapping("fbgemm")
mp = prepare_fx(m, qconfig_dict, example_inputs=example_inputs)
mp(*example_inputs)
mq = convert_fx(mp)
expected_occurrence = {
ns.call_function(torch.matmul): 0,
ns.call_function(torch.ops.quantized.matmul): 1,
}
self.checkGraphModuleNodes(
mq,
expected_node_occurrence=expected_occurrence)
# verify no crash
res = mq(*example_inputs)
def test_pixel_shuffle(self):
class MyBias(nn.Module):
def __init__(self):
super().__init__()
self.bias = nn.Parameter(torch.randn(8))
class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(8, 8, 1, bias=False)
self.bias = MyBias()
def forward(self, x):
x = self.conv(x)
x = nn.functional.pixel_shuffle(x, 2)
x = x.view(-1, 8, 2, 2)
bias = self.bias.bias
return x + bias
backend_config = get_qnnpack_backend_config()
qconfig_mapping = get_default_qconfig_mapping("qnnpack")
model = MyModel()
m = prepare_fx(
model,
qconfig_mapping=qconfig_mapping,
example_inputs=(torch.randn(1, 8, 3, 3),),
backend_config=backend_config
)
m = convert_fx(m)
expected_occurrence = {
ns.call_function(torch.quantize_per_tensor): 2,
ns.call_method("dequantize"): 1,
}
self.checkGraphModuleNodes(m, expected_node_occurrence=expected_occurrence)
def test_pixel_shuffle_module(self) -> None:
class MyBias(nn.Module):
def __init__(self):
super().__init__()
self.bias = nn.Parameter(torch.randn(8))
class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(8, 8, 1, bias=False)
self.ps = nn.PixelShuffle(upscale_factor=2)
self.bias = MyBias()
def forward(self, x):
x = self.conv(x)
x = self.ps(x)
x = x.view(-1, 8, 2, 2)
bias = self.bias.bias
return x + bias
backend_config = get_qnnpack_backend_config()
qconfig_mapping = get_default_qconfig_mapping("qnnpack")
model = MyModel()
m = prepare_fx(
model,
qconfig_mapping=qconfig_mapping,
example_inputs=(torch.randn(1, 8, 3, 3),),
backend_config=backend_config
)
m = convert_fx(m)
expected_occurrence = {
ns.call_function(torch.quantize_per_tensor): 2,
ns.call_method("dequantize"): 1,
ns.call_module(nn.PixelShuffle): 1,
}
self.checkGraphModuleNodes(m, expected_node_occurrence=expected_occurrence)
def test_pixel_unshuffle(self):
class MyBias(nn.Module):
def __init__(self):
super().__init__()
self.bias = nn.Parameter(torch.randn(64))
class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(8, 8, 1, bias=False)
self.bias = MyBias()
def forward(self, x):
x = self.conv(x)
x = nn.functional.pixel_unshuffle(x, 2)
bias = self.bias.bias
return x + bias
for backend in ["fbgemm", "qnnpack"]:
if backend == "fbgemm":
backend_config = get_fbgemm_backend_config()
else:
backend_config = get_qnnpack_backend_config()
qconfig_mapping = get_default_qconfig_mapping(backend)
model = MyModel()
m = prepare_fx(
model,
qconfig_mapping=qconfig_mapping,
example_inputs=(torch.randn(1, 8, 6, 6),),
backend_config=backend_config
)
m = convert_fx(m)
expected_occurrence = {
ns.call_function(torch.quantize_per_tensor): 2,
ns.call_method("dequantize"): 1,
}
self.checkGraphModuleNodes(m, expected_node_occurrence=expected_occurrence)
def test_pixel_unshuffle_module(self) -> None:
class MyBias(nn.Module):
def __init__(self):
super().__init__()
self.bias = nn.Parameter(torch.randn(64))
class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(8, 8, 1, bias=False)
self.unshuffle = nn.PixelUnshuffle(downscale_factor=2)
self.bias = MyBias()
def forward(self, x):
x = self.conv(x)
x = self.unshuffle(x)
bias = self.bias.bias
return x + bias
for backend in ["fbgemm", "qnnpack"]:
if backend == "fbgemm":
backend_config = get_fbgemm_backend_config()
else:
backend_config = get_qnnpack_backend_config()
qconfig_mapping = get_default_qconfig_mapping(backend)
model = MyModel()
m = prepare_fx(
model,
qconfig_mapping=qconfig_mapping,
example_inputs=(torch.randn(1, 8, 6, 6),),
backend_config=backend_config
)
m = convert_fx(m)
expected_occurrence = {
ns.call_function(torch.quantize_per_tensor): 2,
ns.call_method("dequantize"): 1,
ns.call_module(nn.PixelUnshuffle): 1,
}
self.checkGraphModuleNodes(m, expected_node_occurrence=expected_occurrence)
def test_narrow(self):
class MyBias(nn.Module):
def __init__(self):
super().__init__()
self.bias = nn.Parameter(torch.randn(4))
class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(8, 8, 1, bias=False)
self.bias = MyBias()
def forward(self, x):
x = self.conv(x)
x = torch.narrow(x, 1, 0, 4)
bias = self.bias.bias
return x + bias
for backend in ["fbgemm", "qnnpack"]:
if backend == "fbgemm":
backend_config = get_fbgemm_backend_config()
else:
backend_config = get_qnnpack_backend_config()
qconfig_mapping = get_default_qconfig_mapping(backend)
model = MyModel()
m = prepare_fx(
model,
qconfig_mapping=qconfig_mapping,
example_inputs=(torch.randn(1, 8, 3, 3),),
backend_config=backend_config
)
m = convert_fx(m)
expected_occurrence = {
ns.call_function(torch.quantize_per_tensor): 2,
ns.call_method("dequantize"): 1,
}
self.checkGraphModuleNodes(m, expected_node_occurrence=expected_occurrence)
class TestQuantizeFxModels(QuantizationTestCase):
@skipIfNoFBGEMM
@unittest.skipIf(not TEST_CUDA, "gpu is not available.")
def test_static_gpu_convert_basic(self):
class Net(nn.Module):
def __init__(self):
super().__init__()
self.relu1 = nn.ReLU()
self.conv1 = nn.Conv2d(1, 6, 5)
self.linear1 = nn.Linear(120, 1)
def forward(self, x):
x = self.relu1(self.conv1(x))
y = self.linear1(x.view(-1))
return y
input = torch.randn((5, 1, 6, 6)).to('cuda')
example_inputs = (input,)
model = Net().to('cuda').eval()
qconfig_dict = {"": torch.ao.quantization.get_default_qconfig('fbgemm')}
model_prepared = prepare_fx(model, qconfig_dict, example_inputs=example_inputs)
model_prepared(*example_inputs)
model_quantized = convert_to_reference_fx(model_prepared)
out = model_quantized(*example_inputs)
self.assertEqual(out.device.type, 'cuda')
@skipIfNoFBGEMM
@unittest.skipIf(not TEST_CUDA, "gpu is not available.")
def test_switch_device_prepare_convert(self):
class Net(nn.Module):
def __init__(self):
super().__init__()
self.relu1 = nn.ReLU()
self.conv1 = nn.Conv2d(1, 6, 5)
self.linear1 = nn.Linear(120, 1)
def forward(self, x):
x = self.relu1(self.conv1(x))
y = self.linear1(x.view(-1))
return y
for device in ['cuda', 'cpu']:
device_after = 'cuda' if device == 'cpu' else 'cpu'
input = torch.randn((5, 1, 6, 6)).to(device)
model = Net().to(device).eval()
qconfig_dict = {"": torch.ao.quantization.get_default_qconfig('fbgemm')}
model_prepared = prepare_fx(model, qconfig_dict, example_inputs=(input,))
model_prepared(input)
model_prepared.to(device_after)
model_quantized = convert_to_reference_fx(model_prepared)
out = model_quantized(input.to(device_after))
self.assertEqual(out.device.type, device_after)
@skipIfNoFBGEMM
@unittest.skipIf(not TEST_CUDA, "gpu is not available.")
def test_prepare_serialize_switch_device_convert(self):
class Net(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 6, 5)
self.linear1 = nn.Linear(120, 1)
def forward(self, x):
x = self.conv1(x)
y = self.linear1(x.view(-1))
return y
for device in ['cuda', 'cpu']:
for device_after in ['cuda', 'cpu']:
input = torch.randn((5, 1, 6, 6)).to(device)
model = Net().to(device).eval()
qconfig_dict = {"": torch.ao.quantization.get_default_qconfig('fbgemm')}
model_prepared_first = prepare_fx(model, qconfig_dict, example_inputs=(input,))
model_prepared_second = prepare_fx(model, qconfig_dict, example_inputs=(input,))
model_prepared_first(input)
state_dict = model_prepared_first.state_dict()
del model_prepared_first
model_prepared_second.load_state_dict(state_dict)
model_prepared_second.to(device_after)
model_quantized = convert_to_reference_fx(model_prepared_second)
out = model_quantized(input.to(device_after))
self.assertEqual(out.device.type, device_after)
@skip_if_no_torchvision
def test_model_dropout(self):
from torchvision import models
m = models.mobilenet_v3_small()
qconfig_mapping = torch.ao.quantization.get_default_qat_qconfig_mapping('fbgemm')
example_inputs = (torch.randn(1, 3, 224, 224),)
mp = prepare_qat_fx(m, qconfig_mapping, example_inputs=example_inputs)
mp(*example_inputs)
with override_quantized_engine("qnnpack") if IS_ARM64 else contextlib.nullcontext():
mq = convert_fx(mp)
mq(*example_inputs)
def _test_model_impl(
self, mode, name, model, eager_quantizable_model,
check_with_eager=True,
diff_of_quant=None,
diff_from_eager=None):
if diff_of_quant is None or diff_from_eager is None:
diff_of_quant = {}
diff_from_eager = {}
if mode not in diff_of_quant or mode not in diff_from_eager:
diff_of_quant[mode] = {}
diff_from_eager[mode] = {}
input_tensor = torch.rand(1, 3, 224, 224)
input_tensor_inception = torch.rand(1, 3, 299, 299)
output_value = torch.randint(0, 1, (1,))
# print('quantizing:', name, ' mode:', mode)
if name == 'inception_v3':
input_value = input_tensor_inception
else:
input_value = input_tensor
qconfig = default_qconfig if mode == 'static' else default_qat_qconfig
qconfig_dict = {'': qconfig}
script = torch.jit.script(model)
# make sure graph module and script module are both runanble
original_out = model(input_value)
is_not_tuple_out = not isinstance(original_out, tuple)
script_out = script(input_value)
# set to train just before quantization
prepare_fx_fn = prepare_fx
if mode != 'static':
model.train()
prepare_fx_fn = prepare_qat_fx
prepared = prepare_fx_fn(model, qconfig_dict)
if mode == 'ddp':
mp.spawn(run_ddp,
args=(world_size, prepared),
nprocs=world_size,
join=True)
elif mode == 'qat':
assert prepared.training, 'prepared must be in training mode for qat'
optimizer = torch.optim.SGD(prepared.parameters(), lr=0.0001)
criterion = nn.CrossEntropyLoss()
train_one_epoch(prepared, criterion, optimizer, [(input_value, output_value)], torch.device('cpu'), 1)
else:
for i in range(10):
prepared(input_value)
# print('after observation root:', prepared.root)
qgraph = convert_fx(prepared)
# print('after quantization root:', qgraph.root)
# print('after quantization code:', qgraph.src)
qgraph.eval()
qgraph_script = torch.jit.script(qgraph)
# print('quantized and scripted:', qgraph_script.graph)
qgraph_out = qgraph(input_value)
qgraph_script = qgraph_script(input_value)
if is_not_tuple_out:
diff_of_quant[mode][name] = (original_out - qgraph_out).abs().max()
assert torch.allclose(qgraph_out, qgraph_script), 'graph, scripted graph'
else:
print('tuple output')
if eager_quantizable_model is not None:
# comparing to eager mode quantization
qeager = eager_quantizable_model
ref_out = qeager(input_value)
qeager.qconfig = qconfig
if mode == 'static':
qeager.fuse_model()
prepare(qeager, inplace=True)
else:
qeager.train()
qeager.fuse_model()
prepare_qat(qeager, inplace=True)
# calibration
if mode == 'ddp':
mp.spawn(run_ddp,
args=(world_size, qeager),
nprocs=world_size,
join=True)
elif mode == 'qat':
assert qeager.training, 'qeager should be in training mode for qat'
optimizer = torch.optim.SGD(qeager.parameters(), lr=0.0001)
train_one_epoch(qeager, criterion, optimizer, [(input_value, output_value)], torch.device('cpu'), 1)
else:
for i in range(10):
qeager(input_value)
# print('ref after observation:', qeager)
convert(qeager, inplace=True)
qeager.eval()
# print('ref after quantization:', qeager)
qeager_out = qeager(input_value)
qeager_script = torch.jit.script(qeager)
qscript_out = qeager_script(input_value)
if is_not_tuple_out:
diff_from_eager[mode][name] = (qeager_out - qgraph_out).abs().max()
if check_with_eager:
self.assertEqual(diff_from_eager[mode][name], 0,
'Result of graph mode quantization and ' +
'eager mode quantization on model: ' + name +
' should match. Mode: ' + mode +
' diff:' + str(diff_from_eager[mode][name]))
def _test_building_block(self, quant_type, BB):
eager = BB().float()
graph = copy.deepcopy(eager)
if quant_type == QuantType.STATIC:
qconfig = default_qconfig
eager_prepare = prepare
graph_prepare = prepare_fx
eager.eval()
graph.eval()
calibrate_or_train = test_only_eval_fn
data = self.img_data_2d
is_qat = False
else:
assert quant_type == QuantType.QAT
qconfig = default_qat_qconfig
eager_prepare = prepare_qat
graph_prepare = prepare_qat_fx
eager.train()
graph.train()
calibrate_or_train = test_only_train_fn
data = self.img_data_2d_train
is_qat = True
if hasattr(eager, "fuse_model"):
eager.fuse_model()
eager = QuantWrapper(eager)
eager.qconfig = qconfig
eager = eager_prepare(eager)
qconfig_dict = {"": qconfig}
graph = graph_prepare(graph, qconfig_dict, example_inputs=(data[0][0],))
eager_out = eager(data[0][0])
graph_out = graph(data[0][0])
# Eager Mode and FX Graph Mode QAT now differ in numerics both
# in Post Training and QAT because FX Graph Mode uses same fake_quant instances
# for input and output of CopyNode
# self.assertEqual(eager_out, graph_out)
calibrate_or_train(eager, data)
calibrate_or_train(graph, data)
eager = convert(eager)
graph = convert_fx(graph)
eager_out = eager(data[0][0])
graph_out = graph(data[0][0])
@override_qengines
def test_resnet_base(self):
models = [ResNetBase]
options = itertools.product(self.static_quant_types, models)
for quant_type, M in options:
self._test_building_block(quant_type, M)
@skip_if_no_torchvision
@skipIfNoFBGEMM
@unittest.skip("skip for now since tbb failed")
def test_torchvision(self):
from torchvision import models
from torchvision.models import quantization as quantized_models
from torchvision.models.quantization.utils import _replace_relu
def get_available_classification_models(models):
return [k for k, v in models.__dict__.items() if callable(v) and k[0].lower() == k[0] and k[0] != "_"]
model_list = get_available_classification_models(models)
quantized_model_list = get_available_classification_models(quantized_models)
quantized_model_list = set(quantized_model_list)
# test eager and graph consistency
model_list = quantized_model_list
# mobilenet/inception_v3/googlenet qat is not working due to AdaptiveAveragePool qat
# we might observe the output of AdaptiveAveragePool in the future
# and re-enable the test
fx_eager_not_matching = [
("mobilenet_v2", "qat"),
("inception_v3", "qat"),
("googlenet", "qat")
] # because relu6 is replaced as relu in mobilenetv2
diff_of_quant = {}
diff_from_eager = {}
modes = ['static', 'qat']
options = itertools.product(modes, model_list)
for mode, name in options:
pretrained = name in quantized_model_list # load pretrained model to compare with quantized model
kwargs = {}
# turn off transform input for inception_v3 since
# it's not quantized in eager mode and in fx graph
# mode we can't skip quantizing a method right now
# (might be supported in the future)
if name in ["inception_v3", "googlenet"]:
kwargs["transform_input"] = False
eager_quantizable_model = None
if name in quantized_model_list:
eager_quantizable_model = quantized_models.__dict__[name](pretrained=False, quantize=False, **kwargs).eval().float()
# compare with eager mode quantized model when it is available
pretrained = eager_quantizable_model is not None
model = models.__dict__[name](pretrained=pretrained, **kwargs).eval().float()
if name == "mobilenet_v2":
_replace_relu(model)
# disable aux logits
if hasattr(model, "aux_logits"):
model.aux_logits = False
model.AuxLogits = None
if eager_quantizable_model:
eager_quantizable_model.aux_logits = False
eager_quantizable_model.AuxLogits = None
check_with_eager = (name, mode) not in fx_eager_not_matching
self._test_model_impl(
mode, name, model, eager_quantizable_model,
check_with_eager,
diff_of_quant, diff_from_eager)
def print_diffs(diffs):
for mode, diffs_for_mode in diffs.items():
print('mode:', mode)
for name, diff in diffs_for_mode.items():
print(name, ':', diff)
# print('differences between float and quantized')
# print_diffs(diff_of_quant)
# print('----------------------')
# print('differences between graph mode and eager mode')
# print_diffs(diff_from_eager)
# print('----------------------')
@skip_if_no_torchvision
@skipIfNoFBGEMM
@unittest.skip("TODO: Test is always failing - https://github.com/pytorch/pytorch/issues/54979")
def test_resnet18_ddp(self):
from torchvision import models
from torchvision.models import quantization as quantized_models
eager_quantizable_model = quantized_models.__dict__[name](pretrained=False, quantize=False).eval().float()
model = models.__dict__[name](pretrained=False).eval().float()
self._test_model_impl(
'ddp', 'resnet18', model, eager_quantizable_model)
@override_qengines
def test_qat_embeddingbag_linear(self):
for device in get_supported_device_types():
class EmbeddingBagLinear(torch.nn.Module):
def __init__(self):
super().__init__()
self.emb = torch.nn.EmbeddingBag(num_embeddings=10, embedding_dim=12, mode='sum')
self.linear = torch.nn.Linear(12, 1).to(dtype=torch.float)
def forward(self, input: torch.Tensor, offsets: Optional[torch.Tensor] = None,
per_sample_weights: Optional[torch.Tensor] = None):
x = self.emb(input, offsets, per_sample_weights)
x = self.linear(x)
return x
qengine = torch.backends.quantized.engine
qconfig_dict = QConfigMapping() \
.set_global(get_default_qat_qconfig(qengine)) \
.set_object_type(torch.nn.EmbeddingBag, default_embedding_qat_qconfig)
train_indices = [[torch.randint(0, 10, (12, 12)), torch.randn((12, 1))] for _ in range(2)]
eval_output = [[torch.randint(0, 10, (12, 1))]]
model = EmbeddingBagLinear().train()
prepared_fx_model = prepare_qat_fx(model, qconfig_dict, example_inputs=(train_indices[0][0],))
test_only_train_fn(prepared_fx_model, train_indices)
quant_model = convert_fx(prepared_fx_model,
qconfig_mapping=qconfig_dict)
def checkQuantized(model):
# Make sure EmbeddingBag is now a quantized EmbeddingBag.
self.assertTrue(type(model.emb), nn.quantized.EmbeddingBag)
# Also test that Linear has been quantized.
self.assertTrue(type(model.linear), nnq.Linear)
test_only_eval_fn(model, eval_output)
self.checkScriptable(model, eval_output)
self.checkNoQconfig(model)
checkQuantized(quant_model)
@override_qengines
def test_qat_embedding_linear(self):
for device in get_supported_device_types():
class EmbeddingLinear(torch.nn.Module):
def __init__(self):
super().__init__()
self.emb = torch.nn.Embedding(num_embeddings=10, embedding_dim=12)
self.linear = torch.nn.Linear(12, 1).to(dtype=torch.float)
def forward(self, input: torch.Tensor):
x = torch.sum(self.emb(input), dim=1)
x = self.linear(x)
return x
qengine = torch.backends.quantized.engine
qconfig_dict = {"": get_default_qat_qconfig(qengine),
"object_type": [(torch.nn.Embedding, default_embedding_qat_qconfig)]}
train_indices = [[torch.randint(0, 10, (12, 12)), torch.randn((12, 1))] for _ in range(2)]
eval_output = [[torch.randint(0, 10, (12, 1))]]
model = EmbeddingLinear().train()
prepared_fx_model = prepare_qat_fx(model, qconfig_dict, example_inputs=(train_indices[0][0],))
test_only_train_fn(prepared_fx_model, train_indices)
quant_model = convert_fx(prepared_fx_model,
qconfig_mapping=qconfig_dict)
def checkQuantized(model):
# Make sure EmbeddingBag is now a quantized EmbeddingBag.
self.assertTrue(type(model.emb), nn.quantized.Embedding)
# Also test that Linear has been quantized.
self.assertTrue(type(model.linear), nnq.Linear)
test_only_eval_fn(model, eval_output)
self.checkScriptable(model, eval_output)
self.checkNoQconfig(model)
checkQuantized(quant_model)
@given(
device=st.sampled_from(
["cpu", "cuda"] if torch.cuda.is_available() else ["cpu"]
)
)
@settings(deadline=None)
@override_qengines
def test_qat_functional_linear(self, device):
if torch.backends.quantized.engine not in ('fbgemm', 'qnnpack'):
return
class Linear(torch.nn.Module):
def __init__(self):
super().__init__()
self.w = torch.ones(5, 5)
self.b = torch.zeros(5)
def forward(self, x):
return torch.nn.functional.linear(x, self.w, self.b)
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.mods1 = torch.nn.Sequential(Linear(), Linear())
self.mods2 = Linear()
def forward(self, x):
x = self.mods1(x)
x = self.mods2(x)
return x
model = M().train()
ref_fake_quant = FakeQuantize.with_args(
observer=MovingAverageMinMaxObserver,
quant_min=0,
quant_max=255,
dtype=torch.quint8,
reduce_range=False,
)
ref_weight_fake_quant = FakeQuantize.with_args(
observer=MovingAverageMinMaxObserver,
quant_min=-128,
quant_max=127,
dtype=torch.qint8,
reduce_range=False,
)
ref_qat_qconfig = QConfig(
activation=ref_fake_quant, weight=ref_weight_fake_quant
)
qconfig_dict = {"": ref_qat_qconfig}
example_inputs = (torch.randn(1, 5),)
prepared_ref = prepare_qat_fx(model, qconfig_dict, example_inputs=example_inputs)
custom_fake_quant = FusedMovingAvgObsFakeQuantize.with_args(
observer=MovingAverageMinMaxObserver,
quant_min=0,
quant_max=255,
dtype=torch.quint8,
reduce_range=False,
)
custom_weight_fake_quant = FusedMovingAvgObsFakeQuantize.with_args(
observer=MovingAverageMinMaxObserver,
quant_min=-128,
quant_max=127,
dtype=torch.qint8,
reduce_range=False,
)
custom_qconfig = QConfig(
activation=custom_fake_quant, weight=custom_weight_fake_quant
)
custom_qconfig_dict = {"": custom_qconfig}
prepared = prepare_qat_fx(model, custom_qconfig_dict, example_inputs=example_inputs)
prepared.to(device)
prepared_ref.to(device)
prepared.apply(torch.ao.quantization.disable_fake_quant)
prepared.apply(torch.ao.quantization.disable_observer)
prepared_ref.apply(torch.ao.quantization.disable_fake_quant)
prepared_ref.apply(torch.ao.quantization.disable_observer)
inp = torch.randn(5, 5, device=device, requires_grad=True)
for i in range(10):
if i == 2:
prepared.apply(torch.ao.quantization.enable_observer)
prepared_ref.apply(torch.ao.quantization.enable_observer)
if i == 4:
prepared.apply(torch.ao.quantization.enable_fake_quant)
prepared_ref.apply(torch.ao.quantization.enable_fake_quant)
inp = torch.randn(5, 5, device=device, requires_grad=True)
out_ref = prepared_ref(inp)
out = prepared(inp)
torch.testing.assert_close(out, out_ref)
# try backward pass
labels = torch.randn(5, 5, device=device)
loss = (out - labels).sum()
grad = torch.autograd.grad(loss, [inp])
loss_ref = (out_ref - labels).sum()
grad_ref = torch.autograd.grad(loss_ref, [inp])
torch.testing.assert_close(grad[0], grad_ref[0])
if 'fbgemm' in torch.backends.quantized.supported_engines:
# During the lowering step in convert, fold_weight calls quantized::linear_prepack
# which doesn't support QuantizedCuda backend
prepared.cpu()
prepared_ref.cpu()
converted = convert_fx(prepared)
converted_ref = convert_fx(prepared_ref)
inp = torch.rand(5, 5)
out = converted(inp)
out_ref = converted_ref(inp)
torch.testing.assert_close(out, out_ref)
if __name__ == '__main__':
raise RuntimeError("This test file is not meant to be run directly, use:\n\n"
"\tpython test/test_quantization.py TESTNAME\n\n"
"instead.")