如何在PyTorch中获取ImageNet预训练模型预测的前5高分类值？

针对你的需求，这里提供两种高效的方法来获取前5个最高概率值及对应类别索引：

方法1：使用torch.topk（推荐） #

torch.topk()是专门用于提取Top N元素的API，一步就能得到目标结果，比先排序再切片更高效：

# 计算softmax得到类别概率
out = F.softmax(out, dim=1)
# 获取前5个最高概率值和对应索引，dim=1指定在类别维度操作
top5_vals, top5_indices = torch.topk(out, k=5, dim=1, largest=True)

# 输出结果
print("前5个最高概率值:", top5_vals)
print("对应类别索引:", top5_indices)

方法2：基于torch.sort处理 #

如果你坚持使用torch.sort()，它返回的是包含排序后数值和对应索引的元组，你可以分别对两个张量切片取前5个：

out = F.softmax(out, dim=1)
# 降序排序，返回(排序后的概率张量, 对应类别索引张量)
sorted_vals, sorted_indices = torch.sort(out, descending=True)

# 提取前5个元素
top5_vals = sorted_vals[:, :5]
top5_indices = sorted_indices[:, :5]

# 输出结果
print("前5个最高概率值:", top5_vals)
print("对应类别索引:", top5_indices)

注意事项 #

• 确保你的输入张量input是批量维度在前的格式（例如[1, 3, 224, 224]，对应单张RGB图像），这样dim=1才会对应ImageNet的1000个类别维度。
• 若需要将索引转为真实类别名称，可以加载ImageNet的类别标签映射表，用索引匹配对应的名称。

内容的提问来源于stack exchange，提问作者Jose Ramon