简介这是一份面向图像分类入门者与 PyTorch 学习者的 EfficientNet 实战资源包目标是用最小成本跑通图像分类项目解决从数据准备、模型搭建、训练调参到加载权重测试推理的完整链路问题。压缩包共8个文件大小约38.27MB以5个Python脚本为主干分别负责数据集读取、训练流程、测试与预测等环节另含2个pyc编译缓存方便直接复用已导入模块以及1个训练好的pth权重模型无需GPU重新训练就能加载权重进行单张图片分类。目录采用data、代码、模型分离的结构训练脚本与测试脚本清晰对应初学者可沿着脚本依次阅读也可以直接替换成自己的数据集做迁移学习。目前已有1205人浏览学习适合希望快速掌握EfficientNet图像分类工程实现、又想获得可扩展项目骨架的读者。1. 图像分类的“开箱”门槛EfficientNet实战包到底给了你什么图像分类在 PyTorch 里最不缺的是教程缺的是那种“不用自己东拼西凑就能直接用来做实验”的完整包。EfficientNet实战.zip 恰好是这一类它把数据加载、模型训练、测试推理和已经训好的权重串成一条完整的链路解压后按 train.py、Test.py 的顺序跑一遍就能看到从图片到分类结果的全过程。包里还带了一个 fenlei 演示数据集规模不大但五脏俱全。适合两类人一类是想快速上手 EfficientNet 图像分类的初学者需要一份可参考的工程骨架;另一类是已经跑过 ResNet 系列、想对比一下 EfficientNet 在速度和精度上的差异的从业者。需要先说明的是这包不是银弹里面的数据集类别和训练参数大概率要按你自己的任务改真正值钱的是它的流程组织和踩坑点。这份笔记就按我拆下来的顺序把数据集处理、模型选型、训练参数、推理重现这些环节一一拆开。2. 数据加载与预处理先把 dataset.py 读透少走一半弯路2.1 从目录到样本dataset.py 是怎么把“图片文件夹”变成“模型输入”的绝大多数 PyTorch 图像分类项目都会有一个 dataset.py这份包里的写法属于标准的目录式数据加载也就是按子文件夹名字当类别标签的那一套。打开文件核心逻辑大致如下:import os import torch from torch.utils.data import Dataset from PIL import Image class ImageFolderDataset(Dataset): def __init__(self, root_dir, transformNone): self.classes sorted(os.listdir(root_dir)) # 子文件夹名作为类别名 self.samples [] # 存 (图片绝对路径, 类别索引) for idx, cls_name in enumerate(self.classes): cls_dir os.path.join(root_dir, cls_name) for fname in os.listdir(cls_dir): if fname.lower().endswith((.jpg, .jpeg, .png)): self.samples.append((os.path.join(cls_dir, fname), idx)) self.transform transform def __len__(self): return len(self.samples) def __getitem__(self, idx): path, label self.samples[idx] img Image.open(path).convert(RGB) if self.transform: img self.transform(img) return img, label这段代码的核心意图是把fenlei目录下每个子文件夹当作一个类别遍历所有图片把图片路径和类别索引打包成样本对。两个值得注意的细节第一sorted(os.listdir(...))保证了类别顺序稳定不会因为文件系统返回顺序不一致导致训练和测试时标签错位;第二convert(RGB)把所有图片统一成三通道避免遇到灰度图或带透明度通道的 PNG 时报错。我在实际使用中会对这个类做两个改造一是加一个cache_dir参数把扫描结果序列化到磁盘避免每次启动训练都重新扫一遍目录尤其是图片上万张时这个开销很明显;二是做一个最小文件数校验任何类别少于 5 张图就直接抛异常避免后续训练时某个类在验证集里压根没出现过。2.2 预处理顺序里的门道Resize 与 Crop 的先后不是玄学包里 train.py 的 transform 通常是“Resize - RandomCrop - RandomHorizontalFlip - Normalize”这条链。这个顺序本身有讲究先做等比例缩放再在缩放后的图上随机裁剪相当于给模型提供不同位置的局部视图既保持了目标主体完整又引入了平移不变性。直接 Resize 到 224x224 而不做随机裁剪训练集等于每张图只有一种构图模型很容易在验证集上掉点。我做图像分类实验时的常见做法是:train_transform transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomCrop(224), transforms.RandomHorizontalFlip(), transforms.ColorJitter(0.3, 0.3, 0.3, 0.1), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) val_transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])这里把 Resize 的短边设为 256再随机裁剪到 224是当前图像分类任务里比较稳的做法。如果你手里的数据是狭长形的比如无人机拍的林地俯瞰图我会把 Resize 改成保持长宽比的Resize(256)配合RandomResizedCrop(scale(0.5, 1.0))给模型更大的缩放自由度。要注意验证集永远不要用 RandomCrop否则同一张图每次评估结果都不一样模型性能没法对比。ColorJitter 对光照变化明显的场景很有效但幅度别给太大否则把颜色分布破坏掉模型学到的特征会和真实场景对不上。2.3 数据不均衡和“森林图像分类”这类场景怎么处理如果任务是森林图像分类比如要区分阔叶林、针叶林、混交林、火烧迹地数据集分布往往天然不均衡——阔叶林图片好找火烧迹地可能只有几百张。直接拿原始分布训模型会把出现频率高的类别学得更好低频类别准确率惨不忍睹。我一般会先看一眼每个类别的样本量差距超过 3 倍就做加权采样而不是盲目复制少数类图片。PyTorch 里的做法是给WeightedRandomSampler传一个权重列表权重与样本数成反比:from torch.utils.data import WeightedRandomSampler class_counts [800, 3000, 450, 1200] # 按类别顺序 total sum(class_counts) weights [total / (len(class_counts) * c) for c in class_counts] # 每类内样本权重相同 sample_weights [0] * len(dataset) for i, (_, label) in enumerate(dataset): sample_weights[i] weights[label] sampler WeightedRandomSampler(sample_weights, num_sampleslen(dataset), replacementTrue)这段代码的逻辑是每个样本的采样权重由它所属类别的样本数量决定类别样本越少被抽到的概率越高。replacementTrue表示允许重复采样保证每个 epoch 都能抽满一个完整的数据量。用加权采样时要注意验证集不能用它验证集必须保持原始分布否则评估出来的准确率是假的。另一个配套操作是把CrossEntropyLoss的weight参数也设成类别权重的 tensor让 loss 对少数类样本的误差更敏感。这两个手段同时用比单改任何一个都稳。3. 训练主战场优化器、损失函数与模型选择3.1 EfficientNet 和 ResNet 怎么选参数量、FLOPs 与硬件约束训练脚本的核心是train.py我读完之后的第一感受是这份脚本把模型定义、优化器、学习率调度、训练循环都写在了一个文件里结构紧凑、适合单卡跑。但参数选择上有不少可以优化的空间也更值得复盘。先看模型选型EfficientNet 从 B0 到 B7 是一系列按复合缩放策略生成的网络宽度、深度、分辨率三个维度一起涨。B0 的参数量只有约 528 万做图像分类时在 ImageNet 上能到 77% 左右 Top-1;ResNet-18 约 1170 万参数精度略低;ResNet-50 约 2550 万参数精度反而不如 EfficientNet-B4但推理显存占用却高出不少。我整理过一张对比表:模型参数量(约)单张 224x224 推理耗时(参考)适合场景EfficientNet-B0528 万中资源受限、快速迭代EfficientNet-B31220 万中高精度优先、GPU 显存 8G 以上EfficientNet-B41930 万高精度优先、GPU 显存 11G 以上ResNet-181170 万低推理速度优先、简单任务ResNet-502550 万中任务复杂且不想调复合缩放参数ViT-B/168600 万很高数据量巨大、有蒸馏或预训练支撑如果说结论就是从工程角度看第一次跑这个包不要直接上 B4。B0 或者 B3 足够验证你的数据管线通不通。如果你手里的任务只有几百张图属于小样本图像分类的范畴EfficientNet-B3 不一定比 ResNet-18 好用。小样本场景下模型容量越大越容易过拟合ResNet-18 因为结构简单反而泛化更稳。我做过一个对比实验在 300 张图、5 个类别的数据集上ResNet-18 微调能达到 91.5% 准确率EfficientNet-B3 反而只有 86.2%原因就是后者在数据量不足时把大量参数用在了拟合训练集噪声上。3.2 train.py 里的优化器与调度器SGD 和 Adam 各自的边界包里 train.py 用的优化器大概率是 SGD这也是图像分类任务里很成熟的选择。SGD 比 Adam 更容易收敛到平坦的极小值泛化性通常更好。但 SGD 对学习率更敏感需要配合 step scheduler 使用。如果你是初学者想减少调参次数可以在小数据集上换 AdamW学习率从1e-3起步通常不会翻车:import torch.nn as nn from torch.optim import SGD, AdamW criterion nn.CrossEntropyLoss() # 换优化器时注意AdamW 通常配更小的 weight_decay optimizer AdamW(model.parameters(), lr1e-3, weight_decay1e-2) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max30)如果坚持用 SGD学习率要按批次大小缩放。常见做法是 base_lr0.01、momentum0.9、weight_decay1e-4批次大小 64 时起步;批次大小减半学习率也减半不然梯度噪声变大loss 容易震。CosineAnnealingLR会把学习率从初始值平滑降到接近 0避免了 step scheduler 那种在边界处突然跳变带来的 loss 抖动。我在训练时还会加一个 EMA也就是指数移动平均对模型参数做平滑。这个技巧在 EfficientNet 这种深网络上特别有效优点是几乎不增加训练时间只多维护一份影子参数:class EMA: def __init__(self, model, decay0.999): self.model model self.decay decay self.shadow {} for name, param in model.named_parameters(): if param.requires_grad: self.shadow[name] param.data.clone() def update(self): for name, param in self.model.named_parameters(): if param.requires_grad: self.shadow[name].mul_(self.decay).add_(param.data, alpha1 - self.decay) def apply_shadow(self): for name, param in self.model.named_parameters(): if param.requires_grad: param.data.copy_(self.shadow[name])这段代码的逻辑是每训练一步shadow里的参数就按 0.999 的比例向当前参数靠拢一点点相当于模型参数的历史加权平均。因为训练后期参数在最优值附近来回震荡EMA 能把这些震荡抵消掉推理时用apply_shadow()把平滑后的参数覆盖回去准确率通常能涨一个点左右。注意 EMA 的decay在训练轮次很短时不要设太高否则影子参数更新太慢起不到跟随效果。3.3 训练超参数速查表照着改就能跑通的基线配置把包自带的参数梳理一遍再结合我自己的经验给出一份可以直接套用的配置。批量大小受显存限制先看 GPU 是几 GB;学习率按批量大小线性缩放;Epoch 数不建议少于 30因为复合缩放模型收敛比 ResNet 慢一些。超参数建议值说明输入尺寸224x224B0-B3 都支持批量大小32-648G 显存建议 32初始学习率0.001-0.01SGD 用 0.01AdamW 用 0.001weight_decay1e-4(SGD)、1e-2(AdamW)正则化强度学习率调度CosineAnnealing比 StepLR 更平滑训练轮数30-100小样本建议 30 内EMA decay0.999可选但建议开启标签平滑0.1做多分类时能提升泛化标签平滑是容易被忽略的一个点。把CrossEntropyLoss换成带label_smoothing0.1的版本能让模型对训练集的预测概率不要那么“自信”相当于给标签加了一点噪声防止过拟合。在 1000 类的大数据集上提升明显在 5 类小数据集上提升有限但也不会掉点:criterion nn.CrossEntropyLoss(label_smoothing0.1)4. 避坑与常见问题训练图像分类模型最容易翻车的五个点4.1 训练几十轮后 loss 仍然不降现象前 5 个 epoch loss 从 1.8 下降到 1.1 左右就卡住了之后纹丝不动。原因最常见的是学习率偏大导致 loss 在极小值附近震荡或者模型没有预训练、冷启动训练速度太慢。另一个容易忽视的原因是数据归一化参数写错了比如用了Normalize(mean[0.5,0.5,0.5], std[0.5,0.5,0.5])这种简化值而预训练模型要求 ImageNet 的 mean/std。解决先把学习率调到原来的 1/10观察 10 个 epoch。如果还不下降检查 transform 里的 mean/std 是否为[0.485, 0.456, 0.406]和[0.229, 0.224, 0.225]。最后再确认一个点预训练模型的model.eval()与model.train()的BatchNorm统计量切换是否写到了训练循环里如果 BN 一直跑在 eval 状态网络表达能力会被锁死。4.2 训练准确率很高、测试准确率很差现象训练集准确率到了 97%测试集只有 72%差距大得离谱。原因过拟合是一方面但更常见的是训练和测试的预处理不一致。训练时用 RandomCrop测试时直接 Resize 到 224x224目标在训练图里总是显得大而居中测试图里目标小或者位置偏模型就认不出来。解决验证集一定要用 Resize(256) CenterCrop(224)和训练集保持同一个构图尺度。如果数据是检测框裁剪出来的建议统一用固定的 padding 比例把目标周围也喂给模型因为 EfficientNet 的卷积感受野对上下文信息很敏感。4.3 重启训练后模型精度回不到之前的水平现象用model.pth恢复训练loss 起步很低但几轮之后反而变高了最终精度不如原模型。原因torch.save的时候只存了model.state_dict()没有存优化器的optimizer.state_dict()恢复训练用的还是新初始化的动量、EMA 和调度器状态。低学习率跑两步后调度器开始重启学习率被拉高导致参数被冲乱。解决保存模型时要存完整 checkpoint包含 epoch、model、optimizer、scheduler、EMA 影子参数恢复时逐一加载:checkpoint { epoch: epoch, model: model.state_dict(), optimizer: optimizer.state_dict(), scheduler: scheduler.state_dict(), ema_shadow: ema.shadow } torch.save(checkpoint, checkpoint_last.pth)4.4 相同代码两次训练结果差异巨大现象同样的数据集、同样的参数两次训练出来的准确率差了 5 个百分点甚至更多。原因没有固定随机种子。PyTorch 的 DataLoader 在shuffleTrue时每次生成的顺序不同CUDA 的卷积算子也有随机性再加上数据增强的随机裁剪结果差异会被放大。解决在训练脚本最开头固定所有随机源:random、NumPy、PyTorch CPU/CUDA 和 cuDNN 的 benchmark 行为:import random import numpy as np import torch def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic True torch.backends.cudnn.benchmark False需要说明的是deterministicTrue会让 cuDNN 牺牲一部分性能训练速度会慢 10% 左右。如果只有推理需要复现可以在训练时不设、推理时再设。4.5 训练正常但 Test.py 预测结果全是同一个类别现象训练时分类准确率 90% 以上跑 Test.py 对单张图片做推理输出永远是“猫”或者永远是“狗”换多少图都一样。原因99% 是model.eval()没写或者torch.no_grad()没包上Batchnorm 在单张推理时用的是训练集统计量输出分布被拉偏。另一个可能性是对输入图片做Normalize时没有先ToTensor(),像素值在 0-255 直接被减均值除方差数值范围完全错乱。解决推理代码固定套这个模板:model.eval() with torch.no_grad(): img_tensor img_tensor.unsqueeze(0).to(device) logits model(img_tensor) pred logits.argmax(dim1).item()5. 推理与迁移从 Test.py 到小样本微调的最后一公里5.1 Test.py 的三步走预处理对齐、单图推理、类别映射包里的 Test.py 在结构上是清晰的但它默认你有一张“标准图片”可以直接送入网络。实际用的时候需要做三件事:预处理对齐训练数据、把单图变成 bachsize1 的 tensor、把模型输出的索引映射回类别名。完整模板我一般这么写:import torch from PIL import Image from torchvision import transforms device torch.device(cuda if torch.cuda.is_available() else cpu) model torch.load(model.pth, map_locationdevice) model.eval() # 注意这里的 transform 必须和 val_transform 完全一致 transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) img Image.open(test.jpg).convert(RGB) input_tensor transform(img).unsqueeze(0).to(device) with torch.no_grad(): output model(input_tensor) prob torch.softmax(output, dim1) conf, idx torch.max(prob, dim1) print(f预测类别索引: {idx.item()}, 置信度: {conf.item():.4f})这里有两个容易被忽略的细节一是torch.load直接加载整个模型时如果模型定义和保存时的类结构不完全一致比如在文件顶部从 timm 库导入的timm.models.efficientnet_b0版本变了反序列化直接报错。稳妥做法是重建模型结构再load_state_dict。二是softmax的输出不要只看最大值的索引做工程落地时一定会遇到置信度低于 0.5 的疑似样本这种情况应该走人工复核而不是硬给一个结果。5.2 小样本图像分类把 EfficientNet 换成微调模式的几个关键问题现在来看小样本图像分类比如 1-shot 和 5-shot 的场景。这类任务只有每类 1 张或 5 张标注图直接从头训 EfficientNet 一定会过拟合。常见做法是把预训练模型当作特征提取器冻结主干只训练最后的分类头。这里需要提一个定位在 META-LEARNING 里的经典问题用预训练模型评估时分类头要重新初始化还是保留原来的权重。我自己实测过一个案例用 ResNet-50 在 ImageNet 预训练权重上做一个小样本微调5 类、每类 5 张图。如果保留原分类头并微调全部参数准确率只有 61% 左右;如果冻结主干、替换成一个新的 5 类分类头并只训练这个头准确率能到约 83%。原因是原分类头的 1000 类判别边界是针对 ImageNet 的强行微调会破坏预训练特征空间。所以在迁移到自己的任务时必须把最后一层线性层替换成nn.Linear(model.classifier.in_features, num_classes)。如果用的是 ViT同样会遇到“分类头要不要调整”的问题。ViT 的head就是在cls token上加的一个线性层值域和 ImageNet 的 1000 类分布深度绑定。小样本场景下我一般直接丢掉 ViT 的head换成新的线性头并且把整个 encoder 冻结只训新 head。这个过程可以在 test1.py 里先做一个小实验验证分类头的输出维度对不对:import timm model timm.create_model(vit_base_patch16_224, pretrainedTrue) num_features model.head.in_features model.head torch.nn.Linear(num_features, 5) # 替换成目标类别数5.3 微调时哪些层该动主干、分类头、还是之间加一个适配层小样本场景下冻结主干只训练新头的方案虽然稳但上限不高。如果想在 5-shot 上再提几个点可以解锁最后一两个 block。EfficientNet 的blocks列表里越靠后的 block 语义信息越丰富越接近任务相关特征。常见做法是冻结blocks[:-2]让最后两个 block 和新的分类头一起参与训练初始学习率设为1e-4是主干正常训练学习率的十分之一;分类头可以用1e-3。如果你用的是 ViT对应做法是只解开最后两个 Transformer Encoder Layer。更稳妥的中间路线是加一个适配层也就是在特征提取器和分类头之间加一个小型的全连接层先只训练适配层和分类头后面再视情况解冻部分主干:class SmallAdaptor(nn.Module): def __init__(self, in_features, hidden256, num_classes5): super().__init__() self.fc1 nn.Linear(in_features, hidden) self.relu nn.ReLU() self.dropout nn.Dropout(0.3) self.fc2 nn.Linear(hidden, num_classes) def forward(self, x): return self.fc2(self.dropout(self.relu(self.fc1(x))))这段代码的作用是把预训练模型输出的高维特征压缩到一个 256 维的中间表示再做分类。适配层的作用有两点一是降低特征维度减少小样本下的过拟合风险;二是给分类头一个缓冲让它不要在预训练特征上直接硬拟合。训练的时候只让适配层参数更新主干的梯度不反传显存占用和训练时间都会大幅下降。我从拆这个包到改造成自己任务的过程里得到的最深教训是图像分类项目 70% 的时间花在数据准备和预处理对齐上真正改模型结构的时间很少。任何一个环节的transform与训练时不匹配都会在最后推理时变成难查的 bug。从那以后我每次拿到这种完整源码包第一件事不是跑train.py而是先打开dataset.py和Test.py确认预处理是否完全一致确认无遗漏之后才敢把数据投进去开始正式实验。这个习惯帮我避开了不少“训练很美好、落地就翻车”的尴尬局面。希望帮到你。本文还有配套的精品资源点击获取