mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-26 19:52:38 +00:00
fdce4fa6af
Description: Format all python files under onnxruntime with black and isort. After checking in, we can use .git-blame-ignore-revs to ignore the formatting PR in git blame. #11315, #11316
109 lines
4.4 KiB
Python
109 lines
4.4 KiB
Python
import numpy as np
|
|
import onnx
|
|
from onnx import TensorProto, helper
|
|
|
|
|
|
# Create graph with Add and Sub nodes that can be used to test partitioning when one of the operators
|
|
# can run using the test EP and the other cannot.
|
|
# As the operators take 2 inputs and produce one output we can easily create edges to test different scenarios
|
|
def create_model_1():
|
|
# Assume the EP can handle either Add or Sub but not both, and we need to minimize the partitions for nodes
|
|
# the EP can handle (as going to/from the EP has significant performance cost).
|
|
#
|
|
# graph inputs
|
|
# / \ \
|
|
# a1 s1 \
|
|
# | / \ \
|
|
# | a2 s2 \
|
|
# \ / |
|
|
# \ / |
|
|
# a3 a4
|
|
#
|
|
# Assuming the initial topological sort is top down, left to right, we get a1, s1, a2, s2, a3, a4
|
|
#
|
|
# Naively creating groups based on iterating this order and whether a node is supported gives the following groups
|
|
# (a1), (s1), (a2), (s2), (a3, a4). This is similar to what most EPs do currently.
|
|
#
|
|
# To improve on that, we may consider all reachable supported nodes when iterating the topological ordering.
|
|
# In this model, regardless of whether Add or Sub is supported, we get the groups (s1, s2), (a1, a2, a3, a4).
|
|
# One of those groups is the resulting partition.
|
|
#
|
|
# So if this model is loaded in a partitioning test, there should only be one partition running on the EP regardless
|
|
# of whether Add or Sub is supported by it.
|
|
graph = helper.make_graph(
|
|
nodes=[
|
|
helper.make_node("Add", ["input0", "input1"], ["1"], "A1"),
|
|
helper.make_node("Sub", ["input0", "input1"], ["2"], "S1"),
|
|
helper.make_node("Add", ["2", "input1"], ["3_out"], "A2"),
|
|
helper.make_node("Sub", ["2", "input1"], ["4"], "S2"),
|
|
helper.make_node("Add", ["1", "4"], ["5_out"], "A3"),
|
|
helper.make_node("Add", ["input1", "input2"], ["6_out"], "A4"),
|
|
],
|
|
name="graph",
|
|
inputs=[
|
|
helper.make_tensor_value_info("input0", TensorProto.INT64, [1]),
|
|
helper.make_tensor_value_info("input1", TensorProto.INT64, [1]),
|
|
helper.make_tensor_value_info("input2", TensorProto.INT64, [1]),
|
|
],
|
|
outputs=[
|
|
helper.make_tensor_value_info("3_out", TensorProto.INT64, [1]),
|
|
helper.make_tensor_value_info("5_out", TensorProto.INT64, [1]),
|
|
helper.make_tensor_value_info("6_out", TensorProto.INT64, [1]),
|
|
],
|
|
initializer=[],
|
|
)
|
|
|
|
model = helper.make_model(graph)
|
|
return model
|
|
|
|
|
|
def create_model_2():
|
|
# Create a model where there's a node that can't be run breaking up the partition.
|
|
# Partition aware topo sort should give us
|
|
# s1, [a1, a2, a3, a5], s2, [a4, a6, a7]
|
|
#
|
|
# graph inputs
|
|
# s1
|
|
# |
|
|
# a1
|
|
# / \
|
|
# a2 s2
|
|
# / \ /
|
|
# a3 a4
|
|
# | |
|
|
# a5 a6
|
|
# \ /
|
|
# a7
|
|
graph = helper.make_graph(
|
|
nodes=[
|
|
helper.make_node("Sub", ["input0", "input1"], ["s1_out"], "S1"),
|
|
helper.make_node("Add", ["s1_out", "input2"], ["a1_out"], "A1"),
|
|
helper.make_node("Add", ["a1_out", "input0"], ["a2_out"], "A2"),
|
|
helper.make_node("Sub", ["a1_out", "input1"], ["s2_out"], "S2"),
|
|
helper.make_node("Add", ["a2_out", "input2"], ["a3_out"], "A3"),
|
|
helper.make_node("Add", ["a2_out", "s2_out"], ["a4_out"], "A4"),
|
|
helper.make_node("Add", ["a3_out", "input0"], ["a5_out"], "A5"),
|
|
helper.make_node("Add", ["a4_out", "input1"], ["a6_out"], "A6"),
|
|
helper.make_node("Add", ["a5_out", "a6_out"], ["a7_out"], "A7"),
|
|
],
|
|
name="graph",
|
|
inputs=[
|
|
helper.make_tensor_value_info("input0", TensorProto.INT64, [1]),
|
|
helper.make_tensor_value_info("input1", TensorProto.INT64, [1]),
|
|
helper.make_tensor_value_info("input2", TensorProto.INT64, [1]),
|
|
],
|
|
outputs=[
|
|
helper.make_tensor_value_info("a7_out", TensorProto.INT64, [1]),
|
|
],
|
|
initializer=[],
|
|
)
|
|
|
|
model = helper.make_model(graph)
|
|
return model
|
|
|
|
|
|
if __name__ == "__main__":
|
|
model = create_model_1()
|
|
onnx.save(model, "ep_partitioning_test_1.onnx")
|
|
model = create_model_2()
|
|
onnx.save(model, "ep_partitioning_test_2.onnx")
|