To pick the most confident data samples(top 3) per class based on the softmax output

Question

So I have trained a neural network on a particular dataset of 50 classes. During testing, I want to pick the best three samples(most confident ones, i.e: top 3 or top 5 samples) from each class in the testing dataset, How can I do that?

Example: If I have 50 classes, I want to pick the 3 best samples from the testing dataset(based on the softmax probabilities) for every class. Hence we will have 50*3 data samples picked. How can I do this in the most efficient way?

Thank you.

score 2 · Accepted Answer · answered Jan 17 '20 at 09:37

Assuming you have your final scores in a tensor probabilities of shape size of the test data × number of classes, I would do something like this:

best_scores, best_classes = probabilities.max(dim=1)

per_class_examples = []
for class_id in range(50):
    # mask telling where class_id class is
    class_positions = best_classes == class_id

    # make sure there are at least three examples,
    # if not, rather take less
    k = min(3, class_positions.sum())

    if k == 0:
        per_class_examples.append([])
    else:
        # set zero score to everything that is not class_id
        _, best_examples = torch.topk(best_scores * class_positions, k, dim=1)
        per_class_examples.append(best_examples)

To pick the most confident data samples(top 3) per class based on the softmax output

1 Answers1