☰
PyTorch混淆矩阵绘制全指南:从原理到出版级可视化
2026/10/22 22:08:32 网站建设 项目流程

从项目标题来看,这是一篇“技术复现型”的分享。只要动过PyTorch做图像分类或文本分类的人,应该都经历过这种场景:训练完模型,打印出准确率,看到一个90%甚至95%的数字,心里挺高兴,结果模型一上线,发现某些类别几乎全错。准确率这个指标太“宏观”了,它会掩盖很多细节。所以我才坚持每次跑完分类任务,都要画一张混淆矩阵出来看看。这篇文章就把我平时在PyTorch项目里画混淆矩阵和做可视化的一套完整流程整理出来,从最底层的矩阵计数原理,到最终的出版级图像导出,尽量覆盖到每一个坑。

1. 为什么要画混淆矩阵:从准确率说不清的问题说起

1.1 准确率看着挺高,模型却可能在“装模作样”

先讲个我自己踩过的案例。当时做一个六分类的纹理识别任务,训练集里样本分布很不均衡,其中一类占了总样本的将近一半。训练结束后验证集准确率到了92%,看起来一切正常。结果把测试集结果按类别拆开一统计,那个占了50%的类别准确率接近99%,而另外几个稀少类别准确率只有60%到70%。如果只看准确率,根本发现不了这个问题。混淆矩阵的价值就在这:它不给你一个笼统的“好”或“坏”,而是把“哪一类被认对了、哪一类被认成了哪一类”全部摊开。

有些读者可能会问,精确率、召回率、F1值不是也能看类别的单独表现吗?确实能。但混淆矩阵是最原始的“计数表”,精确率和召回率其实是基于这个计数表二次计算出来的。有了混淆矩阵,你可以自己计算任何想要的指标,还可以定位到具体的“错误对”,比如A类被错认成B类到底是A类的特征太接近B类,还是数据标注本身就有问题。这是单独一个F1值做不到的。

1.2 混淆矩阵到底在告诉你什么

对于一个N分类任务,混淆矩阵是一个N乘N的方阵。行通常代表真实标签,列代表预测标签。矩阵的第i行第j列的值,表示“真实类别是i,却被模型预测成类别j”的样本个数。对角线上的值越大,说明这一类被正确识别的比例越高。对角线外的值就是错误分布,它们会清晰地告诉你错误主要发生在哪些类别之间。

举个例子,三分类问题假设有150个测试样本,每个类别各50个。如果矩阵长这样:

真实\预测类别0类别1类别2
类别04811
类别15405
类别22840

对角线加起来是128,所以准确率是85.3%。但这张矩阵透露的信息更多:类别0表现最好,只有两个样本出错;类别1有5个样本被认成类别0,5个被认成类别2;类别2则有8个样本被认成类别1。这告诉你类别1和类别2之间的混淆比较严重,可能是特征相似,也可能是训练样本不足。

1.3 什么时候必须画混淆矩阵

我个人的习惯是:只要是分类任务,无论是图像分类、文本分类还是音频分类,跑了测试集之后都要画一张混淆矩阵。不是说每个项目都要给客户看,而是你自己必须心里有数。尤其是以下几种情况,不画混淆矩阵几乎等于白跑测试:

第一,类别不均衡时。就像前面举的例子,准确率会被大类“带节奏”。第二,类别数量较多时,比如几十个类,即使每个类别准确率在大数上看起来差不多,也可能存在个别类别之间系统性错认,只有矩阵才能暴露这种规律。第三,迭代优化过程中。每次调完参、加完数据,画一张矩阵对比两次的错误分布,能直观看出优化是否真的有效,而不是只看准确率有没有涨。第四,需要向团队或客户汇报时。一张带颜色深浅的混淆矩阵图,比一堆数字直观得多。

2. 动手前的准备:环境、模型输出与数据格式

2.1 最小可用环境清单

画混淆矩阵本身对硬件没有任何要求,CPU机器完全能跑。核心依赖就是PyTorch、scikit-learn、matplotlib,如果想画得更精致一点,可以再加seaborn。我用过的几个版本组合供参考:

