EfficientNet-Pytorch实战：用预训练权重训练自己的图像分类模型

简介：面向希望快速将 EfficientNet 迁移到自定义图像分类任务的开发者，这份源码包提供了一个极简可运行的演示项目，并明确给出了数据集的组织方式——将训练集与测试集分别按不同类别文件夹存放图片，与主流分类任务的数据加载习惯一致。项目核心基于 EfficientNet-Pytorch 实现，用户只需整理好数据、修改脚本中的类别数和路径设置，即可启动训练与测试；脚本中还预留了是否自动下载预训练权重的选项，便于在无网络环境下手动加载。资源包共 5 个文件，包括 4 个 Python 脚本与 1 个 Markdown 说明文档，分别用于模型定义、工具函数、主训练流程和使用指导，整体压缩后仅 10KB，轻量简洁。已有 3780 人学习下载，适合具备 Python 和 PyTorch 基础、希望在 EfficientNet 上快速尝试自备数据的读者，通过源码中的注释和示例可掌握数据加载、模型调用、训练测试及参数调整的完整思路。

1. 用 EfficientNet-Pytorch 训练自己的数据集：先想清楚三个问题再动手

手里攒了一批图，想跑通一个像样的图像分类模型，EfficientNet-Pytorch 这个方案恐怕是你绕不开的起点。它把 EfficientNet 的 PyTorch 实现和 ImageNet 预训练权重打包成一个开箱即用的库，你在from_pretrained里传一个num_classes，它就把分类头换好，剩下的活就是把自己的数据集按目录摆好、写个训练循环，然后等指标一点点涨起来。这个标题解决的是「怎么把 ImageNet 上训好的通用特征迁移到自己的小数据集上」这一件事，适合还没跑通完整流程的人，也适合在换数据集时想有个稳定基线的熟手。我先把三个问题放前面：你的数据集有多大、你打算用哪个规模的 EfficientNet、以及你愿不愿意花时间调学习率。这三个问题没想清楚，后面跑出来的结果大概率是玄学。

2. EfficientNet 到底强在哪：复合缩放与 B0 到 B7 的选型逻辑

2.1 复合缩放：为什么不是单方向加深度或加宽度

EfficientNet 的核心贡献是提出了一套比「无脑加深加宽」更精细的模型缩放方法。早在 2019 年之前，放大网络的做法无非是加深 ResNet 层数、把卷积通道数翻倍、或者把输入分辨率从 224 提到 320，每次只动一个维度。EfficientNet 用神经架构搜索先挖出一个基线网络 B0，然后把深度、宽度、分辨率三个维度绑在一起，用一个复合系数 φ 统一控制。φ 每加 1，三个维度按固定比例同步放大，FLOPs 大约翻一倍。

这套思路的实际收益是：在同样的 FLOPs 预算下，EfficientNet 的精度比当时主流的 ResNet 和 DenseNet 都高出一截。对我们这种做落地的人来说，它意味着两件事。第一，想冲精度不一定要换模型族，把 B0 换成 B2、B4 就能看到收益；第二，每一档模型的输入分辨率是定死的，B0 是 224，B4 以上是 380，切换档位时要连带着把数据加载那一段一起改，不是换个名字就行。

2.2 EfficientNet-Pytorch 这个库和 torchvision 里自带版本怎么选

如果你装的是新版 PyTorch，torchvision.models里其实已经提供efficientnet_b0到efficientnet_b7，还带训练好的权重。那为什么还要单独用efficientnet_pytorch这个库？主要是 API 风格差别。torchvision的写法是先取模型再替换分类头，你得自己数一下全连接层的输入维度；而 EfficientNet-Pytorch 的思路更偏向微调场景，EfficientNet.from_pretrained('efficientnet-b0', num_classes=你的类别数)一行就把预训练权重和新的分类头都处理好了。

另一个差别在权重来源。EfficientNet-Pytorch 仓库里放的是原论文作者用 TensorFlow 训练后转出来的权重，和 TF 官方实现完全对齐；torchvision 那版是用 PyTorch 重新训练的，精度略有差异但在多数任务上感觉不出来。做工程时我不太纠结这两个版本谁的基准精度高零点几，真正决定选型的因素是可控性。EfficientNet-Pytorch 的源码干干净净，MBConvBlock里 SE 模块、drop connect、Swish 激活都写在明处，想改哪一层直接改 Python 文件就行。torchvision 那版为了兼容性封装得比较死。我自己的习惯是：能用这个库解决的就不绕路，毕竟标题就叫 EfficientNet-Pytorch。

