check specific op (lstm, gru) in converted graph

Author: wayuanho

tests/backend_test_base.py

@@ -93,7 +93,7 @@ def _run_backend(self, g, outputs, input_dict):
 
     def run_test_case(self, feed_dict, input_names_with_port, output_names_with_port, rtol=1e-07, atol=0.,
                       convert_var_to_const=True, constant_fold=True, check_value=True, check_shape=False,
-                      check_dtype=True, process_args=None, onnx_feed_dict=None):
+                      check_dtype=True, process_args=None, onnx_feed_dict=None, graph_validator=None):
         # optional - passed to process_tf_graph
         if process_args is None:
             process_args = {}
@@ -150,6 +150,9 @@ def run_test_case(self, feed_dict, input_names_with_port, output_names_with_port
             if check_shape:
                 self.assertEqual(expected_val.shape, actual_val.shape)
 
+        if graph_validator:
+            self.assertTrue(graph_validator(g))
+
         return g
 
     def save_onnx_model(self, model_proto, feed_dict, postfix=""):

tests/common.py

@@ -224,3 +224,15 @@ def group_nodes_by_type(graph):
     for node in graph.get_nodes():
         res[node.type].append(node)
     return res
+
+
+def check_op_count(graph, op_type, expected_count):
+    return len(group_nodes_by_type(graph)[op_type]) == expected_count
+
+
+def check_lstm_count(graph, expected_count):
+    return check_op_count(graph, "LSTM", expected_count)
+
+
+def check_gru_count(graph, expected_count):
+    return check_op_count(graph, "GRU", expected_count)

tests/test_gru.py

@@ -16,7 +16,7 @@
 from tensorflow.python.ops import init_ops
 from tensorflow.python.ops import variable_scope
 from backend_test_base import Tf2OnnxBackendTestBase
-from common import unittest_main
+from common import unittest_main, check_gru_count
 
 
 # pylint: disable=missing-docstring,invalid-name,unused-argument,using-constant-test
@@ -47,7 +47,8 @@ def test_single_dynamic_gru(self):
         input_names_with_port = ["input_1:0"]
         feed_dict = {"input_1:0": x_val}
         output_names_with_port = ["output:0", "cell_state:0"]
-        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-03, atol=1e-06)
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-03, atol=1e-06,
+                           graph_validator=lambda g: check_gru_count(g, 1))
 
     def test_multiple_dynamic_gru(self):
         units = 5
@@ -93,7 +94,8 @@ def test_multiple_dynamic_gru(self):
         feed_dict = {"input_1:0": x_val}
         input_names_with_port = ["input_1:0"]
         output_names_with_port = ["output:0", "cell_state:0"]
-        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3, atol=1e-06)
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3, atol=1e-06,
+                           graph_validator=lambda g: check_gru_count(g, 2))
 
     def test_single_dynamic_gru_seq_length_is_const(self):
         units = 5
@@ -119,7 +121,8 @@ def test_single_dynamic_gru_seq_length_is_const(self):
         feed_dict = {"input_1:0": x_val}
         input_names_with_port = ["input_1:0"]
         output_names_with_port = ["output:0", "cell_state:0"]
-        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3, atol=1e-06)
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3, atol=1e-06,
+                           graph_validator=lambda g: check_gru_count(g, 1))
 
     def test_single_dynamic_gru_seq_length_is_not_const(self):
         units = 5
@@ -148,7 +151,8 @@ def test_single_dynamic_gru_seq_length_is_not_const(self):
         feed_dict = {"input_1:0": x_val, "input_2:0": y_val}
         input_names_with_port = ["input_1:0", "input_2:0"]
         output_names_with_port = ["output:0", "cell_state:0"]
-        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-03, atol=1e-06)
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-03, atol=1e-06,
+                           graph_validator=lambda g: check_gru_count(g, 1))
 
     def test_single_dynamic_gru_placeholder_input(self):
         units = 5
@@ -172,7 +176,8 @@ def test_single_dynamic_gru_placeholder_input(self):
         feed_dict = {"input_1:0": x_val}
         input_names_with_port = ["input_1:0"]
         output_names_with_port = ["output:0", "cell_state:0"]
-        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-03, atol=1e-06)
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-03, atol=1e-06,
+                           graph_validator=lambda g: check_gru_count(g, 1))
 
     def test_single_dynamic_gru_ch_zero_state_initializer(self):
         units = 5
@@ -201,7 +206,8 @@ def test_single_dynamic_gru_ch_zero_state_initializer(self):
         feed_dict = {"input_1:0": x_val}
         input_names_with_port = ["input_1:0"]
         output_names_with_port = ["output:0", "cell_state:0"]
-        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-03, atol=1e-06)
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-03, atol=1e-06,
+                           graph_validator=lambda g: check_gru_count(g, 1))
 
     @unittest.skip("FIXME: disable for now for accuracy problem")
     def test_single_dynamic_gru_random_weights(self):