依赖库我常用的版本区间说明
Python3.8 ~ 3.11老版本要注意scikit-learn兼容性
PyTorch1.13 或 2.x分类模型推理
scikit-learn1.0 ~ 1.3提供confusion_matrix等函数
matplotlib3.5 ~ 3.8基础绘图
seaborn0.12 ~ 0.13热力图可视化更省事

安装命令我就不多写了,pip install torch scikit-learn matplotlib seaborn一步到位。如果是在已有的训练环境里补装,注意不要随意升级PyTorch版本,避免影响之前训练好的模型权重加载。

2.2 模型输出到底长什么样:logits、概率和预测类别

刚入门的朋友最容易在这里搞混。PyTorch分类模型在全连接层之后通常输出一个向量,这个向量在进入Softmax之前叫logits,维度是批量大小乘以类别数。举例来说,一个批次的输入经过模型后,输出形状通常是[batch_size, num_classes]。对每个样本来说,这个向量里的值表示该样本属于各个类别的“原始得分”,数值越大代表模型越倾向于这个类别。取最大值的索引,就是预测类别。

有一种常见操作是先把logits经过Softmax转成概率,再取最大值索引。这两者其实结果完全一致,因为Softmax是单调递增函数,不会改变最大值的位置。所以如果你只是要取预测类别,直接做pred = torch.argmax(logits, dim=1)就够了,不需要先算Softmax。但如果你想把概率值也记录下来,比如用于分析置信度,那就需要先过Softmax了。

2.3 拿到预测结果后先做这件事:类别索引与名称对应关系检查

模型输出的是索引,比如数字0、1、2,而不是字符串“猫”“狗”“鸟”。所以动手画图前,先确认索引和类别名称的对应关系。我遇到过不止一次,训练时数据集是用字典映射的,类别名称在数据集类的classes属性里,导出模型做推理时忘了保存映射关系,结果画出来的矩阵行列全是数字,根本不知道哪一行对应哪个类别。

这里给一个最小验证片段,确保测试集的类别是连续且完整的:

# 以 torchvision.datasets.ImageFolder 加载的数据为例 dataset_classes = test_dataset.classes print(dataset_classes) # 例如 ['cat', 'dog', 'bird']

如果你的数据是自定义Dataset,一定要检查__getitem__返回的标签是从0开始的连续整数。如果标签不连续,比如0、1、5、9,混淆矩阵的行列就会错位。这一点后面还会详细说。

3. 零基础到手写实现:混淆矩阵的两种构建方式

3.1 方式一:手写计数,彻底搞懂矩阵逻辑

虽然现成库一行就能搞定,但我强烈建议新手先手写一次。理解了矩阵填充逻辑,后面用库函数出现问题时,你才有能力排查。

首先准备真实标签列表y_true和预测标签列表y_pred,两者必须是等长的Python列表、NumPy数组或PyTorch张量。假设有4个类别,先初始化一个4乘4的全零矩阵,然后遍历每一对真实标签和预测标签,把对应位置加1:

import torch def build_confusion_matrix(y_true, y_pred, num_classes): matrix = torch.zeros(num_classes, num_classes, dtype=torch.int64) for t, p in zip(y_true, y_pred): matrix[t, p] += 1 return matrix # 示例数据:5个样本,3个类别 y_true = [0, 1, 2, 1, 0] y_pred = [0, 2, 2, 1, 0] cm = build_confusion_matrix(y_true, y_pred, num_classes=3) print(cm)

输出:

tensor([[2, 0, 0], [0, 1, 1], [0, 0, 1]])

解释一下:第一个样本真实是0预测是0,所以matrix[0,0]加1;第二个样本真实是1预测是2,所以matrix[1,2]加1。这个逻辑非常简单,遍历完所有样本就得到了完整的计数矩阵。如果类别数很多,用循环加法的效率不高,但画图前统计测试集一般也就几千个样本,性能完全不是问题。

3.2 方式二:使用适合快速统计的现成函数