2.3 B0 到 B7 怎么选：显存、数据量和训练时间的三角权衡

先看一张我从实践里归纳的规格表，再谈选择逻辑：

选择建议：第一次跑通流程，一律从 B0 开始。不是因为 B0 最好，而是因为 B0 输入尺寸小、迭代快，能在一两个小时内验证你的数据管线、标签顺序和学习率是否正常。等整个流程跑通，再按数据量往上爬。这里有一个常见误区：数据集只有几千张，却一上来就用 B5，结果没过拟合先被显存卡死。迁移学习的收益主要来自 ImageNet 预训练的特征提取器，小数据集上 B0 到 B2 的精度差异往往不超过两个点，但训练时间的差异是几倍。量力而行，把省下来的时间拿去调数据增强反而更值。

3. 把数据集喂给 EfficientNet：目录组织、拆分脚本与 transform 边界

3.1 用 ImageFolder 组织自己的数据集：一份可以直接套用的目录结构

EfficientNet-Pytorch 本身不关心数据格式，它只吃 PyTorch 的Dataset。最常见、最不容易出错的做法是用torchvision.datasets.ImageFolder，它要求目录按「类名即文件夹名」的方式组织。

data/ ├── train/ │ ├── cat/ │ │ ├── cat_001.jpg │ │ └── cat_002.jpg │ └── dog/ │ ├── dog_001.jpg │ └── dog_002.jpg └── val/ ├── cat/ │ ├── val_cat_001.jpg │ └── ... └── dog/ └── ...

如果你手里的图是散落在一个大文件夹里、靠 Excel 或 CSV 记录标签，先写个十行脚本按标签拆到上面的目录结构，不要试图在 Dataset 里做复杂映射，后面排查数据错乱会非常痛苦。拆分时注意两个细节。一是类别名的排序，ImageFolder会按文件夹名字母顺序生成类别索引，cat是 0、dog是 1，这个顺序在训练和推理时必须保持一致；二是划分比例，我一般按 8:2 或者 9:1 切，验证集至少保证每类有 50 张以上，太少的话验证指标抖动会大到没法判断模型好坏。

3.2 train 和 val 的 transform：区别不止是数据增强

ImageNet 预训练权重对输入有硬性要求：归一化的 mean 和 std 必须是 ImageNet 统计值(0.485, 0.456, 0.406)和(0.229, 0.224, 0.225)，尺寸至少是 224 的整数倍。这是新手最容易踩的暗坑——用自己数据集的 mean/std 做了归一化，结果加载预训练权重后特征分布对不上，loss 怎么调都降不下来。

我常用的两套 transform 如下：

from torchvision import transforms # 训练集：随机裁剪 + 翻转 + 归一化 train_transform = transforms.Compose([ transforms.RandomResizedCrop(224, scale=(0.08, 1.0)), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2), 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的scale=(0.08, 1.0)是借鉴 ImageNet 训练的标准做法，相当于给模型看不同缩放比例下的目标，对提升泛化能力很有帮助。验证集不引入随机性，保证每次评估的输入完全确定，指标才可比较。如果你的数据集特别小（每类一两百张），可以加RandomRotation(15)或RandomAffine，但要克制，增广过猛会让模型学不到真实分布。还有一个细节：如果从 B0 切到 B4，上面的224和256要同步改成380和420，只改模型不改输入尺寸，等于让模型在错误的分辨率下工作。

3.3 DataLoader 参数：num_workers 和 pin_memory 是白送的加速

数据准备的最后一步是包一层DataLoader：

from torch.utils.data import DataLoader from torchvision import datasets train_dataset = datasets.ImageFolder('data/train', transform=train_transform) val_dataset = datasets.ImageFolder('data/val', transform=val_transform) train_loader = DataLoader( train_dataset, batch_size=32, shuffle=True, num_workers=4, pin_memory=True, drop_last=True ) val_loader = DataLoader( val_dataset, batch_size=64, shuffle=False, num_workers=4, pin_memory=True )

说三个参数的实际意义。num_workers是加载图片的子进程数，Windows 上建议设为 0 或 2，Linux 上设成 CPU 核心数的一半比较稳；如果训练时发现 GPU 利用率时常掉到 50% 以下，优先怀疑num_workers太小，而不是模型有问题。pin_memory=True让数据加载到锁页内存，配合.cuda()能减少一次拷贝，属于不花钱的加速。drop_last=True只在训练集上开，作用是扔掉最后一个不完整的 batch，避免 BN 层在 batch size 过小时统计出偏差；验证集不设是因为评估时需要遍历全部样本。如果你的数据集类别不平衡，还需要在DataLoader里塞sampler，这个放到第五部分讲。

