saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-09 17:28:58 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[QNN Quant] Add preprocessing option to transpose graph inputs/outputs to channel-last (#19731)

Browse source 
### 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) --> <Nodes>
      Updated:
          input0 (N, D1, D2, ..., Dn, C) --> Transpose --> input0_chanfirst (N, C, D1, D2, ..., Dn) --> <Nodes>
      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:
          <Nodes> --> output0 (N, C, D1, D2, ..., Dn)
      Updated:
          <Nodes> --> 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) --> <Nodes>
Updated:
    input0 (N, D1, D2, ..., Dn, C) --> Transpose --> input0_chanfirst (N, C, D1, D2, ..., Dn) --> <Nodes>
```

Example at the output:
```
Original:
   <Nodes> --> output0 (N, C, D1, D2, ..., Dn)
Updated:
   <Nodes> --> output0_chanfirst (N, C, D1, D2, ..., Dn) --> Transpose --> output0 (N, D1, D2, ..., Dn, C)
```
This commit is contained in:
Adrian Lizarraga 2024-03-01 18:39:51 -08:00 committed by GitHub
parent de3158e78d
commit 2d79052ec3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 291 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

198
onnxruntime/python/tools/quantization/execution_providers/qnn/preprocess.py
View file

@ -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) --> <Nodes>
Updated:
input0 (N, D1, D2, ..., Dn, C) --> Transpose --> input0_chanfirst (N, C, D1, D2, ..., Dn) --> <Nodes>
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:
<Nodes> --> output0 (N, C, D1, D2, ..., Dn)
Updated:
<Nodes> --> 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)

93
onnxruntime/test/python/quantization/test_qnn_preprocess_model.py
View file

@ -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()