ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

EfficientNet-PyTorch 图像分类迁移学习实战指南

EfficientNet-PyTorch 图像分类迁移学习实战指南 简介面向深度学习初学者与需要迁移学习实战的开发者这份资源演示了如何基于 EfficientNet-Pytorch 训练自己的分类数据集。资源包共 5 个文件以 4 个 Python 脚本和 1 个 Markdown 说明文档为主压缩后仅 10KB脚本覆盖模型定义、工具函数与训练测试入口结构精简便于快速定位。内容按步骤展开先要求按 train/test 分好类别子目录准备数据集再通过修改指定行选择是否自动下载预训练模型并调整少量参数即可运行训练与测试。读者可借此理解 EfficientNet 迁移学习的完整流程包括数据集组织方式、权重加载机制和脚本修改技巧同时说明文档对训练入口和参数配置作了必要解释能帮助规避常见错误。非常适合初次尝试在该框架上自定义任务的学习者。已有 3780 人学习下载说明该演示具备较好的参考价值。1. EfficientNet-PyTorch 能帮你解决什么问题做图像分类的工程师手里攒了一批自己的图片数据第一反应通常是拿 ResNet 跑个 baseline。但 ResNet 那个年代的设计思路在今天看有点浪费算力——同样的 FLOPs 预算下EfficientNet 在 ImageNet 上的 top-1 准确率能压过 ResNet 一截模型体积还更小。Github 上那个 EfficientNet-Pytorch 项目就是把 Google 的 EfficientNet 用 PyTorch 重新实现了一份接口做得简洁加载预训练权重也不需要去翻 TF Hub 的旧教程直接from_pretrained一行就拿到 ImageNet 权重。它解决的是这么一件事你不需要从零复现 MBConv、SE 模块这些结构只需要把自己的数据整理成分类目录改几行代码就能在这个预训练模型上做迁移学习。这篇笔记就沿着这个路径走先讲清楚 EfficientNet 的缩放逻辑和这个库的加载方式再给一套完整的数据准备和训练脚本最后把我在训练过程中踩过的坑都列出来。适合手里有图片数据、想快速跑出一个可用分类模型的读者也适合那些已经在用 ResNet、想换 EfficientNet 提点准确率的人。2. 选对模型与库EfficientNet 的结构和加载方式2.1 为什么是 EfficientNetcompound scaling 解决了什么问题EfficientNet 的核心不是某一个新的卷积算子而是一套模型缩放方法。此前放大模型的做法很粗暴要么把网络加深要么把每层通道加宽要么把输入分辨率调大一次只动一个维度。Google 的研究者发现这几个维度之间是耦合的——单独加深网络到了某个深度后准确率就饱和了因为网络的感受野和通道容量跟不上单独加宽又容易出现欠拟合。于是他们提出了 compound scaling用一个系数 φ 同时控制深度、宽度和分辨率三个维度。公式是深度d α^φ宽度w β^φ分辨率r γ^φ约束条件是α · β² · γ² ≈ 2同时假设 FLOPS 会增加约2^φ倍这里 α、β、γ 是常数由一个小型网格搜索确定φ 是用户可调的缩放系数。EfficientNet-B0 就是网格搜索出来的 baseline 网络B1 到 B7 分别对应 φ 从 1 到 7。这个缩放方式的精妙之处在于你不需要重新设计网络结构只需要调整 φ就能在算力预算和准确率之间取一个适合自己的点。Net 层结构上EfficientNet 大量使用 MBConv也就是 MobileNetV2 里那个倒残差结构区别是每次卷积后接了 Squeeze-and-Excitation 通道注意力模块激活函数用了 Swish。B0 的整体结构可以简化成一个 stem 卷积接 7 个 stage 的 MBConv 堆叠最后是全局平均池化加全连接分类头。正是这些计算效率高的模块让 B0 只有约 530 万参数却能达到跟 ResNet-50 差不多的 ImageNet 准确率而 ResNet-50 有 2500 多万参数。2.2 安装与加载EfficientNet-Pytorch 库的两种用法这个项目通过 pip 直接安装即可。老项目里有不少人用 Git 克隆方式安装因为早期版本没上 PyPI现在已经可以直接装了pip install efficientnet_pytorch装完之后加载模型有两种方式对应两种不同场景from efficientnet_pytorch import EfficientNet # 方式一加载 ImageNet 预训练权重迁移学习用这个 model EfficientNet.from_pretrained(efficientnet-b0, num_classes10) # 方式二只初始化结构不加载权重从零训练或随机初始化用 model EfficientNet.from_name(efficientnet-b0, num_classes10)from_pretrained会先从网络下载权重文件下载完成后缓存在本地之后加载不需要重复下载。from_name则只搭网络骨架权重随机初始化。这里有个细节如果你是做迁移学习num_classes传入你自己的类别数之后库内部会自动把最后的全连接层替换成对应输出维度的新层同时保留前面所有层的学习率配置。需要注意的是from_pretrained加载的是 ImageNet 权重ImageNet 有 1000 类。你传入num_classes10它会加载前 1000 类的卷积层权重最后的分类层用新初始化的替代。这种设计比手动改model._fc更省心但有一个副作用——如果你后续要冻结某些层做 fine-tune推荐先from_pretrained加载再自己替换model._fc因为旧版本的库直接传num_classes时分类层的初始化分布可能不是你想要的。2.3 结构查看与模型选型参考加载完成之后第一步永远是打印模型结构确认输入输出维度print(model) # 查看模型最终分类层输出维度 print(model._fc.in_features, model._fc.out_features) # 统计参数量确认模型是否符合预期 total_params sum(p.numel() for p in model.parameters()) trainable_params sum(p.numel() for p in model.parameters() if p.requires_grad) print(f总参数量: {total_params / 1e6:.2f}M) print(f可训练参数量: {trainable_params / 1e6:.2f}M)在实际项目中模型选型通常不是越大越好。我的经验是先用 B0 或 B1 跑通整个流程确认数据没问题再根据显存和时间预算升级到 B2 或 B3。下面是几个常用型号的参考值注意不同版本的库或 GPU 环境显存占用会有浮动模型输入分辨率参数量单卡训练 Batch Size 参考8GB 显存适用场景efficientnet-b0224x2245.3M64 左右大部分迁移学习任务首选efficientnet-b1240x2407.8M48 左右数据量中等可接受稍慢训练efficientnet-b2260x2609.2M40 左右需要更高精度显存尚可efficientnet-b3300x30012.2M24 左右数据量大、对准确率要求高efficientnet-b4380x38019.3M12 左右强算力设备或只做推理你的输入分辨率要跟模型匹配。如果训练时用 224x224加载 b3 但输入图保持 224那 b3 的缩放优势反而发挥不出来。模型选好后下一步就是把你自己的数据喂进 DataLoader。3. 把本地图片变成可训练数据集目录规范与增强策略3.1 目录格式与 ImageFolder 读取逻辑这个库的模型入口是torchvision.models的标准风格数据入口自然也用torchvision.datasets.ImageFolder最省事。它要求的目录格式极其简单按类别分文件夹data/ ├── train/ │ ├── cat/ │ │ ├── cat_001.jpg │ │ └── cat_002.jpg │ └── dog/ │ ├── dog_001.jpg │ └── dog_002.jpg └── val/ ├── cat/ └── dog/ImageFolder会自动扫描子目录名作为类别标签并按字母序给每个类别分配一个从 0 开始的索引。这里有一个隐藏的坑如果训练集和验证集的子目录排序不一致或者两边目录名写的不完全一样标签就会错位。常见做法是训练集和验证集共用同一个类别列表或者直接用class_to_idx打印确认from torchvision import datasets train_dataset datasets.ImageFolder(data/train) val_dataset datasets.ImageFolder(data/val) print(train_dataset.class_to_idx) print(val_dataset.class_to_idx) # 两个 dict 必须完全一致 assert train_dataset.classes val_dataset.classes, 类别列表不一致请检查数据目录class_to_idx是{cat: 0, dog: 1}这样的映射确认两边一致后才能进入下一步。3.2 数据划分脚本从零散图片到 train/val很多读者手里的图片不是整理好的可能是一大堆命名混乱的 JPEG外加一个 CSV 标注文件。这时候先写一个划分脚本把图片按类别移动到对应目录同时按比例分出验证集。我一般用shutil和random实现import os import random import shutil from pathlib import Path source_dir Path(raw_images) target_dir Path(data) train_ratio 0.8 random.seed(42) # 假设标注文件格式image_name.jpg,类别名 with open(labels.csv, r) as f: lines f.readlines()[1:] # 跳过表头 samples [] for line in lines: img_name, label line.strip().split(,) samples.append((img_name, label)) # 按类别分组再按比例划分 for img_name, label in samples: src source_dir / img_name if not src.exists(): print(f警告: {src} 不存在已跳过) continue if random.random() train_ratio: dst target_dir / train / label / img_name else: dst target_dir / val / label / img_name dst.parent.mkdir(parentsTrue, exist_okTrue) shutil.copy2(src, dst) print(f共处理 {len(samples)} 个样本按 {int(train_ratio * 100)}% / {int((1 - train_ratio) * 100)}% 划分)random.seed(42)保证每次运行划分结果一致方便复现。这里按样本整体随机划分没有按类别分层。如果某个类别的样本极少比如只有 10 张随机划分后验证集可能只有 2 张模型在验证集上的指标波动会非常大。数据量不平衡时建议改用sklearn.model_selection.train_test_split配合stratify参数做分层划分保证每个类别在训练集和验证集中的占比一致。3.3 Transform 配置迁移学习的数据增强策略数据加载这块训练集和验证集的 transform 要分开写。训练集需要数据增强验证集只需要 resize 和归一化不能在验证集上做随机增强否则验证指标会不稳定且跟训练指标没有可比性。from torchvision import transforms # 通用归一化参数ImageNet 预训练模型的官方推荐值 IMAGENET_MEAN [0.485, 0.456, 0.406] IMAGENET_STD [0.229, 0.224, 0.225] train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.08, 1.0)), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2), transforms.ToTensor(), transforms.Normalize(meanIMAGENET_MEAN, stdIMAGENET_STD) ]) val_transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(meanIMAGENET_MEAN, stdIMAGENET_STD) ])RandomResizedCrop的scale(0.08, 1.0)是 ImageNet 训练的标准配置随机裁剪原图 8% 到 100% 的区域再缩放到 224x224这对尺度变化比较鲁棒。验证集先Resize(256)再CenterCrop(224)比直接Resize(224)多保留一点边缘信息是业内默认操作。这里有一个常见的翻车点如果用的是 efficientnet-b2 或更高版本输入分辨率不是 224而是 260、300、380。transform 里的 crop 尺寸必须跟着模型走否则模型输入尺寸和预训练时的输入分布不一致准确率会掉一截。可以用model._global_params.image_size获取该型号对应的推荐输入尺寸动态替换input_size model._global_params.image_size # b0 是 224b3 是 300拿到正确尺寸后再把 transform 里的 224 都换成input_size这样换模型时不用手动改数字。3.4 DataLoader 参数与数据量检查transform 配好后封装 DataLoaderfrom torch.utils.data import DataLoader train_loader DataLoader( train_dataset, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue, drop_lastTrue ) val_loader DataLoader( val_dataset, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue )shuffleTrue只对训练集开启验证集不需要打乱。drop_lastTrue在训练集上丢弃最后不足一个 batch 的样本避免 BN 统计量在这几个样本上波动。num_workers4在 Windows 上如果报错改成 0 先跑通再逐步调大。Windows 的 DataLoader 多进程还会有一些兼容性问题这是 PyTorch 的老毛病后续避坑章节会细说。数据集准备完毕后做一次数量检查再进训练打印每个类别的样本数确认没有空目录、没有类别数异常。类别数跟模型输出维度对不上是最低级的错误但也是最容易犯的。4. 训练自己的分类器完整训练脚本与关键参数设置4.1 整体训练流程与迁移学习策略从零训练一个 EfficientNet如果想达到 ImageNet 级别的准确率需要数百万张图片和大量 GPU 算力个人开发者基本没这个条件。所以我们的做法必然是迁移学习——加载 ImageNet 预训练权重然后在自己数据集上微调。常见策略分三种第一如果你的数据集跟 ImageNet 内容比较接近比如也是自然物体、场景分类可以只替换分类头冻结主干网络只训练最后一层。这种方式训练极快几分钟就能收敛缺点是准确率上限受限于冻结特征的质量。第二如果数据集有一定差异比如医学影像、卫星图、工业缺陷检测需要解锁主干网络的最后几个 stage 一起训练前面浅层仍然冻结——浅层学到的是边缘、纹理这类通用特征跟任务无关没必要重新训练。第三数据量够大、算力够用就把整个网络全部解冻用一个很小的学习率微调。这是效果最稳的方案但显存和时间消耗也最大。我个人的习惯是第一轮固定训练最后的分类层跑 5 个 epoch 把分类头训到收敛第二轮解冻全部层把学习率调低到之前的十分之一继续训练到验证集指标不再提升。两段式的好处是第一步先让随机的分类头拟合当前特征避免在训练初期梯度回传震荡太大。4.2 完整训练脚本从数据加载到模型保存下面给一套可以直接跑的完整训练脚本。这里用 SGD 加 momentum学习率 0.01配合 warmup 和 cosine 衰减。代码里注释标明了每个关键部分的作用import torch import torch.nn as nn import torch.optim as optim from torch.optim.lr_scheduler import CosineAnnealingLR from efficientnet_pytorch import EfficientNet from torchvision import datasets, transforms # ---------- 配置区 ---------- num_classes 10 batch_size 32 epochs 30 device torch.device(cuda if torch.cuda.is_available() else cpu) lr 0.01 momentum 0.9 weight_decay 1e-4 warmup_epochs 5 save_path best_model.pth # --------------------------- # ---------- 数据加载 ---------- train_loader, val_loader build_dataloaders(batch_size) # 用上一章的 DataLoader 封装 # ---------- 模型初始化 ---------- model EfficientNet.from_pretrained(efficientnet-b0, num_classesnum_classes) model.to(device) # ---------- 冻结主干只训练分类头第一轮 ---------- for param in model.parameters(): param.requires_grad False for param in model._fc.parameters(): param.requires_grad True criterion nn.CrossEntropyLoss(label_smoothing0.1) optimizer optim.SGD(model._fc.parameters(), lrlr, momentummomentum, weight_decayweight_decay) # ---------- warmup cosine 学习率调度 ---------- def lr_lambda(epoch): if epoch warmup_epochs: return (epoch 1) / warmup_epochs else: progress (epoch - warmup_epochs) / (epochs - warmup_epochs) return 0.5 * (1 torch.cos(torch.tensor(progress * 3.14159))) scheduler torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) # ---------- 训练循环 ---------- best_acc 0.0 for epoch in range(epochs): model.train() running_loss 0.0 correct 0 total 0 for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() train_loss running_loss / total train_acc correct / total scheduler.step() # 每个 epoch 结束做验证 val_acc evaluate(model, val_loader, device) print(fEpoch {epoch 1}/{epochs} | fTrain Loss: {train_loss:.4f} | fTrain Acc: {train_acc:.4f} | fVal Acc: {val_acc:.4f}) if val_acc best_acc: best_acc val_acc torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), val_acc: val_acc, }, save_path) print(f已保存最佳模型验证准确率 {val_acc:.4f}) print(f训练完成最佳验证准确率: {best_acc:.4f})label_smoothing0.1是 PyTorch 新版CrossEntropyLoss自带的参数它把 one-hot 硬标签变成 0.9/0.1 分布的软标签可以抑制过拟合对不均衡数据也有一定帮助。CosineAnnealingLR或者LambdaLR二选一都可以我上面用的是自定义的LambdaLR因为标准库的CosineAnnealingLR不支持 warmup 阶段需要额外包一层。训练过程中打印Train Loss下降趋势很关键。正常情况 loss 应该是平滑下降的如果 loss 出现突然跳高再降回去或者完全不下降后面避坑章节会分析原因。4.3 第二阶段解冻全部层低学习率微调第一轮训练结束分类头已经收敛了。这时候解冻全部参数用一个更小的学习率继续训练# 解冻所有层 for param in model.parameters(): param.requires_grad True # 为所有层重新配置优化器学习率降到原来的十分之一 optimizer optim.SGD(model.parameters(), lr0.001, momentummomentum, weight_decayweight_decay) # 重新创建 warmup cosine 调度器epochs 个数改为第二阶段所需轮数 epochs_finetune 20 warmup_epochs_finetune 2 scheduler optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)这里需要注意解冻全部层之后BN 层的running_mean和running_var也需要更新不能冻结 BN。有些教程为了让小 batch 训练稳定会强制 BN 使用全局统计量但那是在 batch size 很小的前提下。正常训练时BN 跟着梯度一起更新即可。第二阶段的学习率推荐在1e-3到1e-4之间。如果第一轮用的 0.01第二轮直接 0.001 是比较稳的。数据量少的话学习率甚至可以低到3e-4。训练轮数不必太长我自己一般设置 15 到 20 轮如果验证集准确率连续 5 个 epoch 不再提升手动停掉就行。4.4 混合精度训练与显存优化如果你的 GPU 支持 AMP现在大多数卡都支持建议一开始就开启混合精度训练。它能直接减少大约一半显存占用同时利用 Tensor Core 加速卷积运算对训练结果是完全无损的from torch.cuda.amp import GradScaler, autocast scaler GradScaler() for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() with autocast(): outputs model(images) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()GradScaler会动态调整 loss 的缩放系数防止梯度在反向传播中下溢为 0。如果不用autocast混合精度显存放在半精度下的计算图就无法构建这一步不能省。AMP 的坑主要有两类一类是 BN 层的统计量在半精度下更新可能不准确不过 PyTorch 的autocast会自动把 BN 保留在 float32这个问题不算大另一类是当 batch size 太小时AMP 的收益不明显因为半精度计算的优势需要大规模并行才能体现。5. 训练避坑指南5 个常见报错与性能问题排查5.1 OSError: unparsable file / 数据集里混入损坏图片这个报错在使用ImageFolder加载图片时非常常见。现象是训练跑到一半DataLoader 突然抛错文件路径指向某一张 JPEG。原因大多是这张图片本身损坏、扩展名是 jpg 但实际是 PNG 格式、或者是网上下载的空文件。解决方法是加载数据前做一次过滤把无法解析的图片踢掉from PIL import Image def verify_image(path): try: img Image.open(path) img.verify() return True except Exception: return False # 遍历数据集目录删除损坏图片 for img_path in Path(data).rglob(*.*): if img_path.suffix.lower() in [.jpg, .jpeg, .png]: if not verify_image(img_path): print(f删除损坏图片: {img_path}) img_path.unlink()Image.verify()只检查文件头是否正确不做完整解码速度很快。但注意verify()之后这个 Image 对象不能再用于加载需要Image.open()重新打开。这个脚本放在训练前跑一次基本能避免 90% 的数据加载报错。如果你的项目里图片来自爬虫或外部系统这一步是必备的。5.2 CUDA out of memory显存不够的三种解法CUDA out of memory是训练中最常见的翻车现场。现象就是在某个 batch 前向传播时报错nvidia-smi 一看显存已经满了。原因无非是 batch size 过大或模型输入分辨率太高但解决方法不止一种。第一种是调小 batch size最简单直接。Batch size 从 32 降到 16显存占用直接减半。第二种是开启上面讲的混合精度通过autocast把激活值存储从 float32 换成 float16显存占用可以下降 40% 左右。第三种是使用梯度累积模拟更大的 batch size同时不增加显存占用accumulation_steps 4 # 相当于 batch size 扩大 4 倍 optimizer.zero_grad() for i, (images, labels) in enumerate(train_loader): with autocast(): outputs model(images) loss criterion(outputs, labels) loss loss / accumulation_steps # 归一化梯度 scaler.scale(loss).backward() if (i 1) % accumulation_steps 0: scaler.step(optimizer) scaler.update() optimizer.zero_grad()注意这里必须把 loss 除以accumulation_steps否则梯度会累加到原始 loss 的accumulation_steps倍等效学习率被放大容易导致训练不稳定。梯度累积的另一个副作用是 BN 层的统计量每个 batch 更新一次而不是累积后更新所以它并不能完全等效于真正的大 batch。对 BN 网络来说如果梯度累积步数超过 8建议同时调小学习率。5.3 加载预训练权重时报错num_classes 不匹配使用from_pretrained(efficientnet-b0, num_classes10)时如果库版本较旧或者你后来手动改了model._fc加载权重可能出现size mismatch的报错。现象是加载过程打印出一长串Missing key(s)和Unexpected key(s)。原因很简单ImageNet 预训练权重的分类层输出是 1000你的分类层输出是 10维度对不上。解决方法是加strictFalse让加载过程跳过不匹配的层model EfficientNet.from_name(efficientnet-b0, num_classes10) state_dict torch.load(efficientnet-b0-355c32eb.pth, map_locationcpu) state_dict.pop(_fc.weight, None) state_dict.pop(_fc.bias, None) model.load_state_dict(state_dict, strictFalse)这个写法是先from_name创建一个随机初始化的模型然后手动加载权重同时把分类层的权重键值移除。load_state_dict里的strictFalse表示只加载能匹配上的参数不匹配的直接跳过。这样处理后模型前面的卷积层拿到 ImageNet 权重分类层保持随机初始化。5.4 训练 Loss 不降或下降极慢如果你的训练 loss 在前几个 epoch 几乎不动或者反而不降反升先检查三件事。第一确认数据归一化是否正确Normalize的mean和std是否用了 ImageNet 的默认值第二检查学习率是否太低特别是迁移学习第二阶段学习率在1e-4以下训练就会非常慢第三查看分类层的初始化分布有些库的_fc层初始化方差过大导致初始 loss 很高且收敛慢。还有一个容易被忽略的原因类别标签映射错误。如果训练集的目录顺序跟验证集不一致模型看到的是错位的标签loss 会一直在一个较高的水平震荡。遇到这种情况打印一个 batch 的标签和图片名称人工检查一下确认标签顺序正确。5.5 验证集准确率比训练集低一大截过拟合还是分布不一致当训练集准确率超过 95%而验证集只有 80%说明模型在训练集上过拟合了。这时优先加强数据增强比如添加RandomRotation、RandomErasing其次是加正则化调大weight_decay到1e-3或1e-2最直接的是增大label_smoothing到 0.2让模型不要太自信。如果训练集和验证集的图片来自不同来源比如训练集是商品图验证集是实拍图那是数据分布不一致增强解决不了需要补充验证集数据或做风格迁移。判断方法很简单在训练集上随机抽一部分做一个小验证集如果小验证集准确率接近训练集而正式验证集落后很多那问题就在数据分布上。6. 从准确率到可用性评测报告、误判分析与微调优化训练完模型验证集准确率看着挺高但在真实场景能不能用还得看更细的指标。分类报告和混淆矩阵是必须做的。用 sklearn 一行就能搞定from sklearn.metrics import classification_report, confusion_matrix import numpy as np def evaluate_full(model, val_loader, class_names): model.eval() all_preds [] all_labels [] with torch.no_grad(): for images, labels in val_loader: images images.to(device) outputs model(images) _, preds torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.numpy()) print(classification_report(all_labels, all_preds, target_namesclass_names)) cm confusion_matrix(all_labels, all_preds) return cmclassification_report会给出每个类别的 precision、recall、F1-score 和样本数。专业向的落地场景里我更关注 recall 最低的类别——它代表最容易被漏掉的样本。比如同一个缺陷类型的良品被误判为次品虽然整体准确率 98%但漏检率可能高达 30%生产线上的风险完全不可接受。拿到混淆矩阵之后下一步是分析误判样本长什么样。我一般会写个脚本把预测错误的图片和它们的真实标签、预测标签、置信度抽出来存到一个文件夹里人工翻看import torch.nn.functional as F def export_misclassified(model, val_loader, class_names, save_dirmisclassified): os.makedirs(save_dir, exist_okTrue) model.eval() with torch.no_grad(): for images, labels in val_loader: images images.to(device) outputs model(images) probs F.softmax(outputs, dim1) for i in range(images.size(0)): prob, pred torch.max(probs[i], 0) if pred.item() ! labels[i].item(): # 保存原图、真实标签、预测标签和置信度 save_name f{labels[i].item()}_{pred.item()}_{prob.item():.4f}_{i}.jpg save_path os.path.join(save_dir, save_name) save_tensor_as_image(images[i], save_path)看误判图的时间永远值得花。你会发现很多误判是有规律的——某个类别被另一个类别系统性误判多半是这两类外形相似需要补充更多难例样本某个类别所有误判的置信度都接近 0.5说明特征不够明显可能需要更高分辨率的输入如果是某些特定角度、光线下的图全错了那就是数据增强里少了这些扰动可以针对性补偿。针对难分类别除了加数据还有两个常用技巧。一个是 focal loss替代CrossEntropyLoss让模型更关注难以分类的样本另一个是在微调阶段对难例样本做 oversampling让模型在一个 epoch 里看到更多困难样本。这两个技巧不能根治分布问题但通常能再提升 1 到 3 个点的召回率。最后说一个我的习惯模型保存时除了保存模型权重一定把数据的class_to_idx、transform 参数、输入分辨率一起存成 JSON。半年之后回来重训或者部署不用对着旧代码猜参数。这个习惯救过我很多次希望也能帮到你。本文还有配套的精品资源点击获取
返回列表