如何从tf.nn.dynamic_rnn中获取非填充序列的最后输出？

Great question! Let's break down your two questions clearly:

Absolutely not! When your input sequences are padded to match the longest sequence in the batch, output[:,-1] grabs the last time step's output of that longest sequence. For shorter sequences that were padded, this position corresponds to the output of padding tokens—not the actual last meaningful step of the original sequence. Using this will introduce invalid, noisy data into your model's downstream tasks, hurting both performance and result accuracy.

There are a couple of clean, efficient ways to do this, leveraging the seq_lengths tensor that stores each input's original sequence length:

Method 1: Use tf.gather_nd (most straightforward) #

This method constructs precise indices for each sample's last valid time step, then pulls the corresponding output directly:

# output shape: [batch_size, max_time, hidden_size]
# seq_lengths shape: [batch_size] (each value is the original length of the sequence)

# Create indices for each sample's last valid step
batch_indices = tf.range(tf.shape(output)[0])  # [0, 1, ..., batch_size-1]
last_step_indices = seq_lengths - 1  # convert to 0-based index
indices = tf.stack([batch_indices, last_step_indices], axis=1)

# Extract the last valid output for each sequence
last_valid_output = tf.gather_nd(output, indices)

Method 2: Use tf.boolean_mask (alternative approach) #

You can create a boolean mask that marks only the last valid step of each sequence, then filter the output with it:

# output shape: [batch_size, max_time, hidden_size]
# seq_lengths shape: [batch_size]

# Create a mask where only the last valid step of each sequence is True
mask = tf.one_hot(seq_lengths - 1, depth=tf.shape(output)[1], dtype=tf.bool)

# Extract and reshape to get the final output
last_valid_output = tf.boolean_mask(output, mask)
last_valid_output = tf.reshape(last_valid_output, [tf.shape(output)[0], tf.shape(output)[2]])

Both methods reliably retrieve the actual last meaningful output for each sequence, avoiding the padded steps entirely. Method 1 is generally preferred for its clarity and efficiency.

内容的提问来源于stack exchange，提问作者Hayk Sargsyan