我已经在一个包含50个类别的特定数据集上训练了一个神经网络。在测试过程中，我想从测试数据集中为每个类别挑选出最好的三个样本（即最有信心的样本，前3名或前5名）。我该如何做到这一点？

举例来说，如果我有50个类别，我想从测试数据集中为每个类别挑选出3个最佳样本（基于softmax概率）。因此，我们将挑选出50*3个数据样本。我怎样才能以最有效的方式做到这一点？

谢谢。

回答：

假设你已经将最终分数存储在一个名为probabilities的张量中，其形状为测试数据大小 × 类别数量，我会这样做：

best_scores, best_classes = probabilities.max(dim=1)per_class_examples = []for class_id in range(50):    # 标记class_id类别的位置    class_positions = best_classes == class_id    # 确保至少有三个样本，如果没有，则取更少的    k = min(3, class_positions.sum())    if k == 0:        per_class_examples.append([])    else:        # 将非class_id类别的分数设为零        _, best_examples = torch.topk(best_scores * class_positions, k, dim=1)        per_class_examples.append(best_examples)