如何获得图像分类器中的前 5 个预测?

社会演员多 pytorch 433

原文标题How can i get top 5 prediction in Image Classifier?

这段代码给了我前 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())

原文链接:https://stackoverflow.com//questions/71450451/how-can-i-get-top-5-prediction-in-image-classifier

回复

我来回复
  • Kevin的头像
    Kevin 评论

    softmax 给出了类的概率分布。 argmax 只取概率最高的类的索引。您也许可以使用 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]])
    

    如您所见,上面输出中的最后一列是概率最高的预测,第二列是概率第二高的预测,依此类推。

    2年前 0条评论