在 Tensorflow 中检索 LSTM 序列的最后一个值

Question

我有不同长度的序列，我想在 Tensorflow 中使用 LSTM 对其进行分类。对于分类，我只需要每个序列的最后一个时间步的 LSTM 输出。

max_length = 10
n_dims = 2
layer_units = 5
input = tf.placeholder(tf.float32, [None, max_length, n_dims])
lengths =  tf.placeholder(tf.int32, [None])
cell = tf.nn.rnn_cell.LSTMCell(num_units=layer_units, state_is_tuple=True)

sequence_outputs, last_states = tf.nn.dynamic_rnn(cell, sequence_length=lengths, inputs=input)

我想用 NumPy 表示法得到：output = sequence_outputs[:,lengths]

是否有任何方法或解决方法可以在 Tensorflow 中获得此行为？

---更新---

按照How to select rows from a 3-D Tensor in TensorFlow? 这篇帖子，似乎可以使用tf.gather 并操纵索引以有效的方式解决问题。唯一的要求是必须事先知道批量大小。以下是引用帖子对这个具体问题的改编：

max_length = 10
n_dims = 2
layer_units = 5
batch_size = 2
input = tf.placeholder(tf.float32, [batch_size, max_length, n_dims])
lengths =  tf.placeholder(tf.int32, [batch_size])
cell = tf.nn.rnn_cell.LSTMCell(num_units=layer_units, state_is_tuple=True)

sequence_outputs, last_states = tf.nn.dynamic_rnn(cell,
                                                  sequence_length=lengths, inputs=input)

#Code adapted from @mrry response in StackOverflow:
#https://stackoverflow.com/questions/36088277/how-to-select-rows-from-a-3-d-tensor-in-tensorflow
rows_per_batch = tf.shape(input)[1]
indices_per_batch = 1

# Offset to add to each row in indices. We use `tf.expand_dims()` to make 
# this broadcast appropriately.
offset = tf.range(0, batch_size) * rows_per_batch

# Convert indices and logits into appropriate form for `tf.gather()`. 
flattened_indices = lengths - 1 + offset
flattened_sequence_outputs = tf.reshape(self.sequence_outputs, tf.concat(0, [[-1],
                             tf.shape(sequence_outputs)[2:]]))

selected_rows = tf.gather(flattened_sequence_outputs, flattened_indices)
last_output  = tf.reshape(selected_rows,
                          tf.concat(0, [tf.pack([batch_size, indices_per_batch]),
                          tf.shape(self.sequence_outputs)[2:]]))

@petrux 选项 (Get the last output of a dynamic_rnn in TensorFlow) 似乎也可以工作，但在 for 循环中构建列表的需求可能没有优化，尽管我没有执行任何基准测试来支持此语句。

Answer 1

This 可能是一个答案。我不认为有什么类似于你指出的 NumPy 表示法，但效果是一样的。

Answer 2

这是一个解决方案，使用gather_nd，不需要提前知道批量大小。

def extract_axis_1(data, ind):
    """
    Get specified elements along the first axis of tensor.
    :param data: Tensorflow tensor that will be subsetted.
    :param ind: Indices to take (one for each element along axis 0 of data).
    :return: Subsetted tensor.
    """

    batch_range = tf.range(tf.shape(data)[0])
    indices = tf.stack([batch_range, ind], axis=1)
    res = tf.gather_nd(data, indices)

    return res

output = extract_axis_1(sequence_outputs, lengths - 1)

现在output 是一个维度为[batch_size, num_cells] 的张量。