多分类文本分类：无循环高效获取Top3预测概率及对应类别

嘿，我来帮你搞定这个需求！你现在已经能拿到单条样本的最高概率类别了，要获取前3个概率最高的结果及其对应类别，完全不需要用循环，借助NumPy的向量操作就能高效完成，这里有两种实用方法供你选：

方法1：用argsort（简单直观） #

你之前已经用到了argsort，其实只需要调整索引，取每个样本排序后的最后3个位置就行——因为argsort默认是升序排列，末尾的就是概率最大的：

import numpy as np
from sklearn.metrics import accuracy_score

# 先获取所有类别的概率矩阵
probabilities = classifier.predict_proba(X_test)

# 对每个样本的概率进行升序排序，得到索引矩阵
order = np.argsort(probabilities, axis=1)

# 取每个样本最后3个索引（对应概率最高的3个类别），倒序后让概率从高到低排列
top3_indices = order[:, -3:][:, ::-1]

# 获取对应的类别
top3_classes = classifier.classes_[top3_indices]

# 获取对应的概率值，同样倒序保证从高到低
top3_probs = probabilities[np.arange(probabilities.shape[0])[:, None], top3_indices]

# 验证最高概率类别的准确率（和你原来的逻辑一致）
top1_classes = classifier.classes_[order[:, -1]]
print(accuracy_score(top1_classes, y_test))

方法2：用argpartition（更高效，适合类别数多的场景） #

如果你的分类任务类别数量很大，argsort会对所有类别概率全排序，有点浪费算力。argpartition只需要把前k个最大的元素放到对应位置，不需要全排序，效率会更高：

import numpy as np
from sklearn.metrics import accuracy_score

probabilities = classifier.predict_proba(X_test)
k = 3

# 找到每个样本中概率最大的k个元素的索引
top3_indices_part = np.argpartition(-probabilities, k, axis=1)[:, :k]

# argpartition不会给这k个元素排序，所以对每个样本的这k个概率再排序，保证从高到低
top3_sorted_indices = np.argsort(-probabilities[np.arange(probabilities.shape[0])[:, None], top3_indices_part], axis=1)

# 得到最终按概率从高到低排列的索引
top3_indices = top3_indices_part[np.arange(probabilities.shape[0])[:, None], top3_sorted_indices]

# 获取对应的类别和概率
top3_classes = classifier.classes_[top3_indices]
top3_probs = probabilities[np.arange(probabilities.shape[0])[:, None], top3_indices]

# 验证top1准确率
top1_classes = classifier.classes_[top3_indices[:, 0]]
print(accuracy_score(top1_classes, y_test))

小提示 #

• 两种方法都是纯向量操作，没有循环，效率拉满
• 如果类别数不多（比如10个以内），argsort和argpartition性能差不多，选哪个都行；如果类别数很多（比如上百个），优先选argpartition
• 直接打印top3_classes和top3_probs就能看到结果，每个样本对应一个长度为3的数组，分别是概率从高到低的类别和对应概率值

内容的提问来源于stack exchange，提问作者Poornesh V