基于多维数组计算多分类NLP模型的混淆矩阵

针对你的多分类场景，需要先把one-hot编码的真实标签和模型输出的多维预测分数转换成一维的类别索引，再计算混淆矩阵。下面是具体实现方案：

步骤1：转换多维数组为一维类别索引 #

处理真实标签（label_ids） #

你的result.label_ids是one-hot格式（每个样本对应长度为4的数组，正确类别为1），用np.argmax可以快速提取每个样本的真实类别索引（0-3）：

import numpy as np

# 转换one-hot标签为一维类别索引
true_labels = np.argmax(result.label_ids, axis=1)

处理模型预测结果（predictions） #

根据你的描述，模型预测值为正时对应预测类别为1，否则为0。但如果是标准多分类任务，模型输出的logits（未激活的分数）更适合用np.argmax取分数最高的类别作为预测（即使所有分数为负，也能得到模型最倾向的类别）：

# 方案1：按你要求的"正值为预测类别"提取索引
pred_labels = np.where(result.predictions > 0)[1]

# 方案2：更通用的多分类逻辑——取分数最高的类别（推荐）
pred_labels = np.argmax(result.predictions, axis=1)

步骤2：计算混淆矩阵 #

方法一：修改你原有的函数 #

适配转换后的一维类别索引，修改后的函数如下：

def compute_confusion_matrix(labels, true_onehot, pred_scores):
    # 转换为一维类别索引
    true_labels = np.argmax(true_onehot, axis=1)
    # 选择上面的方案1或方案2获取预测类别
    pred_labels = np.argmax(pred_scores, axis=1)  # 这里用推荐的方案2
    
    K = len(labels)
    conf_matrix = np.zeros((K, K), dtype=int)
    
    for t, p in zip(true_labels, pred_labels):
        conf_matrix[t][p] += 1
    
    return conf_matrix

# 使用示例：假设你的4个标签为['类别0', '类别1', '类别2', '类别3']
class_labels = ['类别0', '类别1', '类别2', '类别3']
conf_matrix = compute_confusion_matrix(class_labels, result.label_ids, result.predictions)

方法二：使用sklearn工具（更简洁） #

直接调用sklearn.metrics.confusion_matrix，无需自己实现循环：

from sklearn.metrics import confusion_matrix

true_labels = np.argmax(result.label_ids, axis=1)
pred_labels = np.argmax(result.predictions, axis=1)

conf_matrix = confusion_matrix(true_labels, pred_labels)

混淆矩阵解读 #

得到的conf_matrix是一个4×4的矩阵，其中conf_matrix[i][j]表示真实类别为i，被预测为j的样本数量，可以直接用来分析每个标签的分类准确率、召回率等指标。

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