def test_find_node_in_nx_graph_by_scope():
model = TwoConvTestModel()
nncf_model = NNCFNetwork(deepcopy(model),
input_infos=[ModelInputInfo([1, 1, 4, 4])
]) # type: NNCFNetwork
nncf_graph = nncf_model.get_original_graph()
# Valid scopes should be successfully found
valid_nncf_modules = nncf_model.get_nncf_modules()
nodes_list = list(nncf_graph._nx_graph.nodes)
for module_scope, _ in valid_nncf_modules.items():
graph_node = nncf_graph.find_node_in_nx_graph_by_scope(module_scope)
assert graph_node is not None
assert isinstance(graph_node, dict)
assert graph_node['key'] in nodes_list
fake_model = BasicConvTestModel()
fake_nncf_model = NNCFNetwork(deepcopy(fake_model),
input_infos=[ModelInputInfo([1, 1, 4, 4])])
# Not valid scopes shouldn't be found
fake_nncf_modules = fake_nncf_model.get_nncf_modules()
for module_scope, _ in fake_nncf_modules.items():
graph_node = nncf_graph.find_node_in_nx_graph_by_scope(module_scope)
assert graph_node is None
def test_activation_shape_tracing(input_shape: Tuple):
model = ModelForTest()
input_info = ModelInputInfo(input_shape)
graph_builder = GraphBuilder(create_dummy_forward_fn([
input_info,
]))
graph = graph_builder.build_graph(model)
shape1 = (input_shape[0], ModelForTest.CONV1_OUT_CHANNELS, input_shape[2],
input_shape[3])
ref_node_ids_and_output_shapes = [
# TODO: extend with checking input tensor size once proper input node marking is implemented
("0 ModelForTest/Conv2d[conv1]/conv2d", [shape1]),
("1 ModelForTest/BatchNorm2d[bn1]/batch_norm", [shape1]),
("2 ModelForTest/ReLU[relu1]/RELU", [shape1, shape1]),
("3 ModelForTest/max_pool2d",
[(shape1[0], shape1[1], shape1[2] // ModelForTest.MAXPOOL_SIZE,
shape1[3] // ModelForTest.MAXPOOL_SIZE)]),
("4 ModelForTest/ConvTranspose2d[convt1]/conv_transpose2d",
[input_shape]),
("5 ModelForTest/cat",
[(input_shape[0], ModelForTest.CONV2_IN_CHANNELS, input_shape[2],
input_shape[3])])
# TODO: extend with checking output tensor size once proper output node marking is implemented
]
for node_id, ref_output_shapes in ref_node_ids_and_output_shapes:
# pylint:disable=protected-access
output_edges = graph._get_nncf_graph_pattern_input_output([
node_id,
]).output_edges
output_shapes = [x.tensor_shape for x in output_edges]
assert output_shapes == ref_output_shapes, "Failed for {}".format(
node_id)
def test_build_graph(self, desc: ModelDesc):
net = desc.model_builder()
input_sample_sizes = desc.input_sample_sizes
if isinstance(input_sample_sizes, tuple):
input_info_list = [
ModelInputInfo(sample_size)
for sample_size in input_sample_sizes
]
else:
input_info_list = [ModelInputInfo(input_sample_sizes)]
dummy_forward_fn = desc.dummy_forward_fn
if not dummy_forward_fn:
dummy_forward_fn = create_dummy_forward_fn(input_info_list)
graph_builder = GraphBuilder(custom_forward_fn=dummy_forward_fn)
graph = graph_builder.build_graph(net)
check_graph(graph, desc.dot_filename, 'original')
def test_check_correct_modules_replacement():
model = TwoConvTestModel()
nncf_model = NNCFNetwork(TwoConvTestModel(),
input_infos=[ModelInputInfo([1, 1, 4, 4])
]) # type: NNCFNetwork
_, nncf_modules = check_correct_nncf_modules_replacement(model, nncf_model)
assert set(nncf_modules) == set(nncf_model.get_nncf_modules())
def test_operator_metatype_marking(self):
from nncf.dynamic_graph.operator_metatypes import Conv2dMetatype, BatchNormMetatype, RELUMetatype, \
MaxPool2dMetatype, \
ConvTranspose2dMetatype, DepthwiseConv2dSubtype, AddMetatype, AvgPool2dMetatype, LinearMetatype
ref_scope_vs_metatype_dict = {
"/" + MODEL_INPUT_OP_NAME + "_0": NoopMetatype,
"ModelForMetatypeTesting/NNCFConv2d[conv_regular]/conv2d_0": Conv2dMetatype,
"ModelForMetatypeTesting/BatchNorm2d[bn]/batch_norm_0": BatchNormMetatype,
"ModelForMetatypeTesting/RELU_0": RELUMetatype,
"ModelForMetatypeTesting/MaxPool2d[max_pool2d]/max_pool2d_0": MaxPool2dMetatype,
"ModelForMetatypeTesting/NNCFConvTranspose2d[conv_transpose]/conv_transpose2d_0": ConvTranspose2dMetatype,
"ModelForMetatypeTesting/NNCFConv2d[conv_depthwise]/conv2d_0": DepthwiseConv2dSubtype,
"ModelForMetatypeTesting/__iadd___0": AddMetatype,
"ModelForMetatypeTesting/AdaptiveAvgPool2d[adaptive_avg_pool]/adaptive_avg_pool2d_0": AvgPool2dMetatype,
"ModelForMetatypeTesting/NNCFLinear[linear]/linear_0": LinearMetatype
}
class ModelForMetatypeTesting(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv_regular = torch.nn.Conv2d(in_channels=3,
out_channels=16,
kernel_size=3)
self.bn = torch.nn.BatchNorm2d(num_features=16)
self.max_pool2d = torch.nn.MaxPool2d(kernel_size=2)
self.conv_transpose = torch.nn.ConvTranspose2d(in_channels=16,
out_channels=8,
kernel_size=3)
self.conv_depthwise = torch.nn.Conv2d(in_channels=8, out_channels=8,
kernel_size=5, groups=8)
self.adaptive_avg_pool = torch.nn.AdaptiveAvgPool2d(output_size=1)
self.linear = torch.nn.Linear(in_features=8, out_features=1)
def forward(self, input_):
x = self.conv_regular(input_)
x = self.bn(x)
x = torch.nn.functional.relu(x)
x.transpose_(2, 3)
x = self.max_pool2d(x)
x = self.conv_transpose(x)
x = self.conv_depthwise(x)
x += torch.ones_like(x)
x = self.adaptive_avg_pool(x)
x = self.linear(x.flatten())
return x
model = ModelForMetatypeTesting()
nncf_network = NNCFNetwork(model, [ModelInputInfo([1, 3, 300, 300])])
ip_graph = nncf_network.get_insertion_point_graph()
for node in ip_graph.nodes().values():
if node[InsertionPointGraph.NODE_TYPE_NODE_ATTR] == InsertionPointGraphNodeType.OPERATOR:
nncf_node_ref = node[InsertionPointGraph.REGULAR_NODE_REF_NODE_ATTR]
scope_str = str(nncf_node_ref[NNCFGraph.OP_EXEC_CONTEXT_NODE_ATTR].input_agnostic)
assert scope_str in ref_scope_vs_metatype_dict
ref_metatype = ref_scope_vs_metatype_dict[scope_str]
assert node[InsertionPointGraph.OPERATOR_METATYPE_NODE_ATTR] == ref_metatype
def get_all_node_names(model, input_sample_size, builder=None):
if not builder:
builder = GraphBuilder(
create_dummy_forward_fn([
ModelInputInfo(input_sample_size),
]))
graph = builder.build_graph(model)
return [
node_name.split(' ', 1)[1] for node_name in graph.get_all_node_keys()
]
def test_output_quantization(_quantize_config):
net = test_models.UNet()
ctx = reset_context('orig')
ctx = reset_context('quantized_graphs')
input_shape = (1, 3, 360, 480)
qnet = QuantizedNetwork(net, _quantize_config.quantizer, [ModelInputInfo(input_shape), ],
quantize_outputs=True)
_ = qnet(torch.zeros(*input_shape))
_ = qnet(torch.zeros(*input_shape))
check_graph(ctx.graph, 'unet_qoutput.dot', _quantize_config.graph_dir)
def test_resnet18__with_ignore(_quantize_config):
net = test_models.ResNet18()
ctx = reset_context('orig')
ctx = reset_context('quantized_graphs')
input_shape = (1, 3, 32, 32)
qnet = QuantizedNetwork(net, _quantize_config.quantizer, [ModelInputInfo(input_shape), ],
ignored_scopes=['ResNet/Sequential[layer3]'])
_ = qnet(torch.zeros(*input_shape))
_ = qnet(torch.zeros(*input_shape))
check_graph(ctx.graph, 'resnet18_ignore.dot', _quantize_config.graph_dir)
def test_resnet18__with_not_qinput(_quantize_config):
net = test_models.ResNet18()
ctx = reset_context('orig')
ctx = reset_context('quantized_graphs')
input_shape = (1, 3, 32, 32)
qnet = QuantizedNetwork(net, _quantize_config.quantizer, [ModelInputInfo(input_shape), ],
quantize_inputs=False)
_ = qnet(torch.zeros(*input_shape))
_ = qnet(torch.zeros(*input_shape))
check_graph(ctx.graph, 'resnet18_no_qinput.dot', _quantize_config.graph_dir)
def test_quantize_network(self, model_name, model_builder, forward_fn_, _quantize_config):
net = model_builder()
ctx = reset_context('orig')
ctx = reset_context('quantized_graphs')
qnet = QuantizedNetwork(net, _quantize_config.quantizer,
input_infos=[ModelInputInfo(forward_fn_.keywords["input_size_"]), ],
dummy_forward_fn=forward_fn_)
qnet.to(self.device)
forward_fn_(qnet)
forward_fn_(qnet)
check_graph(ctx.graph, model_name, _quantize_config.graph_dir)
def test_disable_shape_matching():
class MatMulModel(nn.Module):
def __init__(self):
super().__init__()
self.dummy_param = torch.nn.Parameter(torch.ones([1]))
def forward(self, inputs):
half1, half2 = torch.chunk(inputs, 2, dim=2)
return torch.bmm(half1, half2.transpose(1, 2))
model = MatMulModel()
input_shape_1 = (3, 32, 32)
input_shape_2 = (4, 64, 64)
qnet_no_shape = NNCFNetwork(
deepcopy(model),
input_infos=[
ModelInputInfo(input_shape_1),
],
scopes_without_shape_matching=['MatMulModel']) # type: NNCFNetwork
_ = qnet_no_shape(torch.zeros(*input_shape_1))
graph_1 = deepcopy(qnet_no_shape.get_graph())
_ = qnet_no_shape(torch.zeros(*input_shape_2))
graph_2 = deepcopy(qnet_no_shape.get_graph())
keys_1 = list(graph_1.get_all_node_keys())
keys_2 = list(graph_2.get_all_node_keys())
assert len(keys_1) == 2 # 1 input node + 1 operation node
assert keys_1 == keys_2
qnet = NNCFNetwork(model, input_infos=[
ModelInputInfo(input_shape_1),
]) # type: NNCFNetwork
_ = qnet(torch.zeros(*input_shape_1))
_ = qnet(torch.zeros(*input_shape_2))
# The second forward run should have led to an increase in registered node counts
# since disable_shape_matching was False and the network was run with a different
# shape of input tensor
assert qnet.get_graph().get_nodes_count() > graph_1.get_nodes_count()
def test_custom_quantizable_subgraph_patterns(_quantize_config):
net = test_models.SENet18()
ctx = reset_context('orig')
ctx = reset_context('quantized_graphs')
input_shape = (1, 3, 32, 32)
qnet = QuantizedNetwork(net, _quantize_config.quantizer, [ModelInputInfo(input_shape), ],
quantize_outputs=False,
quantizable_subgraph_patterns=(("sigmoid", "__mul__"),
("__iadd__", "batch_norm")))
_ = qnet(torch.zeros(*input_shape))
_ = qnet(torch.zeros(*input_shape))
check_graph(ctx.graph, 'senet_custom_patterns.dot', _quantize_config.graph_dir)
def get_model_and_ctrl_with_applied_hw_config_quantization(model: torch.nn.Module, hw_config_dict: dict,
should_be_quantize_inputs: bool = True):
nncf_config = get_quantization_config_without_range_init(model_size=1)
nncf_config["compression"].update({"quantize_inputs": should_be_quantize_inputs})
nncf_config["hw_config_type"] = "mock"
net = NNCFNetwork(model, input_infos=[ModelInputInfo([1, 2, 1, 1])])
hw_config = HWConfig.from_dict(hw_config_dict)
qbuilder = QuantizationBuilder(nncf_config["compression"], should_init=False)
qbuilder.quantizer_setup_type = QuantizerSetupType.PROPAGATION_BASED
qbuilder.hw_config = hw_config
net = qbuilder.apply_to(net)
ctrl = net.commit_compression_changes()
return net, ctrl
def test_custom_module_registering():
model = TwoConvTestModelWithUserModule()
nncf_model = NNCFNetwork(model, input_infos=[ModelInputInfo([1, 1, 4, 4])
]) # type: NNCFNetwork
from nncf.layers import UNWRAPPED_USER_MODULES
assert ModuleOfUser in UNWRAPPED_USER_MODULES.registry_dict.values()
# pylint: disable=protected-access
assert isinstance(nncf_model.user_module, ModuleOfUser)
assert isinstance(nncf_model.user_module, _NNCFModuleMixin)
assert type(nncf_model.user_module).__name__ == "NNCFUserModuleOfUser"
user_module_attrs = dir(nncf_model.user_module)
for attr in dir(_NNCFModuleMixin):
assert attr in user_module_attrs
def get_all_node_names(model,
input_sample_size,
graph_scope=None,
builder=None):
if graph_scope is None:
graph_scope = 'utils'
reset_context(graph_scope)
if not builder:
builder = GraphBuilder(
create_dummy_forward_fn([
ModelInputInfo(input_sample_size),
]))
graph = builder.build_graph(model, graph_scope)
return [
node_name.split(' ', 1)[1] for node_name in graph.get_all_node_keys()
]
def test_ambiguous_function():
class Model(nn.Module):
def __init__(self):
super().__init__()
self.layers = nn.ModuleList([
nn.Conv2d(1, 1, 1),
nn.Conv2d(1, 1, 1)
])
def forward(self, x):
for layer in self.layers:
x = F.relu(layer(x))
reset_context('orig')
reset_context('quantized_graphs')
mod = Model()
input_info = ModelInputInfo((1, 1, 1, 1))
QuantizedNetwork(mod, create_quantize_module,
input_infos=[input_info, ])
def create_test_quantization_env() -> QuantizationEnv:
model = BasicConvTestModel()
nncf_network = NNCFNetwork(model,
input_infos=[ModelInputInfo([1, 1, 4, 4])])
hw_config_type = HWConfigType.VPU
hw_config_path = HWConfig.get_path_to_hw_config(hw_config_type)
hw_config = HWConfig.from_json(hw_config_path)
setup = PropagationBasedQuantizerSetupGenerator(
NNCFConfig(), nncf_network, hw_config=hw_config).generate_setup()
experimental_builder = ExperimentalQuantizationBuilder(setup, {})
experimental_builder.apply_to(nncf_network)
# pylint:disable=line-too-long
experimental_ctrl = nncf_network.commit_compression_changes(
) # type: ExperimentalQuantizationController
data_loader = create_mock_dataloader({
"sample_size": [1, 1, 4, 4],
})
constraints = HardwareQuantizationConstraints()
for qid in experimental_ctrl.all_quantizations:
qconf_constraint_list = []
qconf = experimental_ctrl.all_quantizations[qid].get_current_config()
bit_set = [8, 4, 2] if 'conv' in str(qid) else [8, 4]
for bits in bit_set:
adj_qconf = deepcopy(qconf)
adj_qconf.bits = bits
qconf_constraint_list.append(adj_qconf)
constraints.add(qid, qconf_constraint_list)
return QuantizationEnv(nncf_network,
experimental_ctrl,
constraints,
data_loader,
lambda *x: 0,
hw_config_type=HWConfigType.VPU,
params=QuantizationEnvParams(
compression_ratio=0.15,
eval_subset_ratio=1.0,
skip_constraint=False,
finetune=False,
bits=[2, 4, 8],
dump_init_precision_data=False))
def test_quantize_has_proper_is_weights_flag():
class Model(nn.Module):
def __init__(self, size=1):
super().__init__()
self.size = size
self.conv = nn.Conv2d(size, size, size)
def forward(self, x):
return self.conv(x)
model = Model()
reset_context('orig')
reset_context('quantized_graphs')
input_info = ModelInputInfo((1, 1, 2, 2))
quant_model = QuantizedNetwork(model, create_quantize_module, input_infos=[input_info, ],
dummy_forward_fn=create_dummy_forward_fn([input_info, ]))
for module in quant_model.modules():
if isinstance(module, NNCFConv2d):
for op in module.pre_ops.values():
assert isinstance(op, (UpdateWeight, UpdateInputs))
assert op.operand.is_weights is isinstance(op, UpdateWeight)
for _, aq in quant_model.activation_quantizers.items():
assert aq.is_weights is False
def test_ambiguous_function():
class Model(nn.Module):
def __init__(self):
super().__init__()
self.layers = nn.ModuleList(
[nn.Conv2d(1, 1, 1), nn.Conv2d(1, 1, 1)])
def forward(self, x):
for layer in self.layers:
x = F.relu(layer(x))
mod = Model()
input_info = ModelInputInfo([1, 1, 1, 1])
graph_builder = GraphBuilder(custom_forward_fn=create_dummy_forward_fn([
input_info,
]))
graph = graph_builder.build_graph(mod)
unique_op_exec_contexts = set()
# pylint:disable=protected-access
for _, node in graph._nx_graph.nodes.items():
node_op_exec_context = node[NNCFGraph.OP_EXEC_CONTEXT_NODE_ATTR]
assert node_op_exec_context not in unique_op_exec_contexts
def test_can_quantize_free_operators(mocker):
class Model(nn.Module):
def __init__(self):
super().__init__()
self.weight = nn.Parameter(torch.ones([1]))
self.bias = nn.Parameter(torch.ones([1]))
def forward(self, x):
return F.linear(x, self.weight, self.bias)
reset_context('orig')
reset_context('quantized_graphs')
mod = Model()
input_shape = (1, 1, 1, 1)
input_info = ModelInputInfo(input_shape)
nncf_module = QuantizedNetwork(mod, create_quantize_module,
input_infos=[input_info, ])
quantizer_list = nncf_module.function_quantizers.values()
assert len(quantizer_list) == 2
for quantizer in quantizer_list:
mocker.spy(quantizer, 'quantize')
_ = nncf_module(torch.ones(input_shape))
for quantizer in quantizer_list:
assert quantizer.quantize.call_count == 1
def forward(arg1=None, arg2=None, arg3=None, arg4=None, arg5=TENSOR_DEFAULT):
pass
class InputWrappingTestStruct:
def __init__(self, input_infos, model_args, model_kwargs,
ref_wrapping_sequence):
self.input_infos = input_infos
self.model_args = model_args
self.model_kwargs = model_kwargs
self.ref_wrapping_sequence = ref_wrapping_sequence
INPUT_WRAPPING_TEST_CASES = [
InputWrappingTestStruct(
input_infos=[ModelInputInfo([1])],
model_args=(TENSOR_1, ),
model_kwargs={},
ref_wrapping_sequence=[TENSOR_1],
),
InputWrappingTestStruct(
input_infos=[ModelInputInfo([1], keyword="arg2")],
model_args=(),
model_kwargs={"arg2": TENSOR_2},
ref_wrapping_sequence=[TENSOR_2],
),
InputWrappingTestStruct(
input_infos=[
ModelInputInfo([1]),
ModelInputInfo([1]),
ModelInputInfo([1], keyword="arg3"),
def test_build_graph(self, model_name, model_builder, input_size):
net = model_builder()
graph_builder = GraphBuilder(create_dummy_forward_fn([ModelInputInfo(input_size), ]))
graph = graph_builder.build_graph(net)
check_graph(graph, model_name, 'original')
def setup(self):
self.compressed_model = NNCFNetwork(
InsertionPointTestModel(),
[ModelInputInfo([1, 1, 10, 10])]) # type: NNCFNetwork
]).output_edges
output_shapes = [x.tensor_shape for x in output_edges]
assert output_shapes == ref_output_shapes, "Failed for {}".format(
node_id)
TEST_KEYWORD_1 = "keyword1"
TEST_KEYWORD_2 = "keyword2"
INPUT_INFO_CONFIG_VS_FORWARD_ARGS = [
({
"sample_size": [2, 3, 300, 300],
"type": "float",
"filler": "zeros"
}, [
ModelInputInfo([2, 3, 300, 300],
type_str="float",
filler=ModelInputInfo.FILLER_TYPE_ZEROS)
]),
([{
"sample_size": [1, 128],
"type": "long",
"filler": "ones"
}, {
"sample_size": [1, 128],
"type": "long",
"filler": "ones"
}, {
"sample_size": [1, 128],
"type": "long",
"filler": "zeros"
}], [