How can i get top 5 prediction in Image Classifier?

Question

this code gives me top 1 prediction but i want top 5. how can i do that?

   # 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

The softmax gives an probability distribution over the classes. argmax only takes the index of the class with the highest probability. You could perhaps use argsort which will return all indices in their sorted position.

An example:

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

output:

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]])

This is might be the probability distribution for a batch of size 5 with 3 targets. An argmax along the last dimension will give:

preds.argmax(1)

output:

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

Whilst an argsort along the last dimension will give:

preds.argsort(1)

output:

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

As you can see, the last column in the output above is the prediction with highest probability and the second column is the prediction with the second highest probability, and so on.