手写矩阵只是教学手段,实际项目里我一般直接用机器学习常用库中现成的混淆矩阵统计函数。它省事、代码少、不易出错:

from sklearn.metrics import confusion_matrix y_true = [0, 1, 2, 1, 0] y_pred = [0, 2, 2, 1, 0] cm = confusion_matrix(y_true, y_pred, labels=[0, 1, 2]) print(cm)

重点是这个labels参数。它用来指定类别索引的完整列表,顺序就是矩阵行列的顺序。如果你的测试集某个类别恰好一个样本都没出现,不传labels的话,矩阵会缺少那一行一列,导致维度对不上。我一般在多分类任务中都会显式传入labels=list(range(num_classes)),强制让矩阵维度固定。

如果真实标签和预测标签是PyTorch张量,先转成Python列表或NumPy数组再传入。

# 从 GPU 张量转到 CPU,再转 numpy y_true_np = y_true.cpu().numpy() y_pred_np = y_pred.cpu().numpy() cm = confusion_matrix(y_true_np, y_pred_np, labels=list(range(num_classes)))

这里我解释一下为什么要先转cpu().numpy()。confusion_matrix函数接收的是类数组结构,PyTorch张量虽然也能在某些情况下被兼容,但为避免type error或device对不上的问题,统一转成NumPy数组最省心。

3.3 手工方式 vs 现成函数:怎么选

有人可能会觉得,既然现成函数这么方便,为什么还要学手写?因为现成函数帮我们省掉的是“计数”这一步,但没有帮我们理解“矩阵的含义”。在实际问题排查中,你需要跟别人解释为什么某个位置数值异常偏高,这时候只有真正理解矩阵每一格的含义,才能基于业务逻辑给出合理解释。两种方式建议这样搭配:学习和调试阶段手写,理解透彻后正式代码里用现成函数。

还有一个细节:手写矩阵用for循环效率较低,如果测试集很大,比如几万甚至几十万样本,建议用torch.zeros和index_put_这类批量操作,或者直接用NumPy的np.add.at。不过对常规项目来说,这个复杂度根本遇不到。

4. 画出真正能用的混淆矩阵图:配色、归一化与标注

4.1 基础热力图:让数字变成颜色

矩阵本身是数字表格,直接看也能看,但远不如热力图直观。热力图的核心思想是用颜色深浅表示数值大小:颜色越深,数值越大。对角线越深,说明分类效果越好。

用seaborn画热力图是我用得最多的方式,因为它的heatmap函数几乎为混淆矩阵量身定制:

import matplotlib.pyplot as plt import seaborn as sns import numpy as np from sklearn.metrics import confusion_matrix # 假设已经得到 cm 和类别名称 class_names plt.figure(figsize=(8, 6)) sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=class_names, yticklabels=class_names, cbar=False) plt.xlabel('Predicted label') plt.ylabel('True label') plt.title('Confusion Matrix') plt.tight_layout() plt.savefig('confusion_matrix.png', dpi=300) plt.show()

参数说明:

  • annot=True:在每个格子里显示数字。
  • fmt='d':因为混淆矩阵是整数计数,用d表示十进制整数格式。如果错用了fmt='.2f'或者不写,可能报错或者显示成小数。
  • cmap='Blues':蓝色渐变,颜色越深数值越大,这是论文里最常见的配色。
  • xticklabels和yticklabels:传入类别名称列表,让坐标轴显示可读的名称而不是数字索引。
  • cbar=False:我比较喜欢关掉右侧的颜色条,因为它只表示颜色深浅对应数值大小,但每个格子已经标了数字,颜色条的信息量有限,还占空间。如果对颜色对比敏感,可以保留。

4.2 归一化:从计数到比例的转换

原始计数矩阵有个问题:如果每个类别的样本数不同,深色的格子分布会失真。例如类别0有1000个测试样本,类别1只有100个,即使两者的识别率都是80%,类别0对应的对角线格子数值是800,颜色会非常深,类别1只有80,看起来很浅,给人造成“类别1表现很差”的错觉。解决方法是按行归一化,归一化后每个格子的值表示“该行真实类别中有多少比例的样本被预测到了这一列”。

