获取TensorFlow中dynamic_rnn的最后输出 [英] Get the last output of a dynamic_rnn in TensorFlow

问题描述

我有一个形状为[batch, None, dim]的3D张量，其中第二维(即时间步长)是未知的.我使用dynamic_rnn处理此类输入，如以下代码片段所示:

I have a 3-D tensor of shape [batch, None, dim] where the second dimension, i.e. the timesteps, is unknown. I use dynamic_rnn to process such input, like in the following snippet:

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)

实际上，使用一些实际数字运行此代码，我得到了一些合理的结果:

Actually, running this snipped with some actual numbers, I have some reasonable results:

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张量?谢谢！

Is there a way to get a 3-D tensor of shape [batch, 1, hidden] with the last relevant output of the dynamic RNN? Thanks!

推荐答案

这是 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]的张量.

这篇关于获取TensorFlow中dynamic_rnn的最后输出的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！