From 6704f116a81acb3e9bfc469744be30ab58654324 Mon Sep 17 00:00:00 2001 From: Harsh Peswani Date: Wed, 11 Dec 2024 11:21:10 +0530 Subject: [PATCH 1/3] Fixed input name and output name in case of sim model Signed-off-by: Harsh Peswani --- .../aimet_onnx/sequential_mse/seq_mse.py | 67 +++++++++++++++++-- 1 file changed, 61 insertions(+), 6 deletions(-) diff --git a/TrainingExtensions/onnx/src/python/aimet_onnx/sequential_mse/seq_mse.py b/TrainingExtensions/onnx/src/python/aimet_onnx/sequential_mse/seq_mse.py index e2289c48e78..3d6198df591 100644 --- a/TrainingExtensions/onnx/src/python/aimet_onnx/sequential_mse/seq_mse.py +++ b/TrainingExtensions/onnx/src/python/aimet_onnx/sequential_mse/seq_mse.py @@ -151,6 +151,36 @@ def __init__(self, self.static_tensor_name_to_proto) self.quantizers_to_be_disabled = self._get_quantizers_to_be_disabled() # check this + def _update_value_info_for_output(self, node): + """ + Updates the value info for output of a node in sim model. + Value info for QcQuantizeOp is not present in _sim_extractor + + :param node: onnx node + """ + + input_name = node.input[0] + output_name = node.output[0] + if input_name in self._sim_extractor.vimap and output_name not in self._sim_extractor.vimap: + value_info_for_output = copy.deepcopy(self._sim_extractor.vimap[input_name]) + value_info_for_output.name = node.output[0] + self._sim_extractor.vimap[node.output[0]] = value_info_for_output + + def _update_value_info_for_input(self, node): + """ + Updates the value info for input of a node in sim model. + Value info for QcQuantizeOp is not present in _sim_extractor + + :param node: onnx node + """ + + input_name = node.input[0] + output_name = node.output[0] + if output_name in self._sim_extractor.vimap and input_name not in self._sim_extractor.vimap: + value_info_for_input = copy.deepcopy(self._sim_extractor.vimap[output_name]) + value_info_for_input.name = node.input[0] + self._sim_extractor.vimap[node.input[0]] = value_info_for_input + def _update_value_info(self): """ Updates the value info for sim model. @@ -159,11 +189,8 @@ def _update_value_info(self): for node in self.sim.model.nodes(): if node.op_type == "QcQuantizeOp": - input_name = node.input[0] - if input_name in self._sim_extractor.vimap: - value_info_for_output = copy.deepcopy(self._sim_extractor.vimap[input_name]) - value_info_for_output.name = node.output[0] - self._sim_extractor.vimap[node.output[0]] = value_info_for_output + self._update_value_info_for_output(node) + self._update_value_info_for_input(node) def _fill_static_tensor_name_to_proto(self): """ @@ -501,6 +528,33 @@ def _get_inputs_from_dependencies(self, dependency_node: DependencyNode): return float_inputs, sim_inputs + def _get_input_and_output_names_for_sim(self, input_names, output_names): + """ + Returns the input_name and output_name for sim model, input_name and + output_name can be different from float model + + :param input_names: input names for float model + :param output_names: output names for float model + :return: modified input names and modified output names + """ + + model_output_names = [model_output.name for model_output in self.model.model.graph.output] + input_names_for_sim = list() + for input_name in input_names: + if input_name not in model_output_names: + input_names_for_sim.append(input_name) + else: + input_names_for_sim.append(input_name + '_updated') + + output_names_for_sim = list() + for output_name in output_names: + if output_name not in model_output_names: + output_names_for_sim.append(output_name) + else: + output_names_for_sim.append(output_name + '_updated') + + return input_names_for_sim, output_names_for_sim + def _split_onnx_graph(self, input_names, output_names): """ Splits the onnx graph from input names to output names using extractor @@ -510,7 +564,8 @@ def _split_onnx_graph(self, input_names, output_names): :return: float split model and sim split model """ float_split_model = self._float_extractor.extract_model(list(input_names), list(output_names)) - sim_split_model = self._sim_extractor.extract_model(list(input_names), list(output_names)) + input_names_for_sim, output_names_for_sim = self._get_input_and_output_names_for_sim(input_names, output_names) + sim_split_model = self._sim_extractor.extract_model(list(input_names_for_sim), list(output_names_for_sim)) return float_split_model, sim_split_model def _run_onnx_graph(self, model, inputs): From dbce7f14a45dee10d6955a0dafa1fbbb3c158a59 Mon Sep 17 00:00:00 2001 From: Harsh Peswani Date: Thu, 12 Dec 2024 09:43:31 +0530 Subject: [PATCH 2/3] Bug in extractor api, resolved by adding value info of outputs to vimap of extractor api Signed-off-by: Harsh Peswani --- .../aimet_onnx/sequential_mse/seq_mse.py | 47 +++++++------------ 1 file changed, 17 insertions(+), 30 deletions(-) diff --git a/TrainingExtensions/onnx/src/python/aimet_onnx/sequential_mse/seq_mse.py b/TrainingExtensions/onnx/src/python/aimet_onnx/sequential_mse/seq_mse.py index 3d6198df591..adca9ec7307 100644 --- a/TrainingExtensions/onnx/src/python/aimet_onnx/sequential_mse/seq_mse.py +++ b/TrainingExtensions/onnx/src/python/aimet_onnx/sequential_mse/seq_mse.py @@ -99,7 +99,7 @@ def __init__(self, :param model: float model :param sim: QuantizationSimModel object - :param data_loader: Data loader + :param data_loader: Torch Dataloader :param params: Sequential MSE parameters """ @@ -181,6 +181,20 @@ def _update_value_info_for_input(self, node): value_info_for_input.name = node.input[0] self._sim_extractor.vimap[node.input[0]] = value_info_for_input + def _update_value_info_for_graph_output(self): + """ + Updates the value info for input of a node in sim model. + Value info for QcQuantizeOp is not present in _sim_extractor + + :param node: onnx node + """ + + for value_info in self.model.model.graph.output: + self._float_extractor.vimap[value_info.name] = value_info + + for value_info in self.sim.model.model.graph.output: + self._sim_extractor.vimap[value_info.name] = value_info + def _update_value_info(self): """ Updates the value info for sim model. @@ -191,6 +205,7 @@ def _update_value_info(self): if node.op_type == "QcQuantizeOp": self._update_value_info_for_output(node) self._update_value_info_for_input(node) + self._update_value_info_for_graph_output() def _fill_static_tensor_name_to_proto(self): """ @@ -528,33 +543,6 @@ def _get_inputs_from_dependencies(self, dependency_node: DependencyNode): return float_inputs, sim_inputs - def _get_input_and_output_names_for_sim(self, input_names, output_names): - """ - Returns the input_name and output_name for sim model, input_name and - output_name can be different from float model - - :param input_names: input names for float model - :param output_names: output names for float model - :return: modified input names and modified output names - """ - - model_output_names = [model_output.name for model_output in self.model.model.graph.output] - input_names_for_sim = list() - for input_name in input_names: - if input_name not in model_output_names: - input_names_for_sim.append(input_name) - else: - input_names_for_sim.append(input_name + '_updated') - - output_names_for_sim = list() - for output_name in output_names: - if output_name not in model_output_names: - output_names_for_sim.append(output_name) - else: - output_names_for_sim.append(output_name + '_updated') - - return input_names_for_sim, output_names_for_sim - def _split_onnx_graph(self, input_names, output_names): """ Splits the onnx graph from input names to output names using extractor @@ -564,8 +552,7 @@ def _split_onnx_graph(self, input_names, output_names): :return: float split model and sim split model """ float_split_model = self._float_extractor.extract_model(list(input_names), list(output_names)) - input_names_for_sim, output_names_for_sim = self._get_input_and_output_names_for_sim(input_names, output_names) - sim_split_model = self._sim_extractor.extract_model(list(input_names_for_sim), list(output_names_for_sim)) + sim_split_model = self._sim_extractor.extract_model(list(input_names), list(output_names)) return float_split_model, sim_split_model def _run_onnx_graph(self, model, inputs): From 46934c1254e4ed1a09d947f2dd43dd4510df376c Mon Sep 17 00:00:00 2001 From: Harsh Peswani Date: Wed, 18 Dec 2024 17:01:38 +0530 Subject: [PATCH 3/3] moved self._update_value_info_for_graph_output() call outside the for loop Signed-off-by: Harsh Peswani --- .../onnx/src/python/aimet_onnx/sequential_mse/seq_mse.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/TrainingExtensions/onnx/src/python/aimet_onnx/sequential_mse/seq_mse.py b/TrainingExtensions/onnx/src/python/aimet_onnx/sequential_mse/seq_mse.py index adca9ec7307..41f20e7cea2 100644 --- a/TrainingExtensions/onnx/src/python/aimet_onnx/sequential_mse/seq_mse.py +++ b/TrainingExtensions/onnx/src/python/aimet_onnx/sequential_mse/seq_mse.py @@ -205,7 +205,8 @@ def _update_value_info(self): if node.op_type == "QcQuantizeOp": self._update_value_info_for_output(node) self._update_value_info_for_input(node) - self._update_value_info_for_graph_output() + + self._update_value_info_for_graph_output() def _fill_static_tensor_name_to_proto(self): """