I'm training an LSTM cell on batches of sequences that have different lengths. The tf.nn.rnn has the very convenient parameter sequence_length, but after calling it, I don't know how to pick the output rows corresponding the last time step of each item in the batch.
My code is basically as follows:
lstm_cell = tf.nn.rnn_cell.LSTMCell(num_lstm_units, input_size)
lstm_outputs, state = tf.nn.rnn(lstm_cell, input_list, dtype=tf.float32, sequence_length=sequence_lengths)
lstm_outputs is a list with the LSTM output at each time step. However, each item in my batch has a different length, and so I would like to create a tensor containing the last LSTM output valid for each item in my batch.
If I could use numpy indexing, I would just do something like this:
all_outputs = tf.pack(lstm_outputs)
last_outputs = all_outputs[sequence_lengths, tf.range(batch_size), :]
But it turns out that for the time begin tensorflow doesn't support it (I'm aware of the feature request).
So, how could I get these values?