pytorch实现求一组softmax结果top 2的差值

假设我们有一个大小为(batch_size, num_classes)的输出张量output，表示模型对于每个类别的预测概率。我们可以使用pytorch的softmax函数将其转换为概率分布，然后使用topk函数找到每个样本的top 2类别及其对应的概率。

最后，我们可以使用torch.gather函数将这些概率与其对应的类别标签匹配，然后计算top 2概率的差值并返回结果。

下面是一个示例代码：

import torch

# 假设我们有10个类别
num_classes = 10

# 生成随机输出
batch_size = 16
output = torch.randn(batch_size, num_classes)

# 计算softmax概率分布
probs = torch.softmax(output, dim=1)

# 找到top 2类别及其对应的概率
topk_probs, topk_indices = torch.topk(probs, k=2)

# 找到每个样本的top 2类别标签
labels = torch.arange(num_classes).unsqueeze(0).expand(batch_size, num_classes)
topk_labels = torch.gather(labels, dim=1, index=topk_indices)

# 计算top 2概率的差值
diff = topk_probs[:, 0] - topk_probs[:, 1]

print(diff)

输出结果是一个大小为(batch_size,)的张量，表示每个样本的top 2概率差值。