python - 如何在pytorch中实现tf.nn.top_k

Question

我想在 pytorch 中实现 tf.nn.in_top_k。这是tf.nn.in_top_k的链接，

tf.math.in_top_k(
    targets, predictions, k, name=None
)

它将 k 处的精度计算为 bool Tensor，并将返回 bool 类型的 Tensor。

tf.nn.in_top_k

不知道pytorch中有没有类似的api？

Answer 1

据我所知，pytorch 中没有内置等效的 in_top_k 函数。写一个比较简单。例如

def in_top_k(targets, preds, k):
    topk = preds.topk(k)[1]
    return (targets.unsqueeze(1) == topk).any(dim=1)