tensorflow - 从 tensorflow 中的 2dim 张量收集值

Question

嗨，这里是 tensorflow 初学者...我正在尝试获取 2 dim 张量中某些元素的值，在我的例子中是来自概率矩阵的类分数。

概率矩阵为 (1000,81)，批量大小为 1000，类数为 81。ClassIDs 为 (1000,)，包含每个样本的最高类分数的索引。如何使用 tf.gather 从概率矩阵中获取相应的类别分数？

class_ids = tf.cast(tf.argmax(probs, axis=1), tf.int32)  
class_scores = tf.gather_nd(probs,class_ids)

class_scores 应该是形状为 (1000,) 的张量，其中包含每个样本的最高 class_score。

现在我正在使用如下所示的解决方法:

class_score_count = []
for i in range(probs.shape[0]):
    prob = probs[i,:]
    class_score = prob[class_ids[i]]
    class_score_count.append(class_score)
class_scores = tf.stack(class_score_count, axis=0)

感谢您的帮助!

Answer 1

您可以使用 tf.gather_nd 来完成像这样:

class_ids = tf.cast(tf.argmax(probs, axis=1), tf.int32)
# If shape is not dynamic you can use probs.shape[0].value instead of tf.shape(probs)[0]
row_ids = tf.range(tf.shape(probs)[0], dtype=tf.int32)
idx = tf.stack([row_ids, class_ids], axis=1)
class_scores = tf.gather_nd(probs, idx)

你也可以只使用 tf.reduce_max ，即使它实际上会再次计算最大值，但如果您的数据不是太大，它可能不会慢很多:

class_scores = tf.reduce_max(probs, axis=1)