4. 训练脚本逐段拆解：用 EfficientNet-Pytorch 跑通第一个 epoch

4.1 加载预训练模型：from_pretrained 背后做了什么

代码先走最核心的两行：

import torch from efficientnet_pytorch import EfficientNet num_classes = 10 # 换成你自己的类别数 model = EfficientNet.from_pretrained('efficientnet-b0', num_classes=num_classes)

from_pretrained会按efficientnet-b0这个名字去下载对应的 ImageNet 预训练权重，然后构建模型网络结构，把最后一层全连接换成输出维度等于num_classes的新层。这个库的一个优势在于，替换分类头的同时保留了前面所有卷积层和 BN 层的参数，你得到的是一个「特征提取器冻结、分类头随机」的迁移学习模型。如果只想用结构不想用预训练权重，就用EfficientNet.from_name('efficientnet-b0', num_classes=num_classes)，这个接口不会触发下载。

下载权重时需要注意：首次运行会打印下载进度，如果反复卡住或报网络连接错误，去仓库的 releases 页面手动下载对应文件，放进~/.cache/torch/checkpoints/目录即可，名字要和库内部期望的文件名保持一致。加载完成后，建议先把模型切成训练模式再开始，PyTorch 的 Dropout 和 BN 在不同模式下行为不同，这个细节容易忽略，但直接影响训练曲线。

4.2 优化器与学习率：SGD 还是 Adam，迁移学习里怎么选

EfficientNet 原论文在 ImageNet 上用的是 RMSProp，但迁移学习场景下很少有人照搬。常见做法是两个选择：SGD 加 momentum，或者 AdamW。我的经验是，微调不超过几千张的小数据集时，SGD 的表现更稳定，不容易在训练初期就把预训练特征冲坏；Adam 家族胜在收敛快，适合数据集比较大、你不想仔细调学习率的时候。

optimizer = torch.optim.SGD( model.parameters(), lr=0.01, # 迁移学习常用初始值 momentum=0.9, weight_decay=1e-4 ) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode='min', factor=0.5, patience=5, verbose=True ) criterion = torch.nn.CrossEntropyLoss()

lr=0.01是 SGD 微调的常见起点，但如果你发现前几个 epoch 的 loss 完全不降甚至在涨，优先把 lr 降到 0.003 或 0.001，这是迁移学习里最常翻车的参数。ReduceLROnPlateau是「等验证 loss 不再下降就砍一半学习率」的调度器，patience=5表示连续 5 个 epoch 没有改善才降。相比固定步长衰减，这种方式更省心，不用提前猜总共要训多少轮。如果你的显存充裕，想更快收敛，可以换成torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=30)，但 Cosine 的缺点是要预先知道总 epoch 数，改起来比 Plateau 繁琐。

4.3 训练循环与验证循环：一套可以直接抄的主循环

def train_one_epoch(model, loader, optimizer, criterion, device): 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) _, preds = torch.max(outputs, 1) correct += (preds == labels).sum().item() total += labels.size(0) return total_loss / total, correct / total def validate(model, loader, criterion, device): model.eval() total_loss, correct, total = 0.0, 0, 0 with torch.no_grad(): for images, labels in loader: images, labels = images.to(device), labels.to(device) outputs = model(images) loss = criterion(outputs, labels) total_loss += loss.item() * images.size(0) _, preds = torch.max(outputs, 1) correct += (preds == labels).sum().item() total += labels.size(0) return total_loss / total, correct / total

两个循环的结构完全对称，区别只在三处：训练循环里要调用model.train()、显式反向传播并更新梯度；验证循环里要调用model.eval()并包一层torch.no_grad()。model.eval()会关闭 Dropout 并让 BN 层使用训练阶段统计好的均值方差，这一步漏掉会导致验证精度远低于训练精度，而且找不到原因。torch.max(outputs, 1)第二个返回值就是每个样本预测的类别索引，和 label 逐位比较后求和，就是正确样本数。

