如何在图像分类器中获得前5名预测

Question

这段代码给了我前1名的预测，但我想要前5名。我该怎么做？

# Get top 1 prediction for all images

predictions = []
confidences = []

with torch.inference_mode():
for _, (data, target) in enumerate(tqdm(test_loader)):
data = data.cuda()
target = target.cuda()
output = model(data)
pred = output.data.max(1)[1]
probs = F.softmax(output, dim=1)
predictions.extend(pred.data.cpu().numpy())
confidences.extend(probs.data.cpu().numpy())

Answer 1

softmax给出了类上的概率分布。CCD_ 1只取具有最高概率的类的索引。您也许可以使用argsort，它将返回所有处于排序位置的索引。

一个例子：

a = torch.randn(5, 3)
preds = a.softmax(1); preds

输出：

tensor([[0.1623, 0.6653, 0.1724],
[0.4107, 0.1660, 0.4234],
[0.6520, 0.2354, 0.1126],
[0.1911, 0.4600, 0.3489],
[0.4797, 0.0843, 0.4360]])

这可能是具有3个目标的5号批次的概率分布。沿着最后一个维度的argmax将给出：

preds.argmax(1)

输出：

tensor([1, 2, 0, 1, 0])

而沿着最后一个维度的argsort将给出：

preds.argsort(1)

输出：

tensor([[0, 2, 1],
[1, 0, 2],
[2, 1, 0],
[0, 2, 1],
[1, 2, 0]])

正如你所看到的，上面输出中的最后一列是概率最高的预测，第二列是概率第二高的预测，依此类推