PyTorch混淆矩阵绘制全指南：从原理到出版级可视化

从项目标题来看这是一篇“技术复现型”的分享。只要动过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个样本被认成类别05个被认成类别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, dim1)就够了不需要先算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的全零矩阵然后遍历每一对真实标签和预测标签把对应位置加1import torch def build_confusion_matrix(y_true, y_pred, num_classes): matrix torch.zeros(num_classes, num_classes, dtypetorch.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_classes3) 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的话矩阵会缺少那一行一列导致维度对不上。我一般在多分类任务中都会显式传入labelslist(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, labelslist(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, annotTrue, fmtd, cmapBlues, xticklabelsclass_names, yticklabelsclass_names, cbarFalse) plt.xlabel(Predicted label) plt.ylabel(True label) plt.title(Confusion Matrix) plt.tight_layout() plt.savefig(confusion_matrix.png, dpi300) plt.show()参数说明annotTrue在每个格子里显示数字。fmtd因为混淆矩阵是整数计数用d表示十进制整数格式。如果错用了fmt.2f或者不写可能报错或者显示成小数。cmapBlues蓝色渐变颜色越深数值越大这是论文里最常见的配色。xticklabels和yticklabels传入类别名称列表让坐标轴显示可读的名称而不是数字索引。cbarFalse我比较喜欢关掉右侧的颜色条因为它只表示颜色深浅对应数值大小但每个格子已经标了数字颜色条的信息量有限还占空间。如果对颜色对比敏感可以保留。4.2 归一化从计数到比例的转换原始计数矩阵有个问题如果每个类别的样本数不同深色的格子分布会失真。例如类别0有1000个测试样本类别1只有100个即使两者的识别率都是80%类别0对应的对角线格子数值是800颜色会非常深类别1只有80看起来很浅给人造成“类别1表现很差”的错觉。解决方法是按行归一化归一化后每个格子的值表示“该行真实类别中有多少比例的样本被预测到了这一列”。归一化矩阵的计算方法把每一行的所有值除以该行的总和。在scikit-learn中可以直接指定normalizetruecm_normalized confusion_matrix(y_true, y_pred, labelslist(range(num_classes)), normalizetrue)在seaborn绘制时我通常这样处理cm confusion_matrix(y_true, y_pred, labelslist(range(num_classes))) cm_norm cm.astype(float) / cm.sum(axis1, keepdimsTrue) sns.heatmap(cm_norm, annotTrue, fmt.2f, cmapBlues, xticklabelsclass_names, yticklabelsclass_names)注意归一化后fmt要改成.2f显示两位小数。在最终展示时我更倾向于同时展示两种形式一是原始计数矩阵用来精确知道样本数量二是按行归一化的比例矩阵用来比较不同类别之间的表现。可参考的展示方案是先打印原始计数矩阵再打印归一化矩阵图里选一种呈现。4.3 坐标轴字体、标题和图像导出细节可视化不只是把图画出来还要考虑最终使用场景。如果图是要放进论文、技术报告或PPT里的要注意以下几点。首先类别名称如果比较长比如“background_with_noise”这种横坐标文字会重叠。解决方法是把坐标轴文字旋转一下plt.xticks(rotation45, haright) plt.yticks(rotation0)旋转角度一般45度就够haright让文字右对齐避免和相邻文字贴在一起。其次字体大小要适配图片尺寸。默认字号在8乘6英寸的图里还行但如果类别很多比如30个类别格子很小数字会挤在一起。这时候要么调大figsize要么调小字体sns.set(font_scale1.2) # 全局放大字体 # 或者单独控制 annot 字体大小 sns.heatmap(..., annot_kws{size: 10})最后是导出格式。我会优先导出一个矢量图加一个位图。矢量图用PDF或SVG用于论文或报告缩放不失真位图用PNGdpi300以上用于网页展示或微信群聊。示例代码plt.savefig(confusion_matrix.pdf, formatpdf, bbox_inchestight) plt.savefig(confusion_matrix.png, formatpng, dpi300, bbox_inchestight)bbox_inchestight会自动裁剪图片边缘空白保证整张图紧凑这一步我几乎每次都会加。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, dim1)如果不加torch.no_grad()推理过程会构建计算图不仅占用大量显存速度也慢画个混淆矩阵把显卡显存跑爆就尴尬了。model.eval()则能保证BatchNorm使用训练阶段统计的均值和方差Dropout关闭这样推理结果才稳定可控。有一次我忘了加model.eval()连续跑两次同一条测试数据得出的预测结果不一样吓了自己一跳最后排查发现是Dropout在作怪。5.3 数据顺序打乱对矩阵没有影响但对可视化有影响混淆矩阵本质上是对所有样本计数的汇总不关心样本顺序。所以DataLoader里是否shuffleTrue不影响矩阵结果。但要注意如果你同时还想输出“哪些具体样本被分错了”以便人工检查错误案例那就必须记录样本路径或ID而不能只记预测结果。顺序一旦打乱你只知道有几个样本被错认却不知道具体是哪几张图。我通常会在推理时把样本的路径和预测结果一起记录下来便于后续人工核对。这步在图像分类里特别重要因为矩阵只能告诉你错误数量不能告诉你错误原因。6. 从矩阵到洞察进阶分析与实战技巧6.1 发现系统性错误相邻类别的强烈混淆混淆矩阵一个特别有用的场景是发现“系统性错认”。比如在植物叶片分类中两种外观极其相似的病害类型常常互相混淆矩阵里会看到两个对称的非对角线位置数值都很高。这说明模型并没有学会区分它们本质上是特征表达不够或者两类样本之间的差异确实微小。发现这种规律后我的处理步骤通常是这样先统计这两个类别的样本量确认是否不平衡然后检查训练数据里这两个类别的图像是否采集条件差异很大比如光照、角度再考虑是否需要对这两类做特征可视化比如用梯度加权类激活映射看模型关注区域是否有问题。6.2 类别不平衡如何让矩阵“说谎”归一化矩阵能解决一部分“比较”问题但解决不了根本问题样本量太少的类别即使归一化后准确率看着还行实际可信度也不高。比如某个类别只有10个测试样本哪怕只有1个被分错归一化后准确率就只有90%而另一个类别有1000个样本50个分错准确率是95%。单看这两个数字好像后者更好但样本量差异导致前者的置信区间大得多。这种情况下我建议在矩阵图下方额外标注每个类别的测试样本数。代码上可以准备一个样本数列表用plt.text在对应的列下方标注或者干脆在类别名称里带上样本量比如“类别A (n100)”。6.3 结合分类报告一起打包输出混淆矩阵是基础但给别人看的时候通常不能只甩一张矩阵图。我在项目汇报时一般会同时输出一份分类报告包含每个类别的精确率、召回率、F1值和样本数。这样既能从矩阵看全局又能从指标看单类表现论证更充分。from sklearn.metrics import classification_report report classification_report(y_true_np, y_pred_np, target_namesclass_names, digits3) 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, dim1) 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_prefixcm): 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, dim1) 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, labelslist(range(len(class_names)))) cm_norm cm.astype(float) / cm.sum(axis1, keepdimsTrue) print(Classification Report:) print(classification_report(y_true, y_pred, target_namesclass_names, digits3)) fig, axes plt.subplots(1, 2, figsize(16, 7)) sns.heatmap(cm, annotTrue, fmtd, cmapBlues, xticklabelsclass_names, yticklabelsclass_names, axaxes[0], cbarFalse) axes[0].set_title(Confusion Matrix (Counts)) axes[0].set_xlabel(Predicted label) axes[0].set_ylabel(True label) sns.heatmap(cm_norm, annotTrue, fmt.2f, cmapBlues, xticklabelsclass_names, yticklabelsclass_names, axaxes[1], cbarFalse) 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, formatpdf, bbox_inchestight) plt.savefig(f{save_path_prefix}.png, formatpng, dpi300, bbox_inchestight) plt.show() return cm, cm_norm, y_true, y_pred这个函数会把归一化和原始计数矩阵并排画在同一张图里一次导出两种视图。传给它的test_loader不需要shuffleclass_names从dataset.classes里提取。我自己做项目时会把保存的PDF图直接放进技术报告里PNG图放到文档中省去重新绘制的麻烦。用下来这么久我觉得混淆矩阵这个工具比想象中更“耐看”。每次训练完新模型我都会把它当作一个标准动作和打印准确率一样自然。有时候矩阵看起来结果不理想反而能逼着我去找数据里的问题这比模型调参本身更有价值。希望这套流程能帮你省掉一些自己摸索的时间。