获取TensorFlow中dynamic_rnn的最后输出

Question

我有一个3D张量形状[batch, None, dim]，其中第二维（即时间步长）未知。我使用dynamic_rnn以下代码段来处理此类输入：

import numpy as np
import tensorflow as tf

batch = 2
dim = 3
hidden = 4

lengths = tf.placeholder(dtype=tf.int32, shape=[batch])
inputs = tf.placeholder(dtype=tf.float32, shape=[batch, None, dim])
cell = tf.nn.rnn_cell.GRUCell(hidden)
cell_state = cell.zero_state(batch, tf.float32)
output, _ = tf.nn.dynamic_rnn(cell, inputs, lengths, initial_state=cell_state)

实际上，以一些实际数字运行此代码后，我得到了一些合理的结果：

inputs_ = np.asarray([[[0, 0, 0], [1, 1, 1], [2, 2, 2], [3, 3, 3]],
                    [[6, 6, 6], [7, 7, 7], [8, 8, 8], [9, 9, 9]]],
                    dtype=np.int32)
lengths_ = np.asarray([3, 1], dtype=np.int32)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    output_ = sess.run(output, {inputs: inputs_, lengths: lengths_})
    print(output_)

输出为：

[[[ 0.          0.          0.          0.        ]
  [ 0.02188676 -0.01294564  0.05340237 -0.47148666]
  [ 0.0343586  -0.02243731  0.0870839  -0.89869428]
  [ 0.          0.          0.          0.        ]]

 [[ 0.00284752 -0.00315077  0.00108094 -0.99883419]
  [ 0.          0.          0.          0.        ]
  [ 0.          0.          0.          0.        ]
  [ 0.          0.          0.          0.        ]]]

有没有一种方法可以通过动态RNN[batch, 1, hidden]的 最后一个相关输出 获得形状的3-D张量？谢谢！

Answer 1

这就是collect_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(output, lengths - 1)

现在output是维的张量[batch_size, num_cells]。