diff --git a/onnxruntime/python/tools/quantization/execution_providers/qnn/preprocess.py b/onnxruntime/python/tools/quantization/execution_providers/qnn/preprocess.py index b0dab81830..e584a65574 100644 --- a/onnxruntime/python/tools/quantization/execution_providers/qnn/preprocess.py +++ b/onnxruntime/python/tools/quantization/execution_providers/qnn/preprocess.py @@ -24,6 +24,8 @@ def qnn_preprocess_model( external_data_location: str | None = None, external_data_size_threshold: int = 1024, external_data_convert_attribute: bool = False, + inputs_to_make_channel_last: list[str] | None = None, + outputs_to_make_channel_last: list[str] | None = None, ) -> bool: """ If necessary, this method creates a new "pre-processed" model in preparation for @@ -52,6 +54,32 @@ def qnn_preprocess_model( external_data_convert_attribute: Effective only if save_as_external_data is true. Defaults to false. If true, convert all tensors to external data. If false, convert only non-attribute tensors to external data. + inputs_to_make_channel_last: List of graph input names to transpose to be "channel-last". For example, + if "input0" originally has the shape (N, C, D1, D2, ..., Dn), the resulting model will change input0's + shape to (N, D1, D2, ..., Dn, C) and add a transpose node after it. + + Original: + input0 (N, C, D1, D2, ..., Dn) --> + + Updated: + input0 (N, D1, D2, ..., Dn, C) --> Transpose --> input0_chanfirst (N, C, D1, D2, ..., Dn) --> + + This can potentially improve inference latency for QDQ models running on QNN EP because the + additional transpose node may allow other transpose nodes inserted during ORT layout transformation + to cancel out. + outputs_to_make_channel_last: List of graph output names to transpose to be "channel-last". For example, + if "output0" originally has the shape (N, C, D1, D2, ..., Dn), the resulting model will change output0's + shape to (N, D1, D2, ..., Dn, C) and add a transpose node before it. + + Original: + --> output0 (N, C, D1, D2, ..., Dn) + + Updated: + --> output0_chanfirst (N, C, D1, D2, ..., Dn) --> Transpose --> output0 (N, D1, D2, ..., Dn, C) + + This can potentially improve inference latency for QDQ models running on QNN EP because the + additional transpose node may allow other transpose nodes inserted during ORT layout transformation + to cancel out. """ modified = False model = onnx.load_model(model_input) @@ -83,6 +111,19 @@ def qnn_preprocess_model( if fusion_layernorm.apply(): modified = True + # Optionally, transpose inputs and/or outputs to make them "channel-last". + if inputs_to_make_channel_last or outputs_to_make_channel_last: + transpose_node_prefix = "Transpose_channel_" + transpose_node_suffix: int = onnx_model.get_largest_node_name_suffix(transpose_node_prefix) + 1 + update_io_to_channel_last( + onnx_model.model, + inputs_to_make_channel_last, + outputs_to_make_channel_last, + transpose_node_name_prefix=transpose_node_prefix, + transpose_node_name_start_suffix=transpose_node_suffix, + ) + modified = True + # Make sure all nodes have a name. unnamed_node_prefix = "qnn_preproc_node_" available_suffix = onnx_model.get_largest_node_name_suffix(unnamed_node_prefix) + 1 @@ -107,3 +148,160 @@ def qnn_preprocess_model( ) return modified + + +class InputOutputNameMap: + def __init__( + self, + orig_tensor_names: set[str], + orig_graph_inputs: dict[str, onnx.ValueInfoProto], + orig_graph_outputs: dict[str, onnx.ValueInfoProto], + ): + self.orig_tensor_names = orig_tensor_names + self.orig_graph_inputs = orig_graph_inputs + self.orig_graph_outputs = orig_graph_outputs + self.updated_io_names = {} + self.new_value_infos = [] + + def get_new_name(self, orig_name: str): + if orig_name in self.updated_io_names: + return self.updated_io_names[orig_name] + + # Make a new tensor name that is unique among all tensors in the graph. + prefix: str = f"{orig_name}_channel_first_" + suffix: int = -1 + for tensor_name in self.orig_tensor_names: + if tensor_name.startswith(prefix) and tensor_name[len(prefix) :].isdigit(): + index = int(tensor_name[len(prefix) :]) + suffix = max(suffix, index) + + suffix += 1 # This is the first available suffix. + new_name = f"{prefix}{suffix!s}" + + # Add new value_info objects for these new tensors. + orig_value_info = self.orig_graph_inputs.get(orig_name) or self.orig_graph_outputs[orig_name] + value_info_proto = onnx.ValueInfoProto() + value_info_proto.CopyFrom(orig_value_info) + value_info_proto.name = new_name + self.new_value_infos.append(value_info_proto) + + self.updated_io_names[orig_name] = new_name + return self.updated_io_names[orig_name] + + +def update_io_to_channel_last( + model: onnx.ModelProto, + inputs_to_update: list[str] | None, + outputs_to_update: list[str] | None, + transpose_node_name_prefix: str = "Transpose_channel_", + transpose_node_name_start_suffix: int = 0, +): + inputs_to_update = set(inputs_to_update or []) + outputs_to_update = set(outputs_to_update or []) + + if not inputs_to_update and not outputs_to_update: + return + + graph = model.graph + orig_graph_inputs = {ginput.name: ginput for ginput in graph.input} + orig_graph_outputs = {goutput.name: goutput for goutput in graph.output} + + # Check that the user passed in actual input and output names. + for input_name in inputs_to_update: + if input_name not in orig_graph_inputs: + raise ValueError(f"{input_name} is not a graph input") + + for output_name in outputs_to_update: + if output_name not in orig_graph_outputs: + raise ValueError(f"{output_name} is not a graph output") + + orig_tensor_names = set() + orig_tensor_names.update(set(orig_graph_inputs)) + orig_tensor_names.update(set(orig_graph_outputs)) + orig_tensor_names.update(input_name for node in graph.node for input_name in node.input if input_name) + + # Maps original input (or output) name to its updated name used within the graph. + io_map = InputOutputNameMap(orig_tensor_names, orig_graph_inputs, orig_graph_outputs) + + # Update each node's inputs/outputs to use the transposed versions. + for node in graph.node: + for i in range(len(node.input)): + if node.input[i] and node.input[i] in inputs_to_update: + node.input[i] = io_map.get_new_name(node.input[i]) + elif node.input[i] and node.input[i] in outputs_to_update: + node.input[i] = io_map.get_new_name(node.input[i]) + + for i in range(len(node.output)): + if node.output[i] in outputs_to_update: + node.output[i] = io_map.get_new_name(node.output[i]) + + # Update graph inputs to channel-last and a Transpose (to channel-first) after each. + for g_input_name in inputs_to_update: + g_input = orig_graph_inputs[g_input_name] + + if not g_input.type.HasField("tensor_type") or not g_input.type.tensor_type.HasField("shape"): + raise ValueError(f"Expected input {g_input.name} to have a tensor_type with a shape") + + input_shape = g_input.type.tensor_type.shape + input_rank = len(input_shape.dim) + + if input_rank < 3: + raise ValueError(f"Expected input {g_input.name} to be of rank >= 3") + + channel_dim = onnx.TensorShapeProto.Dimension() + channel_dim.CopyFrom(input_shape.dim[1]) + for i in range(1, input_rank - 1): + input_shape.dim[i].CopyFrom(input_shape.dim[i + 1]) + input_shape.dim[input_rank - 1].CopyFrom(channel_dim) + + transpose_perm = list(range(input_rank)) + for i in range(input_rank): + transpose_perm[i] = i if i < 1 else i - 1 + transpose_perm[1] = input_rank - 1 + + transpose_node = onnx.helper.make_node( + "Transpose", + name=f"{transpose_node_name_prefix}{transpose_node_name_start_suffix!s}", + inputs=[g_input.name], + outputs=[io_map.get_new_name(g_input.name)], + perm=transpose_perm, + ) + transpose_node_name_start_suffix += 1 + + graph.node.extend([transpose_node]) + + # Update graph outputs to channel-last and a Transpose (from channel-first) before each. + for g_output_name in outputs_to_update: + g_output = orig_graph_outputs[g_output_name] + if not g_output.type.HasField("tensor_type") or not g_output.type.tensor_type.HasField("shape"): + raise ValueError(f"Expected output {g_output.name} to have a tensor_type with a shape") + + output_shape = g_output.type.tensor_type.shape + output_rank = len(output_shape.dim) + + if output_rank < 3: + raise ValueError(f"Expected output {g_output.name} to be of rank >= 3") + + channel_dim = onnx.TensorShapeProto.Dimension() + channel_dim.CopyFrom(output_shape.dim[1]) + for i in range(1, output_rank - 1): + output_shape.dim[i].CopyFrom(output_shape.dim[i + 1]) + output_shape.dim[output_rank - 1].CopyFrom(channel_dim) + + transpose_perm = list(range(output_rank)) + for i in range(output_rank): + transpose_perm[i] = i if i == 0 else i + 1 + transpose_perm[output_rank - 1] = 1 + + transpose_node = onnx.helper.make_node( + "Transpose", + name=f"{transpose_node_name_prefix}{transpose_node_name_start_suffix!s}", + inputs=[io_map.get_new_name(g_output.name)], + outputs=[g_output.name], + perm=transpose_perm, + ) + transpose_node_name_start_suffix += 1 + + graph.node.extend([transpose_node]) + + graph.value_info.extend(io_map.new_value_infos) diff --git a/onnxruntime/test/python/quantization/test_qnn_preprocess_model.py b/onnxruntime/test/python/quantization/test_qnn_preprocess_model.py index 9b67fd41ca..6503b3223b 100644 --- a/onnxruntime/test/python/quantization/test_qnn_preprocess_model.py +++ b/onnxruntime/test/python/quantization/test_qnn_preprocess_model.py @@ -12,6 +12,7 @@ from pathlib import Path import numpy as np import onnx +import onnxruntime from onnxruntime.quantization.execution_providers.qnn import qnn_preprocess_model from onnxruntime.quantization.quant_utils import model_has_external_data, ms_domain @@ -165,6 +166,98 @@ class TestQnnPreprocessModel(unittest.TestCase): for node in fused_model.graph.node: self.assertIn(node.op_type, expected_op_types) + def build_multi_input_output_model(self, shape): + """ + Returns the following model. + +----------> [X] + | + [A] ---> Add ---> Abs -+-> Mul ---> [Y] + ^ ^ + | | + [B] ------+-----------------+ + """ + input_a = onnx.helper.make_tensor_value_info("A", onnx.TensorProto.FLOAT, shape) + input_b = onnx.helper.make_tensor_value_info("B", onnx.TensorProto.FLOAT, shape) + output_x = onnx.helper.make_tensor_value_info("X", onnx.TensorProto.FLOAT, shape) + output_y = onnx.helper.make_tensor_value_info("Y", onnx.TensorProto.FLOAT, shape) + + add_node = onnx.helper.make_node("Add", ["A", "B"], ["add_out"], name="add_node") + abs_node = onnx.helper.make_node("Abs", ["add_out"], ["X"], name="abs_node") + mul_node = onnx.helper.make_node("Mul", ["X", "B"], ["Y"], name="mul_node") + + graph = onnx.helper.make_graph( + [add_node, abs_node, mul_node], + "multi_io_graph", + [input_a, input_b], + [output_x, output_y], + ) + opset_imports = [ + onnx.helper.make_opsetid("", 18), + ] + model = onnx.helper.make_model(graph, opset_imports=opset_imports) + return onnx.shape_inference.infer_shapes(model) + + def test_make_io_channel_last(self): + """ + Test making a model's inputs and outputs channel-last. + """ + model = self.build_multi_input_output_model((1, 2, 3, 4)) + onnx.save_model(model, "model.onnx") + modified = qnn_preprocess_model( + "model.onnx", + "model.qnn_pp.onnx", + inputs_to_make_channel_last=["A", "B"], + outputs_to_make_channel_last=["X", "Y"], + ) + + self.assertTrue(modified) + + preproc_model = onnx.load_model("model.qnn_pp.onnx") + self.assertEqual(len(preproc_model.graph.node), 7) + + num_transposes = sum(1 for node in preproc_model.graph.node if node.op_type == "Transpose") + self.assertEqual(num_transposes, 4) + + # Check that the outputs of the new model are the same, but transposed. + input_a = np.arange(0.0, 24.0, 1.0, dtype=np.float32).reshape((1, 2, 3, 4)) + input_a_t = input_a.transpose(0, 2, 3, 1) + input_b = np.arange(1.0, 25.0, 1.0, dtype=np.float32).reshape((1, 2, 3, 4)) + input_b_t = input_b.transpose(0, 2, 3, 1) + + orig_session = onnxruntime.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"]) + orig_results = orig_session.run(None, {"A": input_a, "B": input_b}) + + new_session = onnxruntime.InferenceSession( + preproc_model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + new_results = new_session.run(None, {"A": input_a_t, "B": input_b_t}) + + self.assertEqual(len(orig_results), len(new_results)) + for idx, orig_output in enumerate(orig_results): + transposed_output = new_results[idx] + np.testing.assert_allclose( + orig_output, + transposed_output.transpose(0, 3, 1, 2), + err_msg=f"Channel-last model output {idx} differs", + ) + + def test_make_io_channel_last_rank_error(self): + """ + Test making a model's inputs and outputs channel-last with a rank < 3 (error). + """ + model = self.build_multi_input_output_model((1, 2)) + onnx.save_model(model, "model.onnx") + + with self.assertRaises(ValueError) as context: + qnn_preprocess_model( + "model.onnx", + "model.qnn_pp.onnx", + inputs_to_make_channel_last=["A", "B"], + outputs_to_make_channel_last=["X", "Y"], + ) + + self.assertIn("to be of rank >= 3", str(context.exception)) + if __name__ == "__main__": unittest.main()