python - 获取 TensorFlow 中 dynamic_rnn 的最后输出

Question

我有一个形状为 [batch, None, dim] 的 3-D 张量，其中第二维(即时间步长)是未知的。我使用 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

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

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