简介：面向深度学习初学者的 EfficientNet 自定义数据集训练演示包，帮助读者在 EfficientNet-PyTorch 框架下快速搭建图像分类训练与测试流程。资源内含完整 Python 脚本及 Markdown 说明文档，共 5 个文件（4 个 py、1 个 md），压缩包仅 10KB，轻量实用。演示覆盖了从数据集目录规范（train/test 按类别分文件夹）到脚本参数调整、预训练模型自动下载开关等关键环节，读者可按说明修改对应行配置即可启动训练。示例中明确展示了数据集的组织方式：train 和 test 目录下按类别编号存放图片，便于迁移到自己的数据；同时提供了预训练权重下载提示（可放入 eff_weights 文件夹），帮助读者快速获得 EfficientNet 在 ImageNet 上的初始化参数。目前已有 3780 人学习，适合希望掌握 EfficientNet 迁移学习操作细节的入门者参考，既能理解训练流程，也能直接复用代码改造自己的任务。
1. EfficientNet-Pytorch 这条路值不值得走：换骨干之前先想清楚三件事
我最早接触 EfficientNet-Pytorch 这个仓库时，和多数人一样是冲着“ImageNet 上比 ResNet 省参又提点”去的。真正动手才发现，仓库里的 demo 都在 ImageNet 上跑得风生水起，真要把自己的分类数据集换上去，还得自己补一堆代码：目录怎么组织、分辨率配多少、预训练权重怎么改分类头、训练超参数怎么调。这篇文章就是把这条路线完整走一遍，从环境搭建到训练脚本再到踩坑排查，让 EfficientNet-Pytorch 不只是别人仓库里的代码，而是你手里能用的工具。
这篇笔记适合三种人：想把骨干网络从 ResNet 换到 EfficientNet 的算法工程师、复现论文但卡在工程细节上的研究生、以及第一次用 PyTorch 做迁移学习的入门者。先给一个反直觉的结论：EfficientNet 在小数据集上的优势，大部分来自预训练权重和正确的输入分辨率，而网络结构本身带来的增益反而没那么大。这个认知会贯穿后面的所有参数选择。
2. 从复合缩放到 PyTorch 环境搭建：为什么 EfficientNet 比 ResNet 难伺候
2.1 EfficientNet 的缩放逻辑：B0 到 B7 不只是变深变宽
EfficientNet 的核心是 compound scaling，也就是复合缩放。ResNet 加层、加通道、加大输入分辨率，这些操作是分开考虑的，而 EfficientNet 用一个系数 phi 同时控制网络深度、宽度和输入分辨率。这个设计的动机很朴素：网络变深了需要更多通道来传递信息，分辨率变大了需要更多层来扩大感受野，单独缩放某一维很快就碰到收益天花板，三个一起动反而能压着计算量预算走。
PyTorch 里面 EfficientNet 的实现基本沿用了论文结构：stem 卷积 + 一串 MBConvBlock + 最后的分类头。MBConv 是 MobileNet 的倒残差结构，中间还夹着 SE（Squeeze-and-Excitation）注意力，激活函数用的是 Swish。这些模块在 efficientnet_pytorch 这个库里都是现成的，这也是为什么大多数人直接用这个库而不自己复现。
B0 到 B7 的区别不仅是参数量，分辨率也跟着变。这是 EfficientNet 最容易踩坑的地方：训练 B0 用 224，训 B5 就得用 456，不是随便 resize 到 224 就能跑的。
| 模型 | 输入分辨率 | 参数量（大约） | 适用场景 |
|---|---|---|---|
| B0 | 224 | 5.3M | 显存有限、快速验证 |
| B1 | 240 | 7.8M | 小数据集微调 |
| B2 | 260 | 9.2M | 中等数据集 |
| B3 | 300 | 12M | 精度优先但显存中等 |
| B4 | 380 | 19M | 显存充足 |
| B5 | 456 | 30M | 大显存、精度优先 |
| B7 | 600 | 66M | 不建议自己训，推理也贵 |
我的建议是：第一次跑通流程用 B0，把脚本调顺了再考虑换更大的模型。B0 在 224 分辨率下 batch size 32 也就占用 8G 左右显存，调试成本最低。
2.2 Anaconda 配置 PyTorch 环境：CPU 与 GPU 两套最小命令
环境搭建是新手卡住的第一道墙。我一般用 Anaconda 创建独立环境，避免把 base 环境的包搞乱。先给两套安装命令，一套 CPU 一套 GPU：
# 创建环境 conda create -n efficientnet python=3.8 -y conda activate efficientnet # CPU 版本（没独显或只是先跑通流程） pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu # GPU 版本（先确认驱动支持，再选对应 CUDA 版本） pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 # 安装 EfficientNet 库 pip install efficientnet-pytorch # 顺手装个日志工具 pip install tensorboard先说 conda 那行：python=3.8 是我用得最多的版本，PyTorch 和 efficientnet-pytorch 对它的兼容性最好，没必要追新。pip 安装 torch 时注意--index-url指定了下载源，CPU 和 GPU 的包是分开的，装混了会报 "Torch not compiled with CUDA enabled"。GPU 版本选 cu118 还是 cu121，取决于你的显卡驱动版本，我一般直接在终端跑nvidia-smi看右上角的 CUDA Version，只要驱动支持的版本大于等于你装的就行。
验证环境是否装好的命令：
python -c "import torch; print(torch.__version__, torch.cuda.is_available())"输出类似2.1.0+cu118 True才是 GPU 可用状态，这里返回 False 的话，后面训练全在 CPU 上跑，一个 epoch 能慢十几倍。把这行验证命令贴在训练脚本开头也是个好习惯，免得排队半天才发现设备不对。
3. 让自定义数据集适配 EfficientNet：目录结构、Dataset 改写与增强参数
3.1 ImageFolder 目录组织：train / val 分开是底线
EfficientNet-Pytorch 仓库没有自带数据加载逻辑，训练自己的数据集时最省事的方案是 torchvision 的 ImageFolder。它要求目录按类分好，一层是类别名，一层是图片文件。常见的目录组织是这样的：
data/ ├── train/ │ ├── cat/ │ │ ├── cat_001.jpg │ │ └── cat_002.jpg │ └── dog/ │ ├── dog_001.jpg │ └── dog_002.jpg └── val/ ├── cat/ │ └── cat_003.jpg └── dog/ └── dog_003.jpgtrain 和 val 分开是底线，千万别把所有图片都塞进 train 然后用同一份数据又训又验，那叫过拟合表演赛。加载的代碼很简单：
from torchvision import datasets, transforms train_set = datasets.ImageFolder('data/train', transform=train_transform) val_set = datasets.ImageFolder('data/val', transform=val_transform) print(train_set.classes) # 类别名列表 print(train_set.class_to_idx) # 类别到索引的映射这里有个容易被忽略的细节：ImageFolder 的 label 是按类别名字母序排的，不是按你文件夹里的顺序。比如你文件夹创建顺序是 dog 在前 cat 在后，但class_to_idx里 cat 仍然是 0，dog 是 1。训练之前一定要把那行print(train_set.class_to_idx)的输出存下来，后面推理的时候要拿索引反查类别名，没有这个映射你预测出个 0 都不知道是猫还是狗。
如果图片不是按目录组织的，比如存在 CSV 或数据库里，那就得写自定义 Dataset。写法其实不复杂，核心就是让__getitem__返回(图像张量, 标签)，__len__返回样本数。我一般建议先走 ImageFolder，实在不行再自定义。
3.2 尺寸和增强参数：按 B0 的 224 来还是按 B5 的 456 来
图像尺寸直接决定了两个东西：显存占用和模型能看到的细节。EfficientNet 的预训练权重是在固定分辨率下训出来的，迁移到自己的数据集时最稳的做法是沿用 ImageNet 时代的预处理方式。
train_transform = transforms.Compose([ transforms.RandomResizedCrop(224, scale=(0.7, 1.0)), transforms.RandomHorizontalFlip(), transforms.ColorJitter(0.3, 0.3, 0.3), 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]) ])先说RandomResizedCrop(224, scale=(0.7, 1.0))：从原图随机裁一块面积占比 70% 到 100% 的区域，再缩放到 224。这是我做细粒度分类时的习惯，scale 下限设 0.7 而不是默认的 0.08，是为了别把目标裁得太碎。如果你的数据集里目标本身就很小，scale 下限可以降到 0.5，但要做好精度波动的准备。
ColorJitter(0.3, 0.3, 0.3)是亮度、对比度、饱和度的抖动幅度，这个参数看数据。医疗影像、卫星图这种对颜色敏感的任务不建议加，掉落很多。
val 那边Resize(256) + CenterCrop(224)是标准做法，理由很简单：验证集要的是稳定可复现的指标，不能用随机裁剪，否则每次验证结果都不一样。至于为什么要多 resize 到 256 再裁 224，是为了让裁剪区域包含更多的上下文，这是 ImageNet 训练时代留下的经验值。
Normalize 的均值和方差必须用 ImageNet 的[0.485, 0.456, 0.406]和[0.229, 0.224, 0.225]。这不是玄学，是因为预训练权重的数据分布就是按这个标准归一化的，你换了均值方差，相当于给模型看了一张它不认识的图像。等到你哪天从头训练，才需要重新统计自己数据集的均值和方差。
4. 用预训练 EfficientNet 训练自己的数据：核心脚本与参数解释
4.1 加载预训练权重：num_classes 参数和 _fc 替换两种写法
EfficientNet-Pytorch 的杀手锏是from_pretrained，一行代码就能拿到在 ImageNet 上训好的权重。但 ImageNet 有 1000 类，你自己的数据集可能只有 10 类，分类头必须换。有两种写法：
from efficientnet_pytorch import EfficientNet import torch.nn as nn # 写法一：传 num_classes，库内部自动替换分类头 model = EfficientNet.from_pretrained('efficientnet-b0', num_classes=10) # 写法二：先加载默认 1000 类权重，再手动换掉 _fc model = EfficientNet.from_pretrained('efficientnet-b0') model._fc = nn.Linear(model._fc.in_features, 10)两种写法等价，区别在于写法二给了你换分类头的自由度。比如你想在分类头前面加一个 Dropout 层防止过拟合，或者想换两层 MLP 而不是单层线性层，写法二更容易操作。写法一适合不想折腾默认结构的场景。
注意看model._fc.in_features这行——它拿到的是分类头输入维度，B0 是 1280，B5 是 2048，不同模型不一样。不要自己硬编码 1280，万一换模型就废了。
from_pretrained第一次调用会去下载权重文件，这一步可能很慢甚至会超时。我的习惯是先把权重想办法下载好，放进本机缓存目录，比如 Linux 下的~/.cache/torch/hub/checkpoints/，后续再跑就不会卡在下载那一步。下载完的文件名要和代码里预期的一致，不然还是会重新下载。
4.2 训练循环的五段式写法与超参设定
训练脚本看着长，核心其实就是五个动作：前向计算、算 loss、反向传播、更新参数、清零梯度。我给一个自己改过的精简版：
import torch import torch.nn as nn from torch.utils.data import DataLoader device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = model.to(device) criterion = nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters(), lr=1e-3, momentum=0.9, weight_decay=1e-4) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=30) def train_one_epoch(model, loader, optimizer, criterion): model.train() total_loss, correct, total = 0.0, 0, 0 for images, labels in loader: images, labels = images.to(device), labels.to(device) optimizer.zero_grad() # 第一步：清空上一轮的梯度 outputs = model(images) # 第二步：前向计算 loss = criterion(outputs, labels) # 第三步：算损失 loss.backward() # 第四步：反向传播 optimizer.step() # 第五步：更新参数 total_loss += loss.item() * images.size(0) correct += (outputs.argmax(1) == labels).sum().item() total += labels.size(0) return total_loss / total, correct / total逐段说：optimizer.zero_grad()是新手最容易漏的。PyTorch 的梯度是累积的，不清空的话上一批数据的梯度会叠加到这一批上，loss 曲线会像锯齿一样乱七八糟。loss.item() * images.size(0)是在算总 loss，因为loss.item()是当前 batch 的平均 loss，乘上样本数才能算回总数，最后除以总样本数才是全 epoch 的平均。outputs.argmax(1)取每个样本预测概率最大的类别索引，跟 labels 逐位比较，算对多少个。
再说超参。我这里用的是 SGD 加动量而不是 Adam，原因是迁移学习场景下 SGD 的泛化性能通常比 Adam 稳。如果你习惯用 Adam，学习率要调小一个数量级，大约1e-4起步，否则前几个 epoch 的 loss 会跳得没法看。weight_decay 是权重衰减，相当于 L2 正则，1e-4 是分类任务的常用值。
CosineAnnealingLR是余弦退火学习率，T_max 设成你总 epoch 数，学习率会从初始值余弦降到接近 0。它的好处是后期步长小，更容易在最优解附近收敛。不想用调度器的话，可以用固定学习率加ReduceLROnPlateau，验证集 loss 不降了再减半，效果也不差。
如果数据集很小（几千张），我一般会在替换分类头之后冻结骨干网络，只训练分类头。冻结的方法很简单：
for param in model.parameters(): param.requires_grad = False for param in model._fc.parameters(): param.requires_grad = True这样前几个 epoch 相当于只训一个线性分类器，速度快很多。跑完一遍再解冻全部参数，用更小的学习率 3e-4 做全面微调。这是迁移学习的标准操作，能明显降低过拟合风险。
4.3 训练中要盯的四个监控点
训练不是把脚本跑起来等结果，要盯着关键指标判断有没有出事。我在训练时主要看四个点：
第一是 train loss 有没有下降。如果前三四个 epoch 一点不降，先怀疑学习率太大太小，再看数据有没有问题。第二是验证集准确率和训练集准确率的距离。两者差得大就是过拟合，差得小甚至验证比训练还高，有可能是验证集太难或者数据划分有泄漏。第三是每个 epoch 的耗时。如果前后波动超过 20%，可能存在 CPU 数据加载瓶颈或者 GPU 降频。第四是学习率的实际变化曲线，配合 loss 曲线能看出调度器有没有生效。
监控手段我习惯用一个极简的 print 加 TensorBoard：
from torch.utils.tensorboard import SummaryWriter writer = SummaryWriter('runs/efficientnet_b0') # 每个 epoch 结束时 writer.add_scalar('loss/train', train_loss, epoch) writer.add_scalar('acc/val', val_acc, epoch)TensorBoard 不是必须的，但看曲线比看一屏数字直观得多。跑完以后在终端执行tensorboard --logdir runs，浏览器打开本地 6006 端口就能看。
5. 自定义数据集训练最容易翻车的 5 个坑：现象、原因与排查
5.1 输入分辨率不对导致推理报错或精度崩盘
这个坑我踩过不止一次。现象有两种：一种是训练都能跑，但推理时一 forward 就报size mismatch；另一种更隐蔽，不报错，但验证集准确率一直停留在随机水平（比如 10 类稳定在 10% 左右）。
原因通常是训练和推理用的预处理不一致。比如训练时用RandomResizedCrop(224)，推理时直接Resize(224)，没有 CenterCrop，输入尺寸是 224 不假，但图像内容分布完全不同。另一种情况是用灰度图训练，EfficientNet 的 stem 是 3 通道卷积，单通道图要么报错要么重复成三通道但分布不对。
解决：把训练和推理的 transform 完全统一，val_transform 怎么写，推理时就用同一套；检查输入图像是不是 RGB，灰度图要convert('RGB')。我最开始给一批灰度工业图训练时，忘了转三通道，loss 一直不降，检查了一天最后发现是Image.open默认读出来是 L 模式。
5.2 预训练权重下载卡住：缓存目录与手动放置
现象：from_pretrained执行后终端卡住不动，或者报超时错误，进度条永远是 0%。
原因：预训练权重存在国外的对象存储服务上，网络环境不好的时候非常容易超时。尤其是公司内网加防火墙的情况，基本必卡。
解决：把权重文件提前下好，放到本机缓存目录。Linux 下缓存路径是~/.cache/torch/hub/checkpoints/，Windows 在C:\Users\<用户名>\.cache\torch\hub\checkpoints\。文件名必须是代码里预期的那个（从报错信息里能看到）。放好之后再执行from_pretrained，库会先检查本地缓存，有就直接加载不再下载。另一个办法是给from_pretrained传url_map参数，把默认下载地址指向本地文件路径。这个改动只影响加载逻辑，不影响模型结构。
5.3 分类头维度写错：loss 不降、acc 稳定在 1/n
现象：训练 loss 虽然在下降但很慢，验证集准确率稳如泰山地停在类别数的倒数附近。
原因：分类头输出维度和真实类别数对不上。最常见的是用EfficientNet.from_pretrained('efficientnet-b0')后忘了改_fc，模型输出还是 1000 维，而你的 label 只有 10 个。CrossEntropyLoss 不会报错，因为 1000 维输出照样能跟 0-9 的索引计算 loss，但这辈子不可能收敛到好效果。
解决：训练脚本里加两行断言，比事后排查省时间得多：
assert model._fc.out_features == len(train_set.classes) print(f'分类头输出维度: {model._fc.out_features}, 类别数: {len(train_set.classes)}')用from_pretrained直接传num_classes也不会出这个错，但手动替换_fc的时候就容易漏。反正断言一行代码的事，成本几乎为零。
5.4 显存溢出：batch size 与梯度累积的取舍
现象：训练到一半爆CUDA out of memory，程序直接崩掉，前面几个小时的训练白跑。
原因：batch size 设太大，或者模型太大。B0 在 224 分辨率下 batch 32 大约 8G 显存，换成 B5 的 456 分辨率，同样 batch 32 至少 20G 往上。很多人从 ResNet 切过来，保持了原来的 batch size，直接爆显存。
解决：先降 batch size 到 8 或 4，跑通再说。如果 batch 太小影响收敛（尤其是有 BN 层的时候），用梯度累积模拟大 batch：
accum_steps = 4 # 模拟 batch size = 实际 batch size * 4 optimizer.zero_grad() for i, (images, labels) in enumerate(loader): outputs = model(images.to(device)) loss = criterion(outputs, labels.to(device)) / accum_steps loss.backward() if (i + 1) % accum_steps == 0: optimizer.step() optimizer.zero_grad()注意 loss 除以了accum_steps，这是关键，不然累积了 4 个 batch 的梯度，等效学习率变成原来的 4 倍。另外梯度累积和真正的大 batch 不完全等价，BN 层在每个小 batch 上各自统计均值方差，所以精度会有微小的波动，但多数任务下差别不大。
5.5 BN 层在 batch 很小的时候玄学报错
现象：batch_size 设为 1 或 2 时，训练几轮后 loss 变成 NaN，或者直接报错，没有任何预兆。
原因：EfficientNet 的 MBConv 里大量使用 BN 层。BN 在统计均值方差时，如果 batch 内只有 1 个样本，方差为 0，归一化分母直接除以 0，数值变得不稳定。这个在 ResNet 里也存在，但 EfficientNet 对 batch size 更敏感，因为浅层 MBConv 的通道数少，统计量更容易偏差。
解决：batch size 至少设 4，稳妥起见 8 以上。如果你的显存只允许 batch 1-2，那就冻结 BN 层，让它们用预训练权重里带过来的统计量：
def freeze_bn(model): for m in model.modules(): if isinstance(m, torch.nn.BatchNorm2d): m.eval()m.eval()之后 BN 不再更新 running_mean 和 running_var，只使用加载预训练权重时的全局统计量。这样小 batch 下训练也能稳定，前提是你的数据集分布和 ImageNet 不要差太远。数据域差异大的话，冻结 BN 会影响精度，不如直接升级显存。
6. 训练完别急着收工：混淆矩阵、迁移验证与 ONNX 导出
训练完成后的第一步不是看最终准确率数字，而是跑一遍混淆矩阵。准确率会骗人，比如类别严重不均衡时，90% 的准确率可能只是把所有样本都预测成了多数类。混淆矩阵能直接看出哪些类别在互相打架。
from sklearn.metrics import confusion_matrix import torch model.eval() y_true, y_pred = [], [] with torch.no_grad(): for images, labels in val_loader: images = images.to(device) outputs = model(images) y_true.extend(labels.tolist()) y_pred.extend(outputs.argmax(1).cpu().tolist()) print(confusion_matrix(y_true, y_pred))看混淆矩阵有两个重点：一是对角线是否明显高于非对角线，二是哪两个类别最容易混淆。我做过一个缺陷检测项目，混淆矩阵显示“划痕”和“脏污”互相误判率接近 40%，最后发现是标注标准不统一——同一种缺陷有人标划痕有人标脏污。这种问题靠调模型是解决不了的，回去整理数据才是正路。
接着做迁移验证，拿训练时完全没见过的数据来测。我习惯准备一个额外的图片文件夹，跑一遍单张推理，检查 top-1 和 top-5：
from PIL import Image img = Image.open('new_image.jpg').convert('RGB') img = val_transform(img).unsqueeze(0).to(device) with torch.no_grad(): probs = torch.softmax(model(img), dim=1) top5 = torch.topk(probs, 5) for i in range(5): idx = top5.indices[0][i].item() print(f'{val_set.classes[idx]}: {top5.values[0][i].item():.4f}')如果新图片的 top-1 不对但 top-5 里有正确答案，说明模型学到了相近特征但区分度不够；如果 top-5 里压根没有正确答案，而且置信度都很低，那这模型上线是要出事的。
最后做 ONNX 导出。EfficientNet 转 ONNX 有个老坑：Swish 激活函数是自定义的，导出时如果 opset 版本太低，有些推理引擎不认。我一般用 opset 11 以上：
dummy_input = torch.randn(1, 3, 224, 224).to(device) torch.onnx.export(model, dummy_input, 'efficientnet_b0.onnx', opset_version=11, input_names=['input'], output_names=['output'])导出后先用 ONNX Runtime 跑一遍验证结果一致再交付。转完模型后我在 mobile 和 server 端各自测过一遍推理速度，B0 在 CPU 上的单张推理大概几十毫秒，取决于具体硬件；B5 在 CPU 上就比较吃力了，能上 GPU 就上 GPU。有一段时间我只跑 Eigen 那个比较极端的网络，B7 在 CPU 上直接慢得没法用，后来给自己定了个规矩：移动端部署最多用 B0，服务端也只在 GPU 上考虑 B3 以上。希望这些路能帮你少踩几个坑，祝顺利。
本文还有配套的精品资源，点击获取
