Skip to content

Commit

Permalink
add dynamic split optimizer test
Browse files Browse the repository at this point in the history
Signed-off-by: Deyu Huang <[email protected]>
  • Loading branch information
hwangdeyu committed Apr 29, 2022
1 parent 9eb4694 commit 3739c79
Showing 1 changed file with 24 additions and 4 deletions.
28 changes: 24 additions & 4 deletions tests/test_optimizers.py
Original file line number Diff line number Diff line change
Expand Up @@ -133,15 +133,35 @@ def test_transpose_with_split(self, input_shape, perm, inner_perm):
graph = helper.make_graph(
[node1, node2, node3],
"test_transpose_with_split",
[helper.make_tensor_value_info("X", TensorProto.INT64, input_shape_with_trans)],
[helper.make_tensor_value_info("res", TensorProto.INT64, output_shape)],
[helper.make_tensor_value_info("X", TensorProto.FLOAT, input_shape_with_trans)],
[helper.make_tensor_value_info("res", TensorProto.FLOAT, output_shape)],
)

model_proto = self.make_model(graph, producer_name="onnx-tests")

feed_dict = {"X": np.random.randn(*input_shape_with_trans).astype(np.int64)}
feed_dict = {"X": np.random.randn(*input_shape_with_trans).astype(np.float32)}
self.run_transpose_compare(["res"], feed_dict, model_proto, remaining_transpose_num=0)

@parameterized.expand([
((1, -1), (1, 1710), (1710,), [1, 0]),
((3, 1, 1, 5, -1), (3, 1, 1, 5, 6), (3, 5, 6), [0, 2, 3, 4, 1]),
])
@check_opset_max_version(12, "split attribute changed to input in opset 13")
def test_transpose_with_split_dynamic_shape(self, input_shape, specific_input, output_shape, perm):
node1 = helper.make_node("Transpose", ["X"], ["Y"], perm=perm, name="trans")
node2 = helper.make_node("Split", ["Y"], ["Z"], axis=1, split=[1], name="split")
node3 = helper.make_node("Squeeze", ["Z"], ["B"], name="squeeze")

graph = helper.make_graph(
[node1, node2, node3],
"test_transpose_with_split_dynamic_shape",
[helper.make_tensor_value_info("X", TensorProto.FLOAT, input_shape)],
[helper.make_tensor_value_info("B", TensorProto.FLOAT, output_shape)],
)

model_proto = self.make_model(graph, producer_name="onnx-tests")
self.run_transpose_compare(["B"], {"X": np.random.randn(*specific_input).astype(np.float32)},
model_proto, remaining_transpose_num=0)

@parameterized.expand([
((2, 3, 4), [2, 0, 1], [1, 2, 0]),
((2, 3, 4, 5), [0, 2, 3, 1], [0, 3, 1, 2]),
Expand Down

0 comments on commit 3739c79

Please sign in to comment.