pytorch实现求一批softmax结果top 2的差值diff。根据差值diff，如果我用一个阈值，diff大于0.5的为1，否者为0

代码如下：

import torch

# 生成一批概率值
batch_size = 3
num_classes = 5
probs = torch.randn(batch_size, num_classes)
probs = torch.softmax(probs, dim=1)

# 求出top 2的概率值和对应的下标
topk_probs, topk_idxs = torch.topk(probs, k=2, dim=-1)

# 计算差值diff
diff = topk_probs[:, 0] - topk_probs[:, 1]

# 根据阈值生成二值化的结果
threshold = 0.5
binary = torch.where(diff > threshold, torch.tensor(1), torch.tensor(0))
print(binary)

输出：

tensor([0, 0, 0])

解释：

假设生成的概率值为：

tensor([[0.2003, 0.0006, 0.1062, 0.0176, 0.6754],
        [0.3184, 0.2169, 0.0340, 0.2619, 0.1688],
        [0.0337, 0.3602, 0.0099, 0.0853, 0.5109]])

则求出的top 2概率值和下标分别为：

topk_probs: tensor([[0.6754, 0.2003],
                    [0.3184, 0.2619],
                    [0.5109, 0.3602]])

topk_idxs: tensor([[4, 0],
                   [0, 3],
                   [4, 1]])

因此，差值diff为：

diff: tensor([0.4751, 0.0565, 0.1507])

由于只有第一个差值大于0.5，因此生成的二值化结果为：

binary: tensor([0, 0, 0])