归一化矩阵的计算方法:把每一行的所有值除以该行的总和。在scikit-learn中可以直接指定normalize='true':

cm_normalized = confusion_matrix(y_true, y_pred, labels=list(range(num_classes)), normalize='true')

在seaborn绘制时,我通常这样处理:

cm = confusion_matrix(y_true, y_pred, labels=list(range(num_classes))) cm_norm = cm.astype('float') / cm.sum(axis=1, keepdims=True) sns.heatmap(cm_norm, annot=True, fmt='.2f', cmap='Blues', xticklabels=class_names, yticklabels=class_names)

注意归一化后fmt要改成'.2f',显示两位小数。在最终展示时,我更倾向于同时展示两种形式:一是原始计数矩阵,用来精确知道样本数量;二是按行归一化的比例矩阵,用来比较不同类别之间的表现。可参考的展示方案是:先打印原始计数矩阵,再打印归一化矩阵,图里选一种呈现。

4.3 坐标轴字体、标题和图像导出细节

可视化不只是把图画出来,还要考虑最终使用场景。如果图是要放进论文、技术报告或PPT里的,要注意以下几点。

首先,类别名称如果比较长,比如“background_with_noise”这种,横坐标文字会重叠。解决方法是把坐标轴文字旋转一下:

plt.xticks(rotation=45, ha='right') plt.yticks(rotation=0)

旋转角度一般45度就够,ha='right'让文字右对齐,避免和相邻文字贴在一起。

其次,字体大小要适配图片尺寸。默认字号在8乘6英寸的图里还行,但如果类别很多,比如30个类别,格子很小,数字会挤在一起。这时候要么调大figsize,要么调小字体:

sns.set(font_scale=1.2) # 全局放大字体 # 或者单独控制 annot 字体大小 sns.heatmap(..., annot_kws={'size': 10})

最后是导出格式。我会优先导出一个矢量图加一个位图。矢量图用PDF或SVG,用于论文或报告,缩放不失真;位图用PNG,dpi=300以上,用于网页展示或微信群聊。示例代码:

plt.savefig('confusion_matrix.pdf', format='pdf', bbox_inches='tight') plt.savefig('confusion_matrix.png', format='png', dpi=300, bbox_inches='tight')

bbox_inches='tight'会自动裁剪图片边缘空白,保证整张图紧凑,这一步我几乎每次都会加。

5. 验证集、测试集和全量推理:数据到底用哪一份

5.1 三份数据的分工

不少初学者会把训练过程中打印的验证集准确率对应的混淆矩阵当作最终结果,这其实不太严谨。验证集的作用是调参和选模型,它参与了你对超参数的选择过程,所以用它评估模型会略微偏乐观。真正用于最终性能报告的应该是测试集,也就是训练过程中从未参与梯度更新和模型选择的样本。如果连测试集你也反复用来做决策,它也会慢慢“污染”。

所以我的习惯是:训练过程中的混淆矩阵用验证集画,用于观察每个epoch有没有个别类别开始严重错认;最终汇报和论文里的混淆矩阵用测试集画,用于说明模型的真实泛化能力。

5.2 推理时别用模型本身的Dropout和BatchNorm坑

在PyTorch中,如果模型处于训练模式(model.train()),Dropout和BatchNorm的行为与推理模式不同。画混淆矩阵时,必须把模型切换到评估模式:

model.eval() with torch.no_grad(): outputs = model(batch_data) preds = torch.argmax(outputs, dim=1)

如果不加torch.no_grad(),推理过程会构建计算图,不仅占用大量显存,速度也慢,画个混淆矩阵把显卡显存跑爆就尴尬了。model.eval()则能保证BatchNorm使用训练阶段统计的均值和方差,Dropout关闭,这样推理结果才稳定可控。有一次我忘了加model.eval(),连续跑两次同一条测试数据得出的预测结果不一样,吓了自己一跳,最后排查发现是Dropout在作怪。

5.3 数据顺序打乱对矩阵没有影响,但对可视化有影响

