如何使用tf.gather从二维张量中提取对应类别分数？

嘿，我懂你的需求啦！你想用tf.gather或tf.gather_nd从形状为(1000, 81)的概率矩阵里，根据每个样本的类别索引（形状(1000,)的class_ids）提取对应的分数，最终得到(1000,)的结果。

你之前用tf.gather_nd(probs, class_ids)没成功，是因为tf.gather_nd需要的是每个元素的完整坐标索引——对于二维张量来说，每个索引得是[行号, 列号]这样的二维数组，而你只传入了列号（class_ids），缺少了对应的行号，所以TensorFlow没法定位到正确的元素。

正确的实现方式 #

这里给你两种高效的向量化方案，比你写的循环方法性能好得多（尤其在大batch下）：

方法1：用tf.gather_nd构造完整索引

我们需要给每个样本配对上它的行索引（从0到999），然后把行索引和class_ids拼接成形状为(1000, 2)的索引矩阵，再传入tf.gather_nd：

# 生成行索引，形状为(1000,)
row_indices = tf.range(tf.shape(probs)[0], dtype=tf.int32)
# 将行索引和class_ids拼接成(1000, 2)的索引张量
indices = tf.stack([row_indices, class_ids], axis=1)
# 提取对应位置的分数
class_scores = tf.gather_nd(probs, indices)

执行后class_scores的形状就是(1000,)，完全符合你的期望。

方法2：用tf.batch_gather（更简洁）

如果你用的是较新的TensorFlow版本，tf.batch_gather专门用来处理这种按batch提取元素的场景，它会自动匹配每个batch的维度：

# 把class_ids扩展成(1000, 1)的形状，符合batch_gather的输入要求
class_ids_expanded = tf.expand_dims(class_ids, axis=1)
# 按batch提取每个样本对应索引的分数
class_scores = tf.batch_gather(probs, class_ids_expanded)
# 最后去掉多余的维度，得到(1000,)的结果
class_scores = tf.squeeze(class_scores, axis=1)

补充说明 #

你之前写的循环方法虽然能得到正确结果，但TensorFlow的计算图优化对循环支持不好，大batch下会很慢，推荐用上面的向量化操作，既高效又符合TensorFlow的最佳实践。

另外，如果你的需求只是获取每个样本的最高类别分数，其实可以直接用tf.reduce_max(probs, axis=1)，一步到位得到(1000,)的结果——不过既然你已经有了class_ids（比如后续还要用到这个索引），那上面的两种提取方法就更合适啦。

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