From 2d79052ec38b831f3254b20e0f6a42b3f98eabc7 Mon Sep 17 00:00:00 2001 From: Adrian Lizarraga Date: Fri, 1 Mar 2024 18:39:51 -0800 Subject: [PATCH] [QNN Quant] Add preprocessing option to transpose graph inputs/outputs to channel-last (#19731) ### Description Adds the optional parameters `inputs_to_make_channel_last` and `outputs_to_make_channel_last` to the `qnn_preprocess_model()` function. ```python """ 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. """ ``` **NOTE: If you use these options with the quantization scripts, you'll have to make sure your data_reader feeds in transposed input data. It won't happen automatically.** ### Motivation and Context Native QNN operators use the channel-last data layout, but ONNX uses channel-first. To bridge the gap, ORT's layout transformer inserts transposes around layout-sensitive nodes and updates their domain to indicate that they now operate on channel-last data. The transpose optimizer is able to remove most of these inserted transposes, but not all transposes can always be removed (i.e., some could remain at the graph's inputs and outputs). We've found that these extra transpose nodes can significantly degrade inference latency on QNN EP. One workaround (provided by this PR) is to add _additional_ transpose nodes at the graph inputs or outputs. These additional nodes can often help the ORT transpose optimizer cancel out any remaining transpose nodes, which significantly improves latency. Additionally, it may make more sense for some kinds of inputs to just be in channel-last form (e.g., images), avoiding the need to pre-transpose of the input data before inference. Example at the input: ``` Original: input0 (N, C, D1, D2, ..., Dn) --> Updated: input0 (N, D1, D2, ..., Dn, C) --> Transpose --> input0_chanfirst (N, C, D1, D2, ..., Dn) --> ``` Example at the output: ``` Original: --> output0 (N, C, D1, D2, ..., Dn) Updated: --> output0_chanfirst (N, C, D1, D2, ..., Dn) --> Transpose --> output0 (N, D1, D2, ..., Dn, C) ``` --- .../execution_providers/qnn/preprocess.py | 198 ++++++++++++++++++ .../quantization/test_qnn_preprocess_model.py | 93 ++++++++ 2 files changed, 291 insertions(+) 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()