3行代码搞定PyTorch混淆矩阵：从评估指标到业务决策可视化-CSDN博客

3行代码搞定PyTorch混淆矩阵：从评估指标到业务决策可视化

【免费下载链接】pytorch-book PyTorch tutorials and fun projects including neural talk, neural style, poem writing, anime generation (《深度学习框架PyTorch：入门与实战》) 项目地址: https://gitcode.com/gh_mirrors/py/pytorch-book

你是否还在为模型评估报告中的 accuracy 数值纠结？是否遇到过"准确率99%却分错关键样本"的尴尬？混淆矩阵（Confusion Matrix）正是解决这类问题的关键工具。本文将带你用 PyTorch 实现混淆矩阵的完整流程，从5行核心代码到电商商品分类的实战案例，让模型评估从此告别"数字游戏"。

读完本文你将获得：

3分钟上手的混淆矩阵实现模板
4种可视化方案（含TensorBoard/Visdom对比）
电商场景下的误分类分析实战
模型优化的3个决策依据

一、为什么需要混淆矩阵？

在二分类问题中，混淆矩阵能清晰展示真正例（True Positive）、假正例（False Positive）、真负例（True Negative）和假负例（False Negative）的分布。以猫狗分类任务为例：

传统准确率（Accuracy）计算方式为：

准确率 = (正确分类样本数) / (总样本数)

但当样本不平衡时（如99%为猫），模型只需全部预测为猫即可获得99%准确率。而混淆矩阵能揭示：

有多少猫被错误分类为狗（FN）
有多少狗被错误分类为猫（FP）
模型在哪些类别上表现最差

二、PyTorch混淆矩阵实现（5行核心代码）

2.1 基础实现

在 Chapter5/chapter5.ipynb 中，我们可以找到混淆矩阵的基础实现：

def confusion_matrix(preds, labels, num_classes):
    preds = torch.argmax(preds, dim=1)  # 将概率转换为类别索引
    matrix = torch.zeros(num_classes, num_classes, dtype=torch.int32)
    for p, l in zip(preds, labels):
        matrix[p, l] += 1
    return matrix

2.2 批量计算优化

在实际训练中，建议使用向量化操作优化性能：

def batch_confusion_matrix(preds, labels, num_classes):
    preds = torch.argmax(preds, dim=1)
    # 使用torch.bincount计算混淆矩阵
    indices = num_classes * labels + preds
    return torch.bincount(indices, minlength=num_classes**2).reshape(num_classes, num_classes)

三、可视化方案对比

3.1 TensorBoard可视化

PyTorch内置的TensorBoard支持混淆矩阵可视化，在 Chapter5/Chapter5.md 中有详细配置：

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter(log_dir='runs/confusion_matrix_demo')
cm = confusion_matrix(preds, labels, num_classes=2)
writer.add_figure('Confusion Matrix', plot_confusion_matrix(cm), global_step=epoch)

3.2 Visdom交互式可视化

Visdom提供更灵活的交互体验，配置方法见 Chapter5/Chapter5.md：

import visdom
vis = visdom.Visdom(env='confusion_matrix')
vis.heatmap(
    X=cm.numpy(),
    win='confusion_matrix',
    opts=dict(
        title='Confusion Matrix',
        rownames=['Predicted Cat', 'Predicted Dog'],
        columnnames=['Actual Cat', 'Actual Dog']
    )
)

四、电商商品分类实战

4.1 数据准备

使用 Chapter5/data/dogcat_2/ 中的猫狗分类数据：

from torchvision.datasets import ImageFolder
from torchvision import transforms

transform = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5])
])
dataset = ImageFolder('Chapter5/data/dogcat_2/', transform=transform)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

4.2 模型训练与评估

加载预训练模型并计算混淆矩阵：

from torchvision import models

model = models.resnet18(pretrained=True)
model.fc = torch.nn.Linear(512, 2)  # 二分类输出

# 训练循环中计算混淆矩阵
total_cm = torch.zeros(2, 2)
for imgs, labels in dataloader:
    preds = model(imgs)
    cm = batch_confusion_matrix(preds, labels, num_classes=2)
    total_cm += cm

print("最终混淆矩阵:\n", total_cm.numpy())

4.3 结果分析

假设得到混淆矩阵：

[[350,  20],  # 实际为猫：350正确分类，20错误分类为狗
 [ 40, 290]]  # 实际为狗：290正确分类，40错误分类为猫

分析发现：

猫被误分类为狗的比例：20/(350+20)=5.4%
狗被误分类为猫的比例：40/(290+40)=12.1%
模型对狗的识别准确率更低，可能原因：
- 狗的样本数量较少
- 狗的姿态变化更多样
- 训练集中狗的图片质量较差

五、模型优化决策依据

基于混淆矩阵分析，可采取以下优化策略：

数据层面：增加狗类样本，特别是误分类样本的相似样本
算法层面：使用 Chapter5/Chapter5.md#52-预训练模型中的迁移学习策略
损失函数：对狗类样本使用更高的权重

# 加权损失函数示例
weights = torch.tensor([1.0, 1.5])  # 狗类样本权重提高50%
criterion = torch.nn.CrossEntropyLoss(weight=weights)

六、总结与扩展

混淆矩阵是模型评估的"显微镜"，通过它我们能：

发现准确率掩盖的问题
定位模型的薄弱环节
制定针对性的优化策略

在多分类场景下，可扩展为N×N矩阵，并结合 Chapter5/imgs/Tensorboard_embedding.png 中的嵌入可视化，深入分析类别间的混淆关系。

完整代码示例可参考：

希望本文能帮助你构建更健壮的模型评估流程，让每一个准确率数字都经得起业务检验。

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考