如何从RoBERTa文本分类模型预测输出中提取Top3目标类别

你当前训练的是18分类文本模型，model.predict()返回结果的形状为(输入样本数, 18)，数组中每个位置的值对应该样本属于对应ID类别的概率，直接用numpy即可快速提取概率最高的3个分类，具体实现如下：

• 注意：你贴出的预测输出存在多行概率值，不符合单句输入的预期输出（单句输入应该仅返回1行共18个概率值），问题出在roberta_encode编码阶段没有给输入特征增加batch维度，后续可以给编码结果加np.expand_dims(feat_sent, axis=0)修正维度，否则拿到的预测结果是逐token的分类结果，不是整句的分类结果。
• 完整修改后的代码：

import numpy as np

model = None
with self.strategy.scope():
    # 加载模型与权重
    model = self.build_model(18)
    model.summary()
    model.load_weights("./saved_weights.h5")
    # 加载分词器、编码输入语句
    tokenizer = RobertaTokenizer.from_pretrained(self.MODEL_NAME)
    sentence = "entrepreneurs do you know these 4 golden rules for choosing your partner"
    feat_sent = self.roberta_encode(sentence, tokenizer)
    # 修正单句输入的维度问题：如果你的编码方法已经处理了batch维度，可以注释掉下面这行
    if len(feat_sent.shape) == 1:
        feat_sent = np.expand_dims(feat_sent, axis=0)
    
    # 执行预测
    pred_result = model.predict(feat_sent)
    # 取当前单条样本的概率数组
    sample_probs = pred_result[0]
    # 计算Top3分类：argsort返回从小到大排序的元素索引，取最后3位后倒序，得到按概率从高到低排列的分类ID
    top3_indices = np.argsort(sample_probs)[-3:][::-1]
    top3_probs = sample_probs[top3_indices]

    # 打印结果，如果你提前存了分类ID到类别名称的映射字典id2label，可以替换成真实分类名
    for cls_id, cls_prob in zip(top3_indices, top3_probs):
        print(f"分类ID：{cls_id}，预测概率：{cls_prob:.4f}")

• 逻辑说明：
  • np.argsort()不会修改原概率数组，只会返回数组元素从小到大排序对应的原位置索引，取切片[-3:]就能得到概率最高的3个类别对应的ID，再通过[::-1]倒序排列，即可实现按概率从高到低输出结果
  • 如果你提前保存了类别ID与实际业务分类名称的映射关系，例如id2label = {0:"创业指南", 1:"职场经验", ...}，打印时直接取id2label[cls_id]就能输出可读的分类名称。

内容的提问来源于stack exchange，提问作者o akka