@@ -229,7 +235,8 @@ def test_single_dynamic_gru_random_weights(self):
         feed_dict = {"input_1:0": x_val}
         input_names_with_port = ["input_1:0"]
         output_names_with_port = ["output:0", "cell_state:0"]
-        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, 0.0001)
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, 0.0001,
+                           graph_validator=lambda g: check_gru_count(g, 1))
 
     @unittest.skip("FIXME: disable for now for accuracy problem")
     def test_single_dynamic_gru_random_weights2(self):
@@ -256,7 +263,8 @@ def test_single_dynamic_gru_random_weights2(self):
         feed_dict = {"input_1:0": x_val}
         input_names_with_port = ["input_1:0"]
         output_names_with_port = ["output:0", "cell_state:0"]
-        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, 0.01)
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, 0.01,
+                           graph_validator=lambda g: check_gru_count(g, 1))
 
     def test_dynamic_gru_output_consumed_only(self):
         units = 5
@@ -280,7 +288,8 @@ def test_dynamic_gru_output_consumed_only(self):
         feed_dict = {"input_1:0": x_val}
         input_names_with_port = ["input_1:0"]
         output_names_with_port = ["output:0"]
-        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, 0.0001)
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, 0.0001,
+                           graph_validator=lambda g: check_gru_count(g, 1))
 
     def test_dynamic_gru_state_consumed_only(self):
         units = 5
@@ -304,7 +313,8 @@ def test_dynamic_gru_state_consumed_only(self):
         feed_dict = {"input_1:0": x_val}
         input_names_with_port = ["input_1:0"]
         output_names_with_port = ["cell_state:0"]
-        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=0.0001, atol=1e-06)
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=0.0001, atol=1e-06,
+                           graph_validator=lambda g: check_gru_count(g, 1))
 
     def test_dynamic_bigru(self):
         units = 5
@@ -335,7 +345,8 @@ def test_dynamic_bigru(self):
         feed_dict = {"input_1:0": x_val}
         input_names_with_port = ["input_1:0"]
         output_names_with_port = ["output:0", "cell_state:0"]
-        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3, atol=1e-06)
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3, atol=1e-06,
+                           graph_validator=lambda g: check_gru_count(g, 1))
 
     def test_dynamic_bigru_output_consumed_only(self):
         units = 5
@@ -365,7 +376,8 @@ def test_dynamic_bigru_output_consumed_only(self):
         feed_dict = {"input_1:0": x_val}
         input_names_with_port = ["input_1:0"]
         output_names_with_port = ["output:0"]
-        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3, atol=1e-06)
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3, atol=1e-06,
+                           graph_validator=lambda g: check_gru_count(g, 1))
 
     def test_dynamic_bigru_state_consumed_only(self):
         units = 5
@@ -395,7 +407,8 @@ def test_dynamic_bigru_state_consumed_only(self):
         feed_dict = {"input_1:0": x_val}
         input_names_with_port = ["input_1:0"]
         output_names_with_port = ["cell_state:0"]
-        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3, atol=1e-06)
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3, atol=1e-06,
+                           graph_validator=lambda g: check_gru_count(g, 1))
 
     def test_dynamic_bidirectional_but_one_gru(self):
         units = 5
@@ -423,7 +436,8 @@ def test_dynamic_bidirectional_but_one_gru(self):
         feed_dict = {"input_1:0": x_val}
         input_names_with_port = ["input_1:0"]
         output_names_with_port = ["output:0", "cell_state:0"]
-        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3, atol=1e-06)
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3, atol=1e-06,
+                           graph_validator=lambda g: check_gru_count(g, 1))
 
     def test_dynamic_bidirectional_but_one_gru_and_output_consumed_only(self):
         units = 5
@@ -448,7 +462,8 @@ def test_dynamic_bidirectional_but_one_gru_and_output_consumed_only(self):
         feed_dict = {"input_1:0": x_val}
         input_names_with_port = ["input_1:0"]
         output_names_with_port = ["output:0"]
-        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3, atol=1e-06)
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3, atol=1e-06,
+                           graph_validator=lambda g: check_gru_count(g, 1))
 
     def test_dynamic_bidirectional_but_one_gru_and_state_consumed_only(self):
         units = 5
@@ -473,7 +488,8 @@ def test_dynamic_bidirectional_but_one_gru_and_state_consumed_only(self):
         feed_dict = {"input_1:0": x_val}
         input_names_with_port = ["input_1:0"]
         output_names_with_port = ["cell_state:0"]
-        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3, atol=1e-06)
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3, atol=1e-06,
+                           graph_validator=lambda g: check_gru_count(g, 1))
 
 
 if __name__ == '__main__':