将pytorch张量/ numpy数组的两个最高值变为1，将其他值变为0 [重复]

Question

我有一些 pytorch 张量（或 numpy 数组），想将两个最高的数字变为 1，将其他每个数字变为零。所以这个张量

tensor([0.9998, 0.9997, 0.9991, 0.9998, 0.9996, 0.9996, 0.9997, 0.9995],
   dtype=torch.float64)

应该变成这样：

tensor([1, 0, 0, 1, 0, 0, 0, 0],
   dtype=torch.float64)

有一些方法可以将最大数字变为 1，其他数字变为 0，但我需要将两个最大数字变为 1。为什么没有内置函数呢？我有一个分类问题，我知道两个对象属于第 1 类。所以这不应该是一个很常见的问题吗？

Answer 1

您可以使用topk：

x = tensor([0.9998, 0.9997, 0.9991, 0.9998, 0.9996, 0.9996, 0.9997, 0.9995], dtype=torch.float64)

_, idx = x.topk(2)

x.fill_(0)
x[idx] = 1