混淆矩阵本质上是对所有样本计数的汇总,不关心样本顺序。所以DataLoader里是否shuffle=True不影响矩阵结果。但要注意,如果你同时还想输出“哪些具体样本被分错了”,以便人工检查错误案例,那就必须记录样本路径或ID,而不能只记预测结果。顺序一旦打乱,你只知道有几个样本被错认,却不知道具体是哪几张图。

我通常会在推理时把样本的路径和预测结果一起记录下来,便于后续人工核对。这步在图像分类里特别重要,因为矩阵只能告诉你错误数量,不能告诉你错误原因。

6. 从矩阵到洞察:进阶分析与实战技巧

6.1 发现系统性错误:相邻类别的强烈混淆

混淆矩阵一个特别有用的场景是发现“系统性错认”。比如在植物叶片分类中,两种外观极其相似的病害类型常常互相混淆,矩阵里会看到两个对称的非对角线位置数值都很高。这说明模型并没有学会区分它们,本质上是特征表达不够,或者两类样本之间的差异确实微小。

发现这种规律后,我的处理步骤通常是这样:先统计这两个类别的样本量,确认是否不平衡;然后检查训练数据里这两个类别的图像是否采集条件差异很大,比如光照、角度;再考虑是否需要对这两类做特征可视化(比如用梯度加权类激活映射)看模型关注区域是否有问题。

6.2 类别不平衡如何让矩阵“说谎”

归一化矩阵能解决一部分“比较”问题,但解决不了根本问题:样本量太少的类别,即使归一化后准确率看着还行,实际可信度也不高。比如某个类别只有10个测试样本,哪怕只有1个被分错,归一化后准确率就只有90%,而另一个类别有1000个样本,50个分错,准确率是95%。单看这两个数字,好像后者更好,但样本量差异导致前者的置信区间大得多。

这种情况下,我建议在矩阵图下方额外标注每个类别的测试样本数。代码上可以准备一个样本数列表,用plt.text在对应的列下方标注,或者干脆在类别名称里带上样本量,比如“类别A (n=100)”。

6.3 结合分类报告一起打包输出

混淆矩阵是基础,但给别人看的时候,通常不能只甩一张矩阵图。我在项目汇报时一般会同时输出一份分类报告,包含每个类别的精确率、召回率、F1值和样本数。这样既能从矩阵看全局,又能从指标看单类表现,论证更充分。

from sklearn.metrics import classification_report report = classification_report(y_true_np, y_pred_np, target_names=class_names, digits=3) print(report)

这段代码会输出一个表格,每个类一行,最后一行为加权平均。配合矩阵图放在结果部分,信息量完全够用。

6.4 多标签任务不是这个画法

这里要特别提醒一句:本文讨论的是单标签多分类任务,每个样本只能属于一个类别。如果你做的是多标签任务,比如一张图同时包含“猫”和“狗”,那混淆矩阵的画法就完全不同了,需要按标签维度单独画或者使用多标签专用的评估指标。矩阵的行列在这里就对应不上,不要混淆。

7. 常见问题与排查思路实录

7.1 画图时报错“size mismatch”

最常见的错误之一。原因通常是类别数对不上,比如模型输出10类,但labels列表只给了8个;或者测试集里某些类别缺失。排查顺序:先打印model最后一层输出维度,再打印训练时保存的类别映射,然后打印len(np.unique(y_true)),逐项核对。

7.2 归一化矩阵里出现NaN

这类问题一般发生在某个行和为0的情况下,也就是某个类别一个测试样本都没有。因为分母为0,除法结果变成NaN。解决方法是先检查类别分布,确保每个类别都有测试样本;如果确实没有,这一行就不该出现在矩阵里,需要从labels里去掉该类别,或者在计算时加一个极小值防止除零。

7.3 图片上中文类别名称显示成方块

matplotlib默认字体不支持中文,直接写中文标签会变成一个个方框。解决方法有两个:一是直接用英文类别名,最省事;二是给matplotlib设置中文字体。代码片段:

plt.rcParams['font.sans-serif'] = ['SimHei'] # 或 ['Microsoft YaHei'] plt.rcParams['axes.unicode_minus'] = False