4.4 主函数里把这些串起来：记录最优模型是唯一需要上心的细节

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = model.to(device) best_acc = 0.0 num_epochs = 30 for epoch in range(num_epochs): train_loss, train_acc = train_one_epoch( model, train_loader, optimizer, criterion, device ) val_loss, val_acc = validate(model, val_loader, criterion, device) scheduler.step(val_loss) if val_acc > best_acc: best_acc = val_acc torch.save(model.state_dict(), 'best_model.pth') print(f'Epoch {epoch+1}/{num_epochs} | ' f'train_loss={train_loss:.4f} train_acc={train_acc:.4f} | ' f'val_loss={val_loss:.4f} val_acc={val_acc:.4f}')

建议每个 epoch 都在验证集上评估一次，然后用「验证集准确率最高」的那个状态作为最终模型，而不是训完最后一个 epoch 直接拿去用。原因很简单：深度学习训练到后期，验证 loss 往往在一两个点之间反复横跳，最后一个 epoch 不一定是泛化最好的状态。上面这段代码里唯一需要你自己扩展的是「调度器要传训练 loss 还是验证 loss」：ReduceLROnPlateau的mode='min'配合val_loss传参是标准做法，让学习率随验证集表现衰减，比随训练 loss 衰减更可靠。如果数据量比较小，可以把num_epochs从 30 提到 50，配合早停（连续 10 个 epoch 验证 acc 不升就停）能节约不少时间。

5. 训练 EfficientNet 的避坑清单：五个高频翻车点与排查路径

写这一章之前我想先说明，下面每一条都是有人在群里问过、我在自己机器上也撞过的真实问题，按「现象 → 原因 → 解决」的顺序写，方便你对着排查。

5.1 from_pretrained 卡在下载：权重文件到达不了本地

现象：脚本执行到EfficientNet.from_pretrained就不再往下走，或者报网络连接相关错误，重试多次依旧如此。

原因：预训练权重文件托管在 GitHub Releases 上，下载环节受网络环境影响较大，不是你代码的问题。

解决：最简单的方案是手动下载。先去仓库的 releases 页面找到和efficientnet-b0对应的权重文件名，用浏览器或下载工具存到~/.cache/torch/checkpoints/（Linux）或C:/Users/你的用户名/.cache/torch/checkpoints/（Windows），确保文件名完全一致，再重新运行脚本。另一个备用方案是用EfficientNet.from_name初始化随机权重的网络，再自己load_state_dict，但这样要手动对齐权重键名，不如前者省事。我一般会把下载好的权重备份一份到项目目录里，下次直接指定路径加载，省得每个环境都重新下载。

5.2 显存溢出：B4 以上在同一张卡上跑不动

现象：RuntimeError: CUDA out of memory，即使把 batch_size 降到 8 依然报错。

原因：B4 的输入分辨率是 380，同一张图占的显存大约是 B0 的三倍；再加上模型本身参数和激活值变大，同样显存下能跑的 batch 远小于 B0。

解决：先确认你选的档位对应的分辨率，把RandomResizedCrop和CenterCrop的尺寸改到匹配值，然后从 batch_size=8 开始逐步往上试。如果还想加大有效 batch，可以加梯度累积：

accumulation_steps = 4 # 每 4 个 batch 更新一次梯度 optimizer.zero_grad() for i, (images, labels) in enumerate(train_loader): outputs = model(images.to(device)) loss = criterion(outputs, labels.to(device)) loss = loss / accumulation_steps # 平均一下，等价于放大 batch loss.backward() if (i + 1) % accumulation_steps == 0: optimizer.step() optimizer.zero_grad()

梯度累积的代价是训练时间变长，但能把有效 batch 放大到任何你想要的数值，是显存受限时最实用的解法。注意loss / accumulation_steps这步不能省，否则梯度绝对值会随累积步数线性变大，学习率的语义就变了。

5.3 训练 loss 不降或直接变成 NaN

现象：前几个 epoch 的 loss 在 2.3 附近纹丝不动（10 分类随机猜测的交叉熵就是 ln10≈2.3），或者干脆变成 NaN。

原因：学习率过大。预训练权重虽然给了好起点，但分类头是随机初始化的，SGD 默认的 0.01 在部分数据集上会直接把随机初始化的分类头梯度冲飞，进而反传到 backbone 导致整体发散。

解决：先把lr降到 0.001 试一轮，确认 loss 能稳定下降再逐步加回去。如果你用的是 Adam 系优化器，初始学习率建议直接设在 0.0001 到 0.0003 之间，Adam 对学习率更敏感，照搬 SGD 的 0.01 翻车概率很高。另外检查一下有没有在 transform 里用了错误的归一化参数，mean/std 不对也会让特征分布异常，表现出来同样是不收敛。

5.4 训练集指标高、验证集指标上不去：过拟合的典型形态

