test_utils.py get_ops_names

Author: calvinmccarter-at-lightmatter

tests/onnx2pytorch/test_utils.py

@@ -1,12 +1,14 @@
-import torch
-import pytest
 import numpy as np
+import onnx
+import pytest
+import torch
 from torch import nn
 from onnx.backend.test.case.node.pad import pad_impl
 
 from onnx2pytorch.helpers import to_onnx
 from onnx2pytorch.utils import (
     is_constant,
+    get_ops_names,
     get_selection,
     assign_values_to_dim,
     get_activation_value,
@@ -146,3 +148,111 @@ def weight():
 def test_extract_padding_params(weight, onnx_pads, torch_pads):
     out_pads = extract_padding_params(onnx_pads)
     assert out_pads == torch_pads
+
+
+def test_get_ops_names():
+    y_in = onnx.helper.make_tensor_value_info("y_in", onnx.TensorProto.FLOAT, [1])
+    y_out = onnx.helper.make_tensor_value_info("y_out", onnx.TensorProto.FLOAT, [1])
+    scan_out = onnx.helper.make_tensor_value_info(
+        "scan_out", onnx.TensorProto.FLOAT, []
+    )
+    cond_in = onnx.helper.make_tensor_value_info("cond_in", onnx.TensorProto.BOOL, [])
+    cond_out = onnx.helper.make_tensor_value_info("cond_out", onnx.TensorProto.BOOL, [])
+    iter_count = onnx.helper.make_tensor_value_info(
+        "iter_count", onnx.TensorProto.INT64, []
+    )
+
+    x = np.array([1, 2, 3, 4, 5]).astype(np.float32)
+
+    x_const_node = onnx.helper.make_node(
+        "Constant",
+        inputs=[],
+        outputs=["x"],
+        value=onnx.helper.make_tensor(
+            name="const_tensor_x",
+            data_type=onnx.TensorProto.FLOAT,
+            dims=x.shape,
+            vals=x.flatten().astype(float),
+        ),
+    )
+
+    one_const_node = onnx.helper.make_node(
+        "Constant",
+        inputs=[],
+        outputs=["one"],
+        value=onnx.helper.make_tensor(
+            name="const_tensor_one", data_type=onnx.TensorProto.INT64, dims=(), vals=[1]
+        ),
+    )
+
+    i_add_node = onnx.helper.make_node(
+        "Add", inputs=["iter_count", "one"], outputs=["end"]
+    )
+
+    start_unsqueeze_node = onnx.helper.make_node(
+        "Unsqueeze", inputs=["iter_count"], outputs=["slice_start"], axes=[0]
+    )
+
+    end_unsqueeze_node = onnx.helper.make_node(
+        "Unsqueeze", inputs=["end"], outputs=["slice_end"], axes=[0]
+    )
+
+    slice_node = onnx.helper.make_node(
+        "Slice", inputs=["x", "slice_start", "slice_end"], outputs=["slice_out"]
+    )
+
+    y_add_node = onnx.helper.make_node(
+        "Add", inputs=["y_in", "slice_out"], outputs=["y_out"]
+    )
+
+    identity_node = onnx.helper.make_node(
+        "Identity", inputs=["cond_in"], outputs=["cond_out"]
+    )
+
+    scan_identity_node = onnx.helper.make_node(
+        "Identity", inputs=["y_out"], outputs=["scan_out"]
+    )
+
+    loop_body = onnx.helper.make_graph(
+        [
+            identity_node,
+            x_const_node,
+            one_const_node,
+            i_add_node,
+            start_unsqueeze_node,
+            end_unsqueeze_node,
+            slice_node,
+            y_add_node,
+            scan_identity_node,
+        ],
+        "loop_body",
+        [iter_count, cond_in, y_in],
+        [cond_out, y_out, scan_out],
+    )
+
+    node = onnx.helper.make_node(
+        "Loop",
+        inputs=["trip_count", "cond", "y"],
+        outputs=["res_y", "res_scan"],
+        body=loop_body,
+    )
+
+    trip_count = onnx.helper.make_tensor_value_info(
+        "trip_count", onnx.TensorProto.INT64, []
+    )
+    cond = onnx.helper.make_tensor_value_info("cond", onnx.TensorProto.BOOL, [])
+    y = onnx.helper.make_tensor_value_info("y", onnx.TensorProto.FLOAT, [1])
+    res_y = onnx.helper.make_tensor_value_info("res_y", onnx.TensorProto.FLOAT, [1])
+    res_scan = onnx.helper.make_tensor_value_info(
+        "res_scan", onnx.TensorProto.FLOAT, []
+    )
+
+    graph_def = onnx.helper.make_graph(
+        nodes=[node],
+        name="test-model",
+        inputs=[trip_count, cond, y],
+        outputs=[res_y, res_scan],
+    )
+
+    ops_names = set(["Add", "Constant", "Identity", "Loop", "Slice", "Unsqueeze"])
+    assert get_ops_names(graph_def) == ops_names