diff --git a/TrainingExtensions/onnx/src/python/aimet_onnx/adaround/adaround_optimizer.py b/TrainingExtensions/onnx/src/python/aimet_onnx/adaround/adaround_optimizer.py index d75cac0f600..fda62d34688 100644 --- a/TrainingExtensions/onnx/src/python/aimet_onnx/adaround/adaround_optimizer.py +++ b/TrainingExtensions/onnx/src/python/aimet_onnx/adaround/adaround_optimizer.py @@ -105,6 +105,7 @@ def adaround_module(cls, module: ModuleInfo, quantized_input_name: str, # After optimization, set the optimized layer's rounding mode to "Hard rounding" param_to_adaround_tensor_quantizer[module.params['weight'].name].use_soft_rounding = False + # pylint: disable=too-many-statements @classmethod def _optimize_rounding(cls, module: ModuleInfo, quantized_input_name, orig_model: ModelProto, quant_model: QuantizationSimModel, @@ -150,9 +151,10 @@ def _optimize_rounding(cls, module: ModuleInfo, quantized_input_name, out_data_torch.shape) attributes = read_attributes_for_op(module) - if len(attributes['pads']) > 2: - logger.info("Skipping the Convolution layer because padding size of 4 is not supported for optimization") - return + if 'pads' in attributes: + if len(attributes['pads']) > 2: + logger.info("Skipping the Convolution layer because padding size greater than 2 is not supported for optimization") + return if use_cache_acts_data and AdaroundOptimizer.enable_caching_acts_data(): logger.debug("Caching intermediate activations data for optimization.") diff --git a/TrainingExtensions/onnx/src/python/aimet_onnx/meta/connectedgraph.py b/TrainingExtensions/onnx/src/python/aimet_onnx/meta/connectedgraph.py index aa2f079f998..86cc9ada373 100644 --- a/TrainingExtensions/onnx/src/python/aimet_onnx/meta/connectedgraph.py +++ b/TrainingExtensions/onnx/src/python/aimet_onnx/meta/connectedgraph.py @@ -148,15 +148,19 @@ def check_if_node_has_predecessor(node): for node_output in node.output: output_names[node_output] = node + # Capture constant tensors associated to a node that has only contant tensor inputs and are not in the form of a constant node. for node in self.model.graph.node: - for input_name in node.input: - if node.op_type not in OPS_WITH_PARAMS and not check_if_node_has_predecessor(node) and input_name not in output_names: - input_tensors_names.append(input_name) + if node.op_type != 'Identity' and node.op_type not in OPS_WITH_PARAMS and not check_if_node_has_predecessor(node): + for input_name in node.input: + if input_name not in output_names: + input_tensors_names.append(input_name) + # Capture model input tensors. for tensor in self.model.graph.input: if tensor.name not in input_tensors_names and tensor.name in self._input_to_node: input_tensors_names.append(tensor.name) + # Capture nodes having all the inputs as constant tensors and these constants are coming from a constant node. input_ops = [] for node in self.model.graph.node: flag = True @@ -168,13 +172,14 @@ def check_if_node_has_predecessor(node): else: flag = False break - if flag: + if flag and node not in input_ops: input_ops.append(node) for input_tensor_name in input_tensors_names: if input_tensor_name in self._input_to_node: for node in self._input_to_node[input_tensor_name]: - input_ops.append(node) + if node not in input_ops: + input_ops.append(node) return input_ops @@ -560,18 +565,22 @@ def create_batchnorm_params(my_op: Op): op = my_op.get_module() gamma_tensor = ParamUtils.get_param(self.model, op, WEIGHT_INDEX) - create_and_connect_product(gamma_tensor.name, gamma_tensor.dims, my_op, gamma_tensor, 'weight') + if gamma_tensor: + create_and_connect_product(gamma_tensor.name, gamma_tensor.dims, my_op, gamma_tensor, 'weight') beta_tensor = ParamUtils.get_param(self.model, op, BIAS_INDEX) - create_and_connect_product(beta_tensor.name, beta_tensor.dims, my_op, beta_tensor, 'bias') + if beta_tensor: + create_and_connect_product(beta_tensor.name, beta_tensor.dims, my_op, beta_tensor, 'bias') moving_mean_tensor = ParamUtils.get_param(self.model, op, RUNNING_MEAN_INDEX) - create_and_connect_product(moving_mean_tensor.name, moving_mean_tensor.dims, my_op, - moving_mean_tensor, None) + if moving_mean_tensor: + create_and_connect_product(moving_mean_tensor.name, moving_mean_tensor.dims, my_op, + moving_mean_tensor, None) moving_variance_tensor = ParamUtils.get_param(self.model, op, RUNNING_VAR_INDEX) - create_and_connect_product(moving_variance_tensor.name, moving_variance_tensor.dims, my_op, - moving_variance_tensor, None) + if moving_variance_tensor: + create_and_connect_product(moving_variance_tensor.name, moving_variance_tensor.dims, my_op, + moving_variance_tensor, None) def handle_default(my_op: Op): """ Handler for other modules """ diff --git a/TrainingExtensions/onnx/src/python/aimet_onnx/utils.py b/TrainingExtensions/onnx/src/python/aimet_onnx/utils.py index 54025b1be73..970d124a5ec 100644 --- a/TrainingExtensions/onnx/src/python/aimet_onnx/utils.py +++ b/TrainingExtensions/onnx/src/python/aimet_onnx/utils.py @@ -98,7 +98,7 @@ def remove_node(node: ModelProto, onnx_graph: onnx.GraphProto): other_node.input[idx] = node.input[0] # Check if removed node output is an output of the graph for outputs in onnx_graph.output: - if outputs.name in node.output[0] and other_node.output[0] == node.input[0]: + if outputs.name == node.output[0] and other_node.output[0] == node.input[0]: other_node.output[0] = outputs.name inits_to_remove = [] # Remove the node's initializers diff --git a/TrainingExtensions/onnx/test/python/test_adaround_optimizer.py b/TrainingExtensions/onnx/test/python/test_adaround_optimizer.py index 89eaedfb2d9..4e56f35d4c0 100644 --- a/TrainingExtensions/onnx/test/python/test_adaround_optimizer.py +++ b/TrainingExtensions/onnx/test/python/test_adaround_optimizer.py @@ -55,6 +55,7 @@ class TestAdaroundOptimizer: Test functions in utils """ + @pytest.mark.skipif(not torch.cuda.is_available(), reason="This unit-test is meant to be run on GPU") @pytest.mark.parametrize("warm_start", [1.0, 0.2]) def test_optimize_rounding(self, warm_start): if version.parse(torch.__version__) >= version.parse("1.13"): diff --git a/TrainingExtensions/onnx/test/python/test_adaround_weight.py b/TrainingExtensions/onnx/test/python/test_adaround_weight.py index 7c5af66cd2b..414c7cc5c56 100644 --- a/TrainingExtensions/onnx/test/python/test_adaround_weight.py +++ b/TrainingExtensions/onnx/test/python/test_adaround_weight.py @@ -41,6 +41,7 @@ import numpy as np import torch from onnxruntime import SessionOptions, GraphOptimizationLevel, InferenceSession +import pytest from aimet_onnx.adaround.adaround_weight import Adaround, AdaroundParameters import models.models_for_tests as test_models @@ -50,6 +51,7 @@ class TestAdaround: AdaRound Weights Unit Test Cases """ + @pytest.mark.skipif(not torch.cuda.is_available(), reason="This unit-test is meant to be run on GPU") def test_apply_adaround(self): if version.parse(torch.__version__) >= version.parse("1.13"): np.random.seed(0) diff --git a/TrainingExtensions/onnx/test/python/test_bn_fold.py b/TrainingExtensions/onnx/test/python/test_bn_fold.py index c67968f6a7a..4790fd5773a 100644 --- a/TrainingExtensions/onnx/test/python/test_bn_fold.py +++ b/TrainingExtensions/onnx/test/python/test_bn_fold.py @@ -43,6 +43,7 @@ import numpy as np import torchvision import pytest +from packaging import version import torch @@ -80,32 +81,33 @@ class TestBatchNormFold: """ Test methods for BatchNormFold""" def test_find_batch_norms_to_fold(self): - model = MyModel().eval() - initialize_bn_params(model) - - input_shape = (2, 10, 24, 24) - x = torch.randn(*input_shape, requires_grad=True) - - # Export the model - torch.onnx.export(model, # model being run - x, # model input (or a tuple for multiple inputs) - "./model_single_residual.onnx", - # where to save the model (can be a file or file-like object), - training=torch.onnx.TrainingMode.TRAINING, - export_params=True, # store the trained parameter weights inside the model file - opset_version=12, # the ONNX version to export the model to - do_constant_folding=False, # whether to execute constant folding for optimization - input_names=['input'], # the model's input names - output_names=['output']) - model = ONNXModel(load_model('./model_single_residual.onnx')) - - connected_graph = ConnectedGraph(model) - bn_info = _find_conv_bn_pairs(connected_graph) - conv1 = connected_graph.get_op_from_module_name('Conv_0') - conv3 = connected_graph.get_op_from_module_name('Conv_6') - assert len(bn_info.keys()) == 2 - assert connected_graph.get_op_from_module_name('BatchNormalization_1') == bn_info[conv1].output_bn - assert connected_graph.get_op_from_module_name('BatchNormalization_5') == bn_info[conv3].input_bn + if version.parse(torch.__version__) >= version.parse("1.13"): + model = MyModel().eval() + initialize_bn_params(model) + + input_shape = (2, 10, 24, 24) + x = torch.randn(*input_shape, requires_grad=True) + + # Export the model + torch.onnx.export(model, # model being run + x, # model input (or a tuple for multiple inputs) + "./model_single_residual.onnx", + # where to save the model (can be a file or file-like object), + training=torch.onnx.TrainingMode.TRAINING, + export_params=True, # store the trained parameter weights inside the model file + opset_version=12, # the ONNX version to export the model to + do_constant_folding=False, # whether to execute constant folding for optimization + input_names=['input'], # the model's input names + output_names=['output']) + model = ONNXModel(load_model('./model_single_residual.onnx')) + + connected_graph = ConnectedGraph(model) + bn_info = _find_conv_bn_pairs(connected_graph) + conv1 = connected_graph.get_op_from_module_name('/conv1/Conv') + conv3 = connected_graph.get_op_from_module_name('/conv3/Conv') + assert len(bn_info.keys()) == 2 + assert connected_graph.get_op_from_module_name('/bn1/BatchNormalization') == bn_info[conv1].output_bn + assert connected_graph.get_op_from_module_name('/bn2/BatchNormalization') == bn_info[conv3].input_bn def test_find_bn_before_linear(self): x = torch.randn((32, 10)) @@ -117,48 +119,52 @@ def test_find_bn_before_linear(self): assert 'MatMul' in list(bn_info.keys())[0].name def test_find_bn_before_flatten(self): - x = torch.randn((2, 10, 24, 24)) - model = BNBeforeFlattenLinear() - model = _convert_to_onnx_no_fold(model, x) - conn_graph = ConnectedGraph(model) - bn_info = _find_conv_bn_pairs(conn_graph) - linear_layer = conn_graph.get_op_from_module_name('MatMul_5') - assert len(bn_info.keys()) == 1 - assert linear_layer in bn_info.keys() - assert bn_info[linear_layer].input_bn == conn_graph.get_op_from_module_name('BatchNormalization_2') + if version.parse(torch.__version__) >= version.parse("1.13"): + x = torch.randn((2, 10, 24, 24)) + model = BNBeforeFlattenLinear() + model = _convert_to_onnx_no_fold(model, x) + conn_graph = ConnectedGraph(model) + bn_info = _find_conv_bn_pairs(conn_graph) + linear_layer = conn_graph.get_op_from_module_name('/fc2/MatMul') + assert len(bn_info.keys()) == 1 + assert linear_layer in bn_info.keys() + assert bn_info[linear_layer].input_bn == conn_graph.get_op_from_module_name('/bn1/BatchNormalization') def test_find_bn_after_linear(self): - x = torch.randn((32, 10)) - model = BNAfterLinear(bias=True) - model = _convert_to_onnx_no_fold(model, x) - conn_graph = ConnectedGraph(model) - bn_info = _find_conv_bn_pairs(conn_graph) - linear_layer = conn_graph.get_op_from_module_name('Gemm_0') - assert len(bn_info.keys()) == 1 - assert linear_layer in bn_info.keys() - assert bn_info[linear_layer].output_bn == conn_graph.get_op_from_module_name('BatchNormalization_1') + if version.parse(torch.__version__) >= version.parse("1.13"): + x = torch.randn((32, 10)) + model = BNAfterLinear(bias=True) + model = _convert_to_onnx_no_fold(model, x) + conn_graph = ConnectedGraph(model) + bn_info = _find_conv_bn_pairs(conn_graph) + linear_layer = conn_graph.get_op_from_module_name('/fc1/Gemm') + assert len(bn_info.keys()) == 1 + assert linear_layer in bn_info.keys() + assert bn_info[linear_layer].output_bn == conn_graph.get_op_from_module_name('/bn1/BatchNormalization') def test_find_bn_after_convtranspose(self): - x = torch.randn((2, 10, 24, 24)) - model = BNAfterConvTranspose() - model = _convert_to_onnx_no_fold(model, x) - conn_graph = ConnectedGraph(model) - bn_info = _find_conv_bn_pairs(conn_graph) - conv_layer = conn_graph.get_op_from_module_name('ConvTranspose_0') - assert len(bn_info.keys()) == 1 - assert conv_layer in bn_info.keys() - assert bn_info[conv_layer].output_bn == conn_graph.get_op_from_module_name('BatchNormalization_1') + if version.parse(torch.__version__) >= version.parse("1.13"): + x = torch.randn((2, 10, 24, 24)) + model = BNAfterConvTranspose() + model = _convert_to_onnx_no_fold(model, x) + conn_graph = ConnectedGraph(model) + bn_info = _find_conv_bn_pairs(conn_graph) + conv_layer = conn_graph.get_op_from_module_name('/conv1/ConvTranspose') + assert len(bn_info.keys()) == 1 + assert conv_layer in bn_info.keys() + assert bn_info[conv_layer].output_bn == conn_graph.get_op_from_module_name('/bn1/BatchNormalization') def test_find_bn_after_conv1d(self): - x = torch.randn((2, 10, 24)) - model = BNAfterConv1d() - model = _convert_to_onnx_no_fold(model, x) - conn_graph = ConnectedGraph(model) - bn_info = _find_conv_bn_pairs(conn_graph) - conv_layer = conn_graph.get_op_from_module_name('Conv_0') - assert len(bn_info.keys()) == 1 - assert conv_layer in bn_info.keys() - assert bn_info[conv_layer].output_bn == conn_graph.get_op_from_module_name('BatchNormalization_1') + if version.parse(torch.__version__) >= version.parse("1.13"): + x = torch.randn((2, 10, 24)) + model = BNAfterConv1d() + model = _convert_to_onnx_no_fold(model, x) + conn_graph = ConnectedGraph(model) + bn_info = _find_conv_bn_pairs(conn_graph) + conv_layer = conn_graph.get_op_from_module_name('/conv1/Conv') + assert len(bn_info.keys()) == 1 + assert conv_layer in bn_info.keys() + assert bn_info[conv_layer].output_bn == conn_graph.get_op_from_module_name('/bn1/BatchNormalization') def test_filter_bn_before_conv_transpose(self): x = torch.randn((2, 10, 24, 24)) @@ -235,48 +241,51 @@ def test_filter_bn_after_conv(self): assert not bn_conv def test_filter_bn_before_flatten(self): - x = torch.randn((2, 10, 24, 24)) - model = BNBeforeFlattenLinear(bias=True) - model = _convert_to_onnx_no_fold(model, x) - conn_graph = ConnectedGraph(model) - conv_bn, bn_conv = find_all_batch_norms_to_fold(conn_graph) - linear_layer = conn_graph.get_op_from_module_name('Gemm_4') - assert len(bn_conv) == 1 - assert linear_layer.get_module() == bn_conv[0][1] + if version.parse(torch.__version__) >= version.parse("1.13"): + x = torch.randn((2, 10, 24, 24)) + model = BNBeforeFlattenLinear(bias=True) + model = _convert_to_onnx_no_fold(model, x) + conn_graph = ConnectedGraph(model) + conv_bn, bn_conv = find_all_batch_norms_to_fold(conn_graph) + linear_layer = conn_graph.get_op_from_module_name('/fc2/Gemm') + assert len(bn_conv) == 1 + assert linear_layer.get_module() == bn_conv[0][1] def test_fold_bn_before_flatten_no_bias(self): - torch.manual_seed(10) - torch_model = BNBeforeFlattenLinear() - torch_model.eval() - initialize_bn_params(torch_model) + if version.parse(torch.__version__) >= version.parse("1.13"): + torch.manual_seed(10) + torch_model = BNBeforeFlattenLinear() + torch_model.eval() + initialize_bn_params(torch_model) - input_shape = (2, 10, 24, 24) - test_data = np.random.randn(*input_shape).astype(np.float32) + input_shape = (2, 10, 24, 24) + test_data = np.random.randn(*input_shape).astype(np.float32) - model = _convert_to_onnx(torch_model, torch.randn(input_shape)) - layers_orig = len(model.graph().node) - baseline_output, folded_output, pairs = get_outputs_after_fold(model, test_data) + model = _convert_to_onnx(torch_model, torch.randn(input_shape)) + layers_orig = len(model.graph().node) + baseline_output, folded_output, pairs = get_outputs_after_fold(model, test_data) - assert pairs[0][0].name == "Gemm_4" - assert len(model.graph().node) == layers_orig - 2 - assert np.allclose(baseline_output[0], folded_output[0], rtol=1e-2, atol=1e-6) + assert pairs[0][0].name == "/fc2/Gemm" + assert len(model.graph().node) == layers_orig - 2 + assert np.allclose(baseline_output[0], folded_output[0], rtol=1e-2, atol=1e-6) def test_fold_bn_before_flatten_no_bias_with_transpose(self): - torch.manual_seed(10) - torch_model = BNBeforeFlattenLinear() - torch_model.eval() - initialize_bn_params(torch_model) + if version.parse(torch.__version__) >= version.parse("1.13"): + torch.manual_seed(10) + torch_model = BNBeforeFlattenLinear() + torch_model.eval() + initialize_bn_params(torch_model) - input_shape = (2, 10, 24, 24) - test_data = np.random.randn(*input_shape).astype(np.float32) + input_shape = (2, 10, 24, 24) + test_data = np.random.randn(*input_shape).astype(np.float32) - model = _convert_to_onnx_no_fold(torch_model, torch.randn(input_shape)) - layers_orig = len(model.graph().node) - baseline_output, folded_output, pairs = get_outputs_after_fold(model, test_data) + model = _convert_to_onnx_no_fold(torch_model, torch.randn(input_shape)) + layers_orig = len(model.graph().node) + baseline_output, folded_output, pairs = get_outputs_after_fold(model, test_data) - assert pairs[0][0].name == "Gemm_5" - assert len(model.graph().node) == layers_orig - 2 - assert np.allclose(baseline_output[0], folded_output[0], rtol=1e-2, atol=1e-6) + assert pairs[0][0].name == "/fc2/Gemm" + assert len(model.graph().node) == layers_orig - 2 + assert np.allclose(baseline_output[0], folded_output[0], rtol=1e-2, atol=1e-6) def test_fold_resnet18(self): torch.manual_seed(10) @@ -296,21 +305,22 @@ def test_fold_resnet18(self): @pytest.mark.parametrize("bias", [True, False]) def test_fold_bn_before_conv(self, bias): - torch.manual_seed(10) - torch_model = BNBeforeConv(bias=bias) - torch_model.eval() - initialize_bn_params(torch_model) + if version.parse(torch.__version__) >= version.parse("1.13"): + torch.manual_seed(10) + torch_model = BNBeforeConv(bias=bias) + torch_model.eval() + initialize_bn_params(torch_model) - input_shape = (2, 10, 24, 24) - test_data = np.random.randn(*input_shape).astype(np.float32) + input_shape = (2, 10, 24, 24) + test_data = np.random.randn(*input_shape).astype(np.float32) - model = _convert_to_onnx_no_fold(torch_model, torch.randn(input_shape)) - layers_orig = len(model.graph().node) - baseline_output, folded_output, pairs = get_outputs_after_fold(model, test_data) + model = _convert_to_onnx_no_fold(torch_model, torch.randn(input_shape)) + layers_orig = len(model.graph().node) + baseline_output, folded_output, pairs = get_outputs_after_fold(model, test_data) - assert pairs[0][0].name == "Conv_3" - assert len(model.graph().node) == layers_orig - 1 - assert np.allclose(baseline_output[0], folded_output[0], rtol=1e-2, atol=1e-6) + assert pairs[0][0].name == "/conv2/Conv" + assert len(model.graph().node) == layers_orig - 1 + assert np.allclose(baseline_output[0], folded_output[0], rtol=1e-2, atol=1e-6) def test_fold_bn_before_conv_depthwise(self): torch.manual_seed(10) @@ -332,138 +342,146 @@ def test_fold_bn_before_conv_depthwise(self): @pytest.mark.parametrize("padding", [0, 1]) @pytest.mark.parametrize("groups", [1, 5, 20]) def test_fold_bn_after_conv_no_bias(self, bias, padding, groups): - torch.manual_seed(10) - torch_model = BNAfterConv(bias=bias, padding=padding, groups=groups) - torch_model.eval() + if version.parse(torch.__version__) >= version.parse("1.13"): + torch.manual_seed(10) + torch_model = BNAfterConv(bias=bias, padding=padding, groups=groups) + torch_model.eval() + initialize_bn_params(torch_model) - input_shape = (2, 10, 24, 24) - test_data = np.random.randn(*input_shape).astype(np.float32) + input_shape = (2, 10, 24, 24) + test_data = np.random.randn(*input_shape).astype(np.float32) - model = _convert_to_onnx_no_fold(torch_model, torch.randn(input_shape)) - layers_orig = len(model.graph().node) - baseline_output, folded_output, pairs = get_outputs_after_fold(model, test_data) + model = _convert_to_onnx_no_fold(torch_model, torch.randn(input_shape)) + layers_orig = len(model.graph().node) + baseline_output, folded_output, pairs = get_outputs_after_fold(model, test_data) - assert pairs[0][0].name == "Conv_2" - assert len(model.graph().node) == layers_orig - 1 - assert np.allclose(baseline_output[0], folded_output[0], rtol=1e-2, atol=1e-6) + assert pairs[0][0].name == "/conv2/Conv" + assert len(model.graph().node) == layers_orig - 1 + assert np.allclose(baseline_output[0], folded_output[0], rtol=1e-2, atol=1e-6) def test_fold_bn_after_transposed_conv_depthwise(self): - torch.manual_seed(10) - torch_model = BNAfterConvTranspose(groups=10) - torch_model.eval() - initialize_bn_params(torch_model) + if version.parse(torch.__version__) >= version.parse("1.13"): + torch.manual_seed(10) + torch_model = BNAfterConvTranspose(groups=10) + torch_model.eval() + initialize_bn_params(torch_model) - input_shape = (2, 10, 24, 24) - test_data = np.random.randn(*input_shape).astype(np.float32) + input_shape = (2, 10, 24, 24) + test_data = np.random.randn(*input_shape).astype(np.float32) - model = _convert_to_onnx_no_fold(torch_model, torch.randn(input_shape)) - layers_orig = len(model.graph().node) - baseline_output, folded_output, pairs = get_outputs_after_fold(model, test_data) + model = _convert_to_onnx_no_fold(torch_model, torch.randn(input_shape)) + layers_orig = len(model.graph().node) + baseline_output, folded_output, pairs = get_outputs_after_fold(model, test_data) - assert pairs[0][0].name == "ConvTranspose_0" - assert len(model.graph().node) == layers_orig - 1 - assert np.allclose(baseline_output[0], folded_output[0], rtol=1e-2, atol=1e-6) + assert pairs[0][0].name == "/conv1/ConvTranspose" + assert len(model.graph().node) == layers_orig - 1 + assert np.allclose(baseline_output[0], folded_output[0], rtol=1e-2, atol=1e-6) def test_fold_bn_after_transposed_conv1d(self): - torch.manual_seed(10) - torch_model = BNAfterConvTranspose1d() - torch_model.eval() - initialize_bn_params(torch_model) + if version.parse(torch.__version__) >= version.parse("1.13"): + torch.manual_seed(10) + torch_model = BNAfterConvTranspose1d() + torch_model.eval() + initialize_bn_params(torch_model) - input_shape = (2, 10, 24) - test_data = np.random.randn(*input_shape).astype(np.float32) + input_shape = (2, 10, 24) + test_data = np.random.randn(*input_shape).astype(np.float32) - model = _convert_to_onnx_no_fold(torch_model, torch.randn(input_shape)) - layers_orig = len(model.graph().node) - baseline_output, folded_output, pairs = get_outputs_after_fold(model, test_data) + model = _convert_to_onnx_no_fold(torch_model, torch.randn(input_shape)) + layers_orig = len(model.graph().node) + baseline_output, folded_output, pairs = get_outputs_after_fold(model, test_data) - assert pairs[0][0].name == "ConvTranspose_0" - assert len(model.graph().node) == layers_orig - 1 - assert np.allclose(baseline_output[0], folded_output[0], rtol=1e-2, atol=1e-6) + assert pairs[0][0].name == "/conv1/ConvTranspose" + assert len(model.graph().node) == layers_orig - 1 + assert np.allclose(baseline_output[0], folded_output[0], rtol=1e-2, atol=1e-6) def test_fold_bn_before_linear_layer_no_bias(self): - torch.manual_seed(10) - torch_model = BNBeforeLinear(bias=False) - torch_model.eval() - initialize_bn_params(torch_model) + if version.parse(torch.__version__) >= version.parse("1.13"): + torch.manual_seed(10) + torch_model = BNBeforeLinear(bias=False) + torch_model.eval() + initialize_bn_params(torch_model) - input_shape = (32, 10) - test_data = np.random.randn(*input_shape).astype(np.float32) + input_shape = (32, 10) + test_data = np.random.randn(*input_shape).astype(np.float32) - model = _convert_to_onnx_no_fold(torch_model, torch.randn(input_shape)) - layers_orig = len(model.graph().node) - baseline_output, folded_output, pairs = get_outputs_after_fold(model, test_data) + model = _convert_to_onnx_no_fold(torch_model, torch.randn(input_shape)) + layers_orig = len(model.graph().node) + baseline_output, folded_output, pairs = get_outputs_after_fold(model, test_data) - assert pairs[0][0].name == "Gemm_5" - assert len(model.graph().node) == layers_orig - 1 - assert np.allclose(baseline_output[0], folded_output[0], rtol=1e-2, atol=1e-6) + assert pairs[0][0].name == "/fc2/Gemm" + assert len(model.graph().node) == layers_orig - 1 + assert np.allclose(baseline_output[0], folded_output[0], rtol=1e-2, atol=1e-6) - model = _convert_to_onnx(torch_model, torch.randn(input_shape)) - layers_orig = len(model.graph().node) - baseline_output, folded_output, pairs = get_outputs_after_fold(model, test_data) + model = _convert_to_onnx(torch_model, torch.randn(input_shape)) + layers_orig = len(model.graph().node) + baseline_output, folded_output, pairs = get_outputs_after_fold(model, test_data) - assert pairs[0][0].name == "Gemm_3" - assert len(model.graph().node) == layers_orig - 1 - assert np.allclose(baseline_output[0], folded_output[0], rtol=1e-2, atol=1e-6) + assert pairs[0][0].name == "/fc2/Gemm" + assert len(model.graph().node) == layers_orig - 1 + assert np.allclose(baseline_output[0], folded_output[0], rtol=1e-2, atol=1e-6) def test_fold_bn_before_linear_layer_with_bias(self): - torch.manual_seed(10) - torch_model = BNBeforeLinear(bias=True) - torch_model.eval() - initialize_bn_params(torch_model) + if version.parse(torch.__version__) >= version.parse("1.13"): + torch.manual_seed(10) + torch_model = BNBeforeLinear(bias=True) + torch_model.eval() + initialize_bn_params(torch_model) - input_shape = (32, 10) - test_data = np.random.randn(*input_shape).astype(np.float32) + input_shape = (32, 10) + test_data = np.random.randn(*input_shape).astype(np.float32) - model = _convert_to_onnx_no_fold(torch_model, torch.randn(input_shape)) - layers_orig = len(model.graph().node) - baseline_output, folded_output, pairs = get_outputs_after_fold(model, test_data) + model = _convert_to_onnx_no_fold(torch_model, torch.randn(input_shape)) + layers_orig = len(model.graph().node) + baseline_output, folded_output, pairs = get_outputs_after_fold(model, test_data) - assert pairs[0][0].name == "Gemm_3" - assert len(model.graph().node) == layers_orig - 1 - assert np.allclose(baseline_output[0], folded_output[0], rtol=1e-2, atol=1e-6) + assert pairs[0][0].name == "/fc2/Gemm" + assert len(model.graph().node) == layers_orig - 1 + assert np.allclose(baseline_output[0], folded_output[0], rtol=1e-2, atol=1e-6) def test_fold_bn_after_linear_layer_with_bias(self): - torch.manual_seed(10) - torch_model = BNAfterLinear(bias=True) - torch_model.eval() - initialize_bn_params(torch_model) + if version.parse(torch.__version__) >= version.parse("1.13"): + torch.manual_seed(10) + torch_model = BNAfterLinear(bias=True) + torch_model.eval() + initialize_bn_params(torch_model) - input_shape = (32, 10) - test_data = np.random.randn(*input_shape).astype(np.float32) + input_shape = (32, 10) + test_data = np.random.randn(*input_shape).astype(np.float32) - model = _convert_to_onnx_no_fold(torch_model, torch.randn(input_shape)) - layers_orig = len(model.graph().node) - baseline_output, folded_output, pairs = get_outputs_after_fold(model, test_data) + model = _convert_to_onnx_no_fold(torch_model, torch.randn(input_shape)) + layers_orig = len(model.graph().node) + baseline_output, folded_output, pairs = get_outputs_after_fold(model, test_data) - assert pairs[0][0].name == "Gemm_0" - assert len(model.graph().node) == layers_orig - 1 - assert np.allclose(baseline_output[0], folded_output[0], rtol=1e-2, atol=1e-6) + assert pairs[0][0].name == "/fc1/Gemm" + assert len(model.graph().node) == layers_orig - 1 + assert np.allclose(baseline_output[0], folded_output[0], rtol=1e-2, atol=1e-6) def test_fold_bn_after_linear_layer_no_bias(self): - torch.manual_seed(10) - torch_model = BNAfterLinear(bias=False) - torch_model.eval() - initialize_bn_params(torch_model) - - input_shape = (32, 10) - test_data = np.random.randn(*input_shape).astype(np.float32) - - model = _convert_to_onnx_no_fold(torch_model, torch.randn(input_shape)) - layers_orig = len(model.graph().node) - baseline_output, folded_output, pairs = get_outputs_after_fold(model, test_data) - - assert pairs[0][0].name == "Gemm_1" - assert len(model.graph().node) == layers_orig - 1 - assert np.allclose(baseline_output[0], folded_output[0], rtol=1e-2, atol=1e-6) - - model = _convert_to_onnx(torch_model, torch.randn(input_shape)) - layers_orig = len(model.graph().node) - baseline_output, folded_output, pairs = get_outputs_after_fold(model, test_data) - - assert pairs[0][0].name == "Gemm_0" - assert len(model.graph().node) == layers_orig - 1 - assert np.allclose(baseline_output[0], folded_output[0], rtol=1e-2, atol=1e-6) + if version.parse(torch.__version__) >= version.parse("1.13"): + torch.manual_seed(10) + torch_model = BNAfterLinear(bias=False) + torch_model.eval() + initialize_bn_params(torch_model) + + input_shape = (32, 10) + test_data = np.random.randn(*input_shape).astype(np.float32) + + model = _convert_to_onnx_no_fold(torch_model, torch.randn(input_shape)) + layers_orig = len(model.graph().node) + baseline_output, folded_output, pairs = get_outputs_after_fold(model, test_data) + + assert pairs[0][0].name == "/fc1/Gemm" + assert len(model.graph().node) == layers_orig - 1 + assert np.allclose(baseline_output[0], folded_output[0], rtol=1e-2, atol=1e-6) + + model = _convert_to_onnx(torch_model, torch.randn(input_shape)) + layers_orig = len(model.graph().node) + baseline_output, folded_output, pairs = get_outputs_after_fold(model, test_data) + + assert pairs[0][0].name == "/fc1/Gemm" + assert len(model.graph().node) == layers_orig - 1 + assert np.allclose(baseline_output[0], folded_output[0], rtol=1e-2, atol=1e-6) @pytest.mark.parametrize("bias", [True, False]) def test_fold_bn_before_conv1d(self, bias): diff --git a/TrainingExtensions/onnx/test/python/test_connected_graph.py b/TrainingExtensions/onnx/test/python/test_connected_graph.py index 5b28a3c786a..21b0c2a2151 100644 --- a/TrainingExtensions/onnx/test/python/test_connected_graph.py +++ b/TrainingExtensions/onnx/test/python/test_connected_graph.py @@ -34,7 +34,10 @@ # # @@-COPYRIGHT-END-@@ # ============================================================================= +from packaging import version import torch +from packaging import version + from aimet_common.connected_graph.connectedgraph_utils import get_all_input_ops, get_all_ops_with_constant_inputs from aimet_onnx.meta.connectedgraph import ConnectedGraph from models import models_for_tests @@ -53,71 +56,72 @@ def test_simple_model(self): 'conv_w', 'conv_b', 'fc_w', 'fc_b'] == [product for product in products] def test_single_residual_model(self): - model = models_for_tests.single_residual_model() - conn_graph = ConnectedGraph(model) - assert len(conn_graph.get_all_ops()) == 21 - products = conn_graph.get_all_products() - assert len(products) == 34 - assert {'Conv_0_to_Relu_1', 'Relu_1_to_MaxPool_2'}.issubset({product for product in products}) - assert {'Gemm_21_to_output','45', '46', 'conv4.weight'}.issubset({product for product in products}) - input_ops = get_all_input_ops(conn_graph) - assert len(input_ops) == 1 - assert conn_graph._branch_count == 2 - assert conn_graph.ordered_ops[20].transposed_params == True + if version.parse(torch.__version__) >= version.parse("1.13"): + model = models_for_tests.single_residual_model() + conn_graph = ConnectedGraph(model) + assert len(conn_graph.get_all_ops()) == 16 + products = conn_graph.get_all_products() + assert len(products) == 29 + assert {'/conv1/Conv_to_/relu1/Relu', '/relu1/Relu_to_/maxpool/MaxPool'}.issubset({product for product in products}) + assert {'/fc/Gemm_to_output', 'onnx::Conv_43', 'onnx::Conv_44', 'conv4.weight'}.issubset({product for product in products}) + input_ops = get_all_input_ops(conn_graph) + assert len(input_ops) == 1 + assert conn_graph._branch_count == 1 + assert conn_graph.ordered_ops[15].transposed_params == True def test_multi_inputs_model(self): - model = models_for_tests.multi_input_model() - conn_graph = ConnectedGraph(model) - assert len(conn_graph.get_all_ops()) == 15 - - products = conn_graph.get_all_products() - assert len(products) == 27 - assert {'Conv_0_to_MaxPool_1', 'Conv_3_to_MaxPool_4', 'Conv_7_to_MaxPool_8'}.issubset( - {product for product in products}) - assert {'conv1_a.weight', 'conv1_a.bias', 'conv1_b.weight'}.issubset({product for product in products}) - input_ops = get_all_input_ops(conn_graph) - assert len(input_ops) == 2 + if version.parse(torch.__version__) >= version.parse("1.13"): + model = models_for_tests.multi_input_model() + conn_graph = ConnectedGraph(model) + assert len(conn_graph.get_all_ops()) == 15 + + products = conn_graph.get_all_products() + assert len(products) == 27 + assert {'/conv1_a/Conv_to_/maxpool1_a/MaxPool', '/conv1_b/Conv_to_/maxpool1_b/MaxPool', '/conv2/Conv_to_/maxpool2/MaxPool'}.issubset( + {product for product in products}) + assert {'conv1_a.weight', 'conv1_a.bias', 'conv1_b.weight'}.issubset({product for product in products}) + input_ops = get_all_input_ops(conn_graph) + assert len(input_ops) == 2 def test_transposed_conv_model(self): - model = models_for_tests.transposed_conv_model() - conn_graph = ConnectedGraph(model) - assert len(conn_graph.get_all_ops()) == 5 - - products = conn_graph.get_all_products() - assert len(products) == 18 - assert {'bn1.running_mean', - 'bn1.running_var', - 'bn1.weight', - 'bn2.bias', - 'bn2.running_mean', - 'bn2.running_var'}.issubset({product for product in products}) + if version.parse(torch.__version__) >= version.parse("1.13"): + model = models_for_tests.transposed_conv_model() + conn_graph = ConnectedGraph(model) + assert len(conn_graph.get_all_ops()) == 5 + + products = conn_graph.get_all_products() + assert len(products) == 12 + assert {'bn1.weight', + 'bn1.bias'}.issubset({product for product in products}) def test_concat_model(self): - model = models_for_tests.concat_model() - conn_graph = ConnectedGraph(model) - ops = conn_graph.get_all_ops() - assert len(ops) == 11 - assert len(ops['Concat_3'].inputs) == 3 - products = conn_graph.get_all_products() - assert len(products) == 22 - assert conn_graph._branch_count == 1 + if version.parse(torch.__version__) >= version.parse("1.13"): + model = models_for_tests.concat_model() + conn_graph = ConnectedGraph(model) + ops = conn_graph.get_all_ops() + assert len(ops) == 6 + assert len(ops['/Concat'].inputs) == 3 + products = conn_graph.get_all_products() + assert len(products) == 17 + assert conn_graph._branch_count == 0 def test_hierarchical_model(self): - model = models_for_tests.hierarchical_model() - conn_graph = ConnectedGraph(model) - ops = conn_graph.get_all_ops() - assert len(ops) == 95 - assert conn_graph._branch_count == 5 - ordered_ops = conn_graph.ordered_ops - name_to_index = {} - for index, op in enumerate(ordered_ops): - name_to_index[op.name] = index - - # Check in the graph that if A & B are connected and A comes before B in the graph then that should be the case - # in ordered graphs as well - assert name_to_index['Conv_0'] < name_to_index['Concat_18'] - assert name_to_index['Conv_86'] < name_to_index['branch_0'] - assert name_to_index['Conv_40'] < name_to_index['Conv_54'] + if version.parse(torch.__version__) >= version.parse("1.13"): + model = models_for_tests.hierarchical_model() + conn_graph = ConnectedGraph(model) + ops = conn_graph.get_all_ops() + assert len(ops) == 68 + assert conn_graph._branch_count == 0 + ordered_ops = conn_graph.ordered_ops + name_to_index = {} + for index, op in enumerate(ordered_ops): + name_to_index[op.name] = index + + # Check in the graph that if A & B are connected and A comes before B in the graph then that should be the case + # in ordered graphs as well + assert name_to_index['/conv1/conv/Conv'] < name_to_index['/nm1/tm1/Reshape'] + assert name_to_index['/sq/seq_list/seq_list.0/Conv'] < name_to_index['/sq/seq_list/seq_list.5/Conv'] + assert name_to_index['/conv2/conv/Conv'] < name_to_index['/nm2/tm1/conv3/Conv'] def test_matmul_layer_param_creation(self): torch.manual_seed(10) diff --git a/TrainingExtensions/onnx/test/python/test_cross_layer_equalization.py b/TrainingExtensions/onnx/test/python/test_cross_layer_equalization.py index 587b8772f72..7af0b9526ad 100644 --- a/TrainingExtensions/onnx/test/python/test_cross_layer_equalization.py +++ b/TrainingExtensions/onnx/test/python/test_cross_layer_equalization.py @@ -36,6 +36,7 @@ # ============================================================================= import numpy as np import copy +from packaging import version import onnx import torch @@ -54,46 +55,56 @@ class TestCLS: def test_graph_search_utils_single_residual_model(self): - model = models_for_tests.single_residual_model() - connected_graph = ConnectedGraph(model) - ordered_module_list = get_ordered_list_of_conv_modules(connected_graph.starting_ops) - graph_search_utils = GraphSearchUtils(connected_graph, ordered_module_list, cls_supported_layer_types, cls_supported_activation_types) - ordered_layer_groups = graph_search_utils.find_layer_groups_to_scale()[0] - ordered_layer_groups_names = [op.dotted_name for op in ordered_layer_groups] - assert ordered_layer_groups_names == ['Conv_3', 'Conv_5'] + if version.parse(torch.__version__) >= version.parse("1.13"): + model = models_for_tests.single_residual_model() + connected_graph = ConnectedGraph(model) + ordered_module_list = get_ordered_list_of_conv_modules(connected_graph.starting_ops) + graph_search_utils = GraphSearchUtils(connected_graph, ordered_module_list, cls_supported_layer_types, cls_supported_activation_types) + ordered_layer_groups = graph_search_utils.find_layer_groups_to_scale()[0] + ordered_layer_groups_names = [op.dotted_name for op in ordered_layer_groups] + assert ordered_layer_groups_names == ['/conv2/Conv', '/conv3/Conv'] def test_find_cls_sets_depthwise_model(self): - model = models_for_tests.depthwise_conv_model() - - connected_graph = ConnectedGraph(model) - ordered_module_list = get_ordered_list_of_conv_modules(connected_graph.starting_ops) - graph_search_utils = GraphSearchUtils(connected_graph, ordered_module_list, cls_supported_layer_types, - cls_supported_activation_types) - - ordered_layer_groups = graph_search_utils.find_layer_groups_to_scale()[0] - # Find cls sets from the layer groups - cls_sets = graph_search_utils.convert_layer_group_to_cls_sets(ordered_layer_groups) - cls_sets_names = [] - for cls_set in cls_sets: - cls_sets_name = tuple([op.dotted_name for op in cls_set]) - cls_sets_names.append(cls_sets_name) - assert cls_sets_names == [('Conv_0', 'Conv_2', 'Conv_4'), ('Conv_4', 'Conv_6', 'Conv_8'), ('Conv_8', 'Conv_10', 'Conv_12'), ('Conv_12', 'Conv_14', 'Conv_16'), ('Conv_16', 'Conv_18', 'Conv_20'), ('Conv_20', 'Conv_22', 'Conv_24'), ('Conv_24', 'Conv_26', 'Conv_28'), ('Conv_28', 'Conv_30', 'Conv_32')] + if version.parse(torch.__version__) >= version.parse("1.13"): + model = models_for_tests.depthwise_conv_model() + + connected_graph = ConnectedGraph(model) + ordered_module_list = get_ordered_list_of_conv_modules(connected_graph.starting_ops) + graph_search_utils = GraphSearchUtils(connected_graph, ordered_module_list, cls_supported_layer_types, + cls_supported_activation_types) + + ordered_layer_groups = graph_search_utils.find_layer_groups_to_scale()[0] + # Find cls sets from the layer groups + cls_sets = graph_search_utils.convert_layer_group_to_cls_sets(ordered_layer_groups) + cls_sets_names = [] + for cls_set in cls_sets: + cls_sets_name = tuple([op.dotted_name for op in cls_set]) + cls_sets_names.append(cls_sets_name) + assert cls_sets_names == [('/model/model.0/model.0.0/Conv', '/model/model.1/model.1.0/Conv', '/model/model.1/model.1.3/Conv'), + ('/model/model.1/model.1.3/Conv', '/model/model.2/model.2.0/Conv', '/model/model.2/model.2.3/Conv'), + ('/model/model.2/model.2.3/Conv', '/model/model.3/model.3.0/Conv', '/model/model.3/model.3.3/Conv'), + ('/model/model.3/model.3.3/Conv', '/model/model.4/model.4.0/Conv', '/model/model.4/model.4.3/Conv'), + ('/model/model.4/model.4.3/Conv', '/model/model.5/model.5.0/Conv', '/model/model.5/model.5.3/Conv'), + ('/model/model.5/model.5.3/Conv', '/model/model.6/model.6.0/Conv', '/model/model.6/model.6.3/Conv'), + ('/model/model.6/model.6.3/Conv', '/model/model.7/model.7.0/Conv', '/model/model.7/model.7.3/Conv'), + ('/model/model.7/model.7.3/Conv', '/model/model.8/model.8.0/Conv', '/model/model.8/model.8.3/Conv')] def test_find_cls_sets_resnet_model(self): - model = models_for_tests.single_residual_model() - connected_graph = ConnectedGraph(model) - ordered_module_list = get_ordered_list_of_conv_modules(connected_graph.starting_ops) - graph_search_utils = GraphSearchUtils(connected_graph, ordered_module_list, cls_supported_layer_types, - cls_supported_activation_types) - - ordered_layer_groups = graph_search_utils.find_layer_groups_to_scale()[0] - # Find cls sets from the layer groups - cls_sets = graph_search_utils.convert_layer_group_to_cls_sets(ordered_layer_groups) - cls_sets_names = [] - for cls_set in cls_sets: - cls_sets_name = tuple([op.dotted_name for op in cls_set]) - cls_sets_names.append(cls_sets_name) - assert cls_sets_names == [('Conv_3', 'Conv_5')] + if version.parse(torch.__version__) >= version.parse("1.13"): + model = models_for_tests.single_residual_model() + connected_graph = ConnectedGraph(model) + ordered_module_list = get_ordered_list_of_conv_modules(connected_graph.starting_ops) + graph_search_utils = GraphSearchUtils(connected_graph, ordered_module_list, cls_supported_layer_types, + cls_supported_activation_types) + + ordered_layer_groups = graph_search_utils.find_layer_groups_to_scale()[0] + # Find cls sets from the layer groups + cls_sets = graph_search_utils.convert_layer_group_to_cls_sets(ordered_layer_groups) + cls_sets_names = [] + for cls_set in cls_sets: + cls_sets_name = tuple([op.dotted_name for op in cls_set]) + cls_sets_names.append(cls_sets_name) + assert cls_sets_names == [('/conv2/Conv', '/conv3/Conv')] def test_scale_model_residual(self): model = models_for_tests.single_residual_model() @@ -158,6 +169,7 @@ def test_cle(self): def test_cle_conv1D_model(self): x = torch.randn((2, 10, 24)) model = models_for_tests.BNAfterConv1d() + models_for_tests.initialize_bn_params(model) model = models_for_tests._convert_to_onnx_no_fold(model, x) input_shape = (2, 10, 24) test_data = np.random.randn(*input_shape).astype(np.float32) @@ -170,6 +182,7 @@ def test_cle_conv1D_model(self): def test_cle_transpose1D_model(self): x = torch.randn((2, 10, 24)) model = models_for_tests.BNAfterConvTranspose1d() + models_for_tests.initialize_bn_params(model) model = models_for_tests._convert_to_onnx_no_fold(model, x) input_shape = (2, 10, 24) test_data = np.random.randn(*input_shape).astype(np.float32) diff --git a/TrainingExtensions/onnx/test/python/test_quantsim.py b/TrainingExtensions/onnx/test/python/test_quantsim.py index e6769f1ab8b..aea77502ecb 100644 --- a/TrainingExtensions/onnx/test/python/test_quantsim.py +++ b/TrainingExtensions/onnx/test/python/test_quantsim.py @@ -36,13 +36,14 @@ # ============================================================================= import json import os - import onnx import torch import numpy as np from onnx import load_model import onnxruntime as ort import pytest +from packaging import version + from aimet_common import libquant_info from aimet_common.defs import QuantScheme, QuantizationDataType from aimet_common.quantsim_config.utils import get_path_for_per_channel_config @@ -50,9 +51,7 @@ from aimet_onnx.qc_quantize_op import OpMode from aimet_onnx.utils import make_dummy_input from models.models_for_tests import SingleResidual - -from models.models_for_tests import build_dummy_model, single_residual_model, BNAfterConv, multi_input_with_constant_model, multi_output_model -from models.models_for_tests import custom_add_model +from models.models_for_tests import build_dummy_model, single_residual_model, BNAfterConv, multi_input_with_constant_model , multi_output_model, custom_add_model class DummyModel(SingleResidual): @@ -127,6 +126,7 @@ def test_create_quantsim_dynamic_batch_size(self): inputs = torch.randn((2, 10, 24, 24)) torch.onnx.export(model, inputs, '/tmp/dummy_model.onnx', training=torch.onnx.TrainingMode.PRESERVE, + opset_version=12, input_names=['input'], output_names=['output'], dynamic_axes={ 'input': {0: 'batch_size'}, @@ -224,30 +224,46 @@ def dummy_callback(session, args): assert param_encodings_keys == ['bitwidth', 'dtype', 'is_symmetric', 'max', 'min', 'offset', 'scale'] def test_single_residual(self): - model = single_residual_model().model - sim = QuantizationSimModel(model, use_cuda=False) - for quantizer in sim.qc_quantize_op_dict: - sim.qc_quantize_op_dict[quantizer].enabled = True - - def dummy_callback(session, args): - pass - - sim.compute_encodings(dummy_callback, None) - sim.export('/tmp/', 'quant_sim_model') - - with open('/tmp/quant_sim_model.encodings', 'rb') as json_file: - encoding_data = json.load(json_file) - activation_keys = list(encoding_data["activation_encodings"].keys()) - assert activation_keys == ['20', '21', '24', '25', '26', '28', '29', '30', '31', '33', '34', '44', '47', 'input', 'output'] - for act in activation_keys: - act_encodings_keys = list(encoding_data["activation_encodings"][act][0].keys()) - assert act_encodings_keys == ['bitwidth', 'dtype', 'is_symmetric', 'max', 'min', 'offset', 'scale'] - - param_keys = list(encoding_data['param_encodings'].keys()) - assert param_keys == ['45', '46', '48', '49', 'conv3.weight', 'conv4.bias', 'conv4.weight', 'fc.bias', 'fc.weight'] - for param in param_keys: - param_encodings_keys = list(encoding_data["param_encodings"][param][0].keys()) - assert param_encodings_keys == ['bitwidth', 'dtype', 'is_symmetric', 'max', 'min', 'offset', 'scale'] + if version.parse(torch.__version__) >= version.parse("1.13"): + model = single_residual_model().model + sim = QuantizationSimModel(model, use_cuda=False) + for quantizer in sim.qc_quantize_op_dict: + sim.qc_quantize_op_dict[quantizer].enabled = True + + def dummy_callback(session, args): + pass + + sim.compute_encodings(dummy_callback, None) + sim.export('/tmp/', 'quant_sim_model') + + with open('/tmp/quant_sim_model.encodings', 'rb') as json_file: + encoding_data = json.load(json_file) + activation_keys = list(encoding_data["activation_encodings"].keys()) + assert activation_keys == ['/Add_output_0', + '/ada/AveragePool_output_0', + '/ada/Pad_output_0', + '/avgpool/AveragePool_output_0', + '/avgpool/Pad_output_0', + '/conv1/Conv_output_0', + '/conv2/Conv_output_0', + '/conv3/Conv_output_0', + '/conv4/Conv_output_0', + '/maxpool/MaxPool_output_0', + '/relu1/Relu_output_0', + '/relu2/Relu_output_0', + '/relu3/Relu_output_0', + 'input', + 'output'] + + for act in activation_keys: + act_encodings_keys = list(encoding_data["activation_encodings"][act][0].keys()) + assert act_encodings_keys == ['bitwidth', 'dtype', 'is_symmetric', 'max', 'min', 'offset', 'scale'] + + param_keys = list(encoding_data['param_encodings'].keys()) + assert param_keys == ["conv3.weight", "conv4.bias", "conv4.weight", "fc.bias", "fc.weight", "onnx::Conv_43", "onnx::Conv_44", "onnx::Conv_46", "onnx::Conv_47"] + for param in param_keys: + param_encodings_keys = list(encoding_data["param_encodings"][param][0].keys()) + assert param_encodings_keys == ['bitwidth', 'dtype', 'is_symmetric', 'max', 'min', 'offset', 'scale'] @pytest.mark.cuda def test_compare_encodings_cpu_gpu(self): @@ -431,11 +447,12 @@ def callback(session, args): def test_model_with_constants(self): - model = multi_input_with_constant_model() + if version.parse(torch.__version__) >= version.parse("1.13"): + model = multi_input_with_constant_model() - sim = QuantizationSimModel(model) - assert sim.qc_quantize_op_dict['13'].enabled == True - assert sim.qc_quantize_op_dict['7'].enabled == True + sim = QuantizationSimModel(model) + assert sim.qc_quantize_op_dict['/add0/Constant_output_0'].enabled == True + assert sim.qc_quantize_op_dict['/add2/Constant_output_0'].enabled == True def test_multiple_output_quantsim(self): diff --git a/TrainingExtensions/onnx/test/python/test_utils.py b/TrainingExtensions/onnx/test/python/test_utils.py index 746a9247e45..869daba0389 100644 --- a/TrainingExtensions/onnx/test/python/test_utils.py +++ b/TrainingExtensions/onnx/test/python/test_utils.py @@ -35,6 +35,8 @@ # @@-COPYRIGHT-END-@@ # ============================================================================= import onnx +import torch +from packaging import version import aimet_onnx.utils as utils from aimet_onnx.utils import ParamUtils @@ -160,25 +162,24 @@ def test_get_attribute(self): assert utils.get_node_attribute(conv_layer, "kernel_shape") == [3, 3] def test_replace_relu6_with_relu(self): - model = models_for_tests.depthwise_conv_model_with_relu6() - relu6_count = 0 - original_relu_count = 0 - for node in model.model.graph.node: - if node.op_type == 'Clip': - relu6_count += 1 - if node.op_type == 'Relu': - original_relu_count += 1 - - utils.replace_relu6_with_relu(model) - - relu_count = 0 - for node in model.model.graph.node: - if node.op_type == 'Relu': - relu_count += 1 - - hard_swish_count = 1 - - assert relu_count - original_relu_count == relu6_count - hard_swish_count + if version.parse(torch.__version__) >= version.parse("1.13"): + model = models_for_tests.depthwise_conv_model_with_relu6() + relu6_count = 0 + original_relu_count = 0 + for node in model.model.graph.node: + if node.op_type == 'Clip': + relu6_count += 1 + if node.op_type == 'Relu': + original_relu_count += 1 + + utils.replace_relu6_with_relu(model) + + relu_count = 0 + for node in model.model.graph.node: + if node.op_type == 'Relu': + relu_count += 1 + + assert relu_count - original_relu_count == relu6_count def test_create_model_data_single_residual_model(self): model = models_for_tests.transposed_conv_model_without_bn()