Can't convert LSTM cell correctly

[fanghuaqi]

Hi, I tried to convert an LSTM model from TF to ONNX, part of the code related to LSTM is as below:

keep_prob = tf.constant(1.0) #tf.placeholder(tf.float32,name="keepProb")
h_fc1_drop = tf.nn.dropout(h_fc1, keep_prob)

# ==== LSTM layer =====
h_fc1_expanded = tf.expand_dims(h_fc1_drop,1)
h_fc1_unstacked = tf.unstack(h_fc1_expanded,time_steps,0)
cell = tf.contrib.rnn.BasicLSTMCell(100)
cell_state = tf.placeholder(tf.float32, [1, 100], name='cell_state')
hidden_state = tf.placeholder(tf.float32, [1, 100], name='hidden_state')
initial_state = tf.nn.rnn_cell.LSTMStateTuple(cell_state, hidden_state)
outputs, current_state = tf.nn.static_rnn(cell, h_fc1_unstacked, initial_state, dtype=tf.float32)
lstmOutput = tf.convert_to_tensor(outputs)
lstmUnstacked = tf.unstack(lstmOutput,1,1)

I followed the steps here: https://github.com/onnx/tensorflow-onnx/blob/master/examples/rnn_tips.md#commands

Here is what I get after conversion:

Could you help me to figure out which part is wrong?

Thanks Huaqi

[zhijxu-MS]

it looks like tensorflow static_rnn will unroll LSTM(aka the tensorflow graph doesn't contain any control flow op such as merge, switch, loopcond ....). If so, the converter will not recognize it as a LSTM graph.

if you want converter to map tensorflow LSTM cell to ONNX LSTM, please try dynamic_rnn, and make sure tensorflow doesn't unroll it.

[fanghuaqi]

Could you make it possible to recognize such type of LSTM, I also see there is a test case for LSTM, which also used tf.contrib.rnn.BasicLSTMCell cell, thanks.

The difference is that it expose initial_state as input, and use static_rnn not dynamic_rnn cell.

Why I would like make this as a feature, since sometimes customer will only provide us pre-trained pb file to us, without any source code, it is also hard for them to change code to generate new pb file.

Thanks Huaqi

[zhijxu-MS]

yeah, this could be an enhancement, thanks for the suggestion. we will treat it an enhancement work item and contribution is warmly welcome!