需要注意,Linux服务器上可能没有SimHei,可以改用DejaVu Sans配合英文标签。如果是长期做中文项目,我建议从始至终统一使用英文类别名加一份映射字典,避免字体问题的困扰。

7.4 GPU显存不足的坑

推理整个测试集时,如果一次性把全部数据丢给模型,显存很容易溢出。正确做法是使用DataLoader分批推理,每次累积预测结果:

all_preds = [] all_labels = [] for images, labels in test_loader: images = images.to(device) labels = labels.to(device) with torch.no_grad(): outputs = model(images) preds = torch.argmax(outputs, dim=1) all_preds.append(preds.cpu()) all_labels.append(labels.cpu()) all_preds = torch.cat(all_preds).numpy() all_labels = torch.cat(all_labels).numpy()

这里我把预测结果先放到CPU,再在最后合并。注意不要每步都torch.cat,这会复制数据多次,效率低;应该先收集列表,最后一次性合并。

7.5 分类报告和矩阵对不上

这个问题的根源往往是y_true和y_pred的来源不一致。比如y_true来自测试集的原始标签,而y_pred来自某个评测阶段的输出文件,两者顺序没有对齐。排查方法是对比前几个样本的路径和预测结果,确认是否一一对应。我自己的经验是:永远在同一个推理循环里同时收集真实标签和预测标签,不要分开跑两次再合并。

8. 我常用的一个完整可交付模板

最后分享一个我自己项目里常用的模板代码,整体打包可直接复制修改。它做了三件事:推理并收集标签、计算原始和归一化矩阵、绘制并导出图片。

import torch import numpy as np import matplotlib.pyplot as plt import seaborn as sns from sklearn.metrics import confusion_matrix, classification_report def evaluate_and_plot(model, test_loader, class_names, device, save_path_prefix='cm'): model.eval() all_preds = [] all_labels = [] with torch.no_grad(): for images, labels in test_loader: images = images.to(device) labels = labels.to(device) outputs = model(images) preds = torch.argmax(outputs, dim=1) all_preds.append(preds.cpu()) all_labels.append(labels.cpu()) y_pred = torch.cat(all_preds).numpy() y_true = torch.cat(all_labels).numpy() cm = confusion_matrix(y_true, y_pred, labels=list(range(len(class_names)))) cm_norm = cm.astype('float') / cm.sum(axis=1, keepdims=True) print('Classification Report:') print(classification_report(y_true, y_pred, target_names=class_names, digits=3)) fig, axes = plt.subplots(1, 2, figsize=(16, 7)) sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=class_names, yticklabels=class_names, ax=axes[0], cbar=False) axes[0].set_title('Confusion Matrix (Counts)') axes[0].set_xlabel('Predicted label') axes[0].set_ylabel('True label') sns.heatmap(cm_norm, annot=True, fmt='.2f', cmap='Blues', xticklabels=class_names, yticklabels=class_names, ax=axes[1], cbar=False) axes[1].set_title('Confusion Matrix (Row-Normalized)') axes[1].set_xlabel('Predicted label') axes[1].set_ylabel('True label') plt.tight_layout() plt.savefig(f'{save_path_prefix}.pdf', format='pdf', bbox_inches='tight') plt.savefig(f'{save_path_prefix}.png', format='png', dpi=300, bbox_inches='tight') plt.show() return cm, cm_norm, y_true, y_pred

这个函数会把归一化和原始计数矩阵并排画在同一张图里,一次导出两种视图。传给它的test_loader不需要shuffle,class_names从dataset.classes里提取。我自己做项目时,会把保存的PDF图直接放进技术报告里,PNG图放到文档中,省去重新绘制的麻烦。

用下来这么久,我觉得混淆矩阵这个工具比想象中更“耐看”。每次训练完新模型,我都会把它当作一个标准动作,和打印准确率一样自然。有时候矩阵看起来结果不理想,反而能逼着我去找数据里的问题,这比模型调参本身更有价值。希望这套流程能帮你省掉一些自己摸索的时间。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询