Merge pull request #445 from lucienwang1009/gru_rewriter

Refactor gru_rewriter

Author: nbcsm

tests/test_custom_rnncell.py

@@ -198,6 +198,53 @@ def test_attention_wrapper_lstm_encoder(self):
         output_names_with_port = ["output_0:0", "output:0", "final_state:0"]
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, 0.1)
 
+    @check_opset_min_version(8, "Scan")
+    @check_tf_min_version("1.8")
+    def test_attention_wrapper_gru_encoder(self):
+        size = 5
+        time_step = 3
+        input_size = 4
+        attn_size = size
+
+        batch_size = 9
+
+        # shape  [batch size, time step, size]
+        # attention_state: usually the output of an RNN encoder.
+        # This tensor should be shaped `[batch_size, max_time, ...]`
+        encoder_time_step = time_step
+        encoder_x_val = np.random.randn(encoder_time_step, input_size).astype('f')
+        encoder_x_val = np.stack([encoder_x_val] * batch_size)
+        encoder_x = tf.placeholder(tf.float32, encoder_x_val.shape, name="input_1")
+        encoder_cell = tf.nn.rnn_cell.GRUCell(size)
+        output, attr_state = tf.nn.dynamic_rnn(encoder_cell, encoder_x, dtype=tf.float32)
+        _ = tf.identity(output, name="output_0")
+        attention_states = output
+        attention_mechanism = tf.contrib.seq2seq.BahdanauAttention(attn_size,
+                                                                   attention_states)
+
+        match_input_fn = lambda curr_input, state: tf.concat([curr_input, state], axis=-1)
+        cell = tf.nn.rnn_cell.GRUCell(size)
+        match_cell_fw = tf.contrib.seq2seq.AttentionWrapper(cell,
+                                                            attention_mechanism,
+                                                            attention_layer_size=attn_size,
+                                                            cell_input_fn=match_input_fn,
+                                                            output_attention=False)
+
+        decoder_time_step = 6
+        decoder_x_val = np.random.randn(decoder_time_step, input_size).astype('f')
+        decoder_x_val = np.stack([decoder_x_val] * batch_size)
+
+        decoder_x = tf.placeholder(tf.float32, decoder_x_val.shape, name="input_2")
+        output, attr_state = tf.nn.dynamic_rnn(match_cell_fw, decoder_x, dtype=tf.float32)
+
+        _ = tf.identity(output, name="output")
+        _ = tf.identity(attr_state.cell_state, name="final_state")
+
+        feed_dict = {"input_1:0": encoder_x_val, "input_2:0": decoder_x_val}
+        input_names_with_port = ["input_1:0", "input_2:0"]
+        output_names_with_port = ["output_0:0", "output:0", "final_state:0"]
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, 0.1)
+
     @check_opset_min_version(8, "Scan")
     @check_tf_min_version("1.8")
     def test_attention_wrapper_lstm_encoder_input_has_none_dim(self):

tf2onnx/rewriter/__init__.py

@@ -10,7 +10,7 @@
 from tf2onnx.rewriter.random_uniform import rewrite_random_uniform, rewrite_random_uniform_fold_const
 from tf2onnx.rewriter.leakyrelu_rewriter import rewrite_leakyrelu
 from tf2onnx.rewriter.rnn import rewrite_single_direction_lstm, rewrite_bi_direction_lstm, \
-    rewrite_single_direction_gru, rewrite_single_direction_grublock, rewrite_bi_direction_gru, \
+    rewrite_single_direction_gru, rewrite_bi_direction_gru, \
     rewrite_custom_rnn_cell, rewrite_generic_loop
 
 __all__ = [
@@ -21,7 +21,6 @@
     "rewrite_single_direction_lstm",
     "rewrite_bi_direction_lstm",
     "rewrite_single_direction_gru",
-    "rewrite_single_direction_grublock",
     "rewrite_bi_direction_gru",
     "rewrite_custom_rnn_cell",
     "rewrite_generic_loop"

tf2onnx/rewriter/gru_rewriter.py

@@ -193,10 +193,13 @@ def rewrite(self, context):
         return REWRITER_RESULT.FAIL
 
     def run_internal(self):
+        loopcond_ops = []
         for op in self.g.get_nodes():
-            if not is_loopcond_op(op):
-                continue
+            if is_loopcond_op(op):
+                loopcond_ops.append(op)
 
+        # self.g.get_nodes may change inside this loop so that we parse all LoopCond first
+        for op in loopcond_ops:
             log.debug("======================\n handling loop cond node called %s", op.name)
             context = self.create_context()
             context.loop_cond = op