现象：训练 acc 在第 15 个 epoch 到 98%，验证 acc 卡在 82% 附近并开始抖动，继续训练两者差距越拉越大。

原因：数据量太小 + 模型容量过剩 + 正则化不足。B0 已经是几百万参数，几千张图很容易被它硬背下来，Dropout 默认 0.2 在迁移学习里偏低。

解决：三件事按优先级做。第一，确认用的是 B0 而不是更大的档位；第二，把from_pretrained的dropout_rate参数调高到 0.3 或 0.5：

model = EfficientNet.from_pretrained( 'efficientnet-b0', num_classes=num_classes, dropout_rate=0.5 )

第三，给训练集 transform 加更强的增强，比如RandomErasing或AutoAugment。还有一个容易忽略的点：drop_connect_rate是 EfficientNet 特有的正则化参数，这个库默认按深度分配，B0 默认从 0 线性加到 0.2，一般不用动。如果篡改卷积层内部的 drop_connect 参数，反而容易把训练搞崩。

5.5 类别不平衡导致验证指标虚高

现象：总准确率 92%，但看每一类的召回率，其中一类只有 40%；把模型输出的混淆矩阵打出来，发现模型把所有样本都往大类上猜。

原因：CrossEntropyLoss默认按样本均匀计算，大类样本多、贡献梯度大，小类被淹没。

解决：用加权交叉熵，在CrossEntropyLoss里传入权重向量：

class_weights = torch.tensor([1.0, 2.0, 5.0], device=device) # 按类别样本数反比设置 criterion = torch.nn.CrossEntropyLoss(weight=class_weights)

权重怎么算：每类权重 = 总样本数 / (类别数 * 该类样本数)。也可在数据加载端用WeightedRandomSampler让采样器多抽小类样本，但在分类任务里改 loss 权重比改采样器更直接，因为采样器只影响采样频率，不影响损失函数的梯度方向。改完权重后，验证参考指标要换成每一类的 precision 和 recall，而不是只看总 acc。

6. 把训练结果做扎实：从 B0 往上换模型、混淆矩阵验收与学习率曲线诊断

换更大的模型是路径依赖最明显的下一步。B0 跑通之后，如果验证集 acc 稳定且没过拟合，说明模型还有容量余量，按「B0 → B2 → B3」的顺序逐步升级，同时把输入分辨率从 224 提到 260 再提到 300。每次升级后 batch_size 减半，学习率保持不动，观察第一个 epoch 的下降速度。如果换到某个档位后验证指标不升反降，回退上一档，不要硬撑。

训练结束后别急着看总 acc 就收工，我习惯花十分钟画一张混淆矩阵：

import numpy as np from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay import matplotlib.pyplot as plt model.load_state_dict(torch.load('best_model.pth')) model.eval() all_preds, all_labels = [], [] with torch.no_grad(): for images, labels in val_loader: outputs = model(images.to(device)) _, preds = torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.numpy()) cm = confusion_matrix(all_labels, all_preds) disp = ConfusionMatrixDisplay(cm, display_labels=val_dataset.classes) disp.plot(xticks_rotation=45) plt.savefig('confusion_matrix.png', dpi=150, bbox_inches='tight')

看混淆矩阵时要找两类特征：对角线明显偏暗的类别（经常被错认成什么）和某一行整体很暗的类别（模型根本不认识）。这两个线索直接指向数据问题，比如两个类别标注混乱、某一类图片摄自不同场景。我遇到过一次「黑色拉布拉多」和「罗威纳犬」混淆严重，翻原始图片发现训练集里罗威纳的照片全是夜间拍的，模型学到的其实是亮度差异而不是品种特征，这就是数据问题，不是模型问题。

最后养成一个习惯：训练结束后把每个 epoch 的 train_loss、val_loss 和 lr 抽出来画在一张图里。val_loss 在降、train_loss 还在降且两者没有分叉，说明训练状态健康；val_loss 先降后升，train_loss 一路向下，过拟合拐点就是 val_loss 的最低点，早停应该设在那附近；曲线锯齿状剧烈抖动，优先怀疑 lr 太大或 batch_size 太小。这个曲线比 TensorBoard 里那些花哨的嵌入可视化有用得多，也是我每次训练完必看的第一张图。EfficientNet 的训练流程并不复杂，但每一步的参数选择都在决定最终精度上限，把数据集、学习率和档位这三件事调顺了，剩下的交给时间就好。希望帮到你。

本文还有配套的精品资源，点击获取