TensorFlow：如何用一个张量索引另一个张量的切片

嗨，这个场景我太熟悉了！其实完全不需要搞复杂的3D掩码，用tf.gather_nd就能精准拿到每个序列的最后有效输出，和你想要的np.array([outputs[i, seq_lengths[i], :] for i in range(batch_size)])效果完全一致，咱们一步步来实现：

核心思路 #

因为序列是从0开始索引的，每个样本i的最后有效元素位置其实是seq_lengths[i]-1（比如实际长度是5，最后一个元素索引是4）。我们只需要构造一个包含(样本索引, 最后有效位置)的二维索引张量，再用tf.gather_nd从outputs里提取对应值就行。

具体代码实现 #

import tensorflow as tf

# 假设你的outputs形状是[16, 15000, 64]，seq_lengths是形状[16]的张量
batch_size = tf.shape(outputs)[0]

# 第一步：构造每个样本的索引（0到batch_size-1），转成[[0], [1], ..., [15]]的形状
indices_i = tf.expand_dims(tf.range(batch_size), axis=1)

# 第二步：构造每个样本的最后有效位置（seq_lengths减1），同样转成[[len0-1], [len1-1], ...]
indices_j = tf.expand_dims(seq_lengths - 1, axis=1)

# 第三步：拼接成[batch_size, 2]的索引张量，每一行对应一个样本的(i, 最后有效位置)
gather_indices = tf.concat([indices_i, indices_j], axis=1)

# 第四步：提取最后有效输出，结果形状是[16, 64]，正好可以直接连接全连接层
last_valid_outputs = tf.gather_nd(outputs, gather_indices)

额外说明 #

• 为什么不用掩码？因为这里我们只需要每个序列的最后一个有效元素，tf.gather_nd直接定位的方式比掩码更高效，掩码更多用于需要处理所有有效元素（比如计算损失时忽略填充部分）的场景。
• 这个方法完全兼容动态batch_size，哪怕你的batch_size不是固定的16，用tf.shape(outputs)[0]获取当前batch大小就没问题。

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