简介：面向需要在EfficientNet上训练自定义图像分类数据集的深度学习开发者，这套演示源码以EfficientNet-Pytorch为基础，完整覆盖了从数据准备、目录组织到训练测试的全流程。当前版本包含5个文件，包括4个Python脚本（如EfficientNet模型实现、工具函数、示例训练入口）和1个Markdown说明文档，压缩包仅10KB，结构精简且便于阅读。已有3777人学习下载，适合计算机视觉初学者和希望快速将EfficientNet迁移到自有任务的工程师参考。通过源码，读者可掌握数据集按类别子目录放入train/test的规范，了解如何通过修改指定行号控制预训练权重的自动下载或手动放入权重文件夹，并参照注释调整数据路径、类别数等关键参数，最后运行示例训练脚本完成整套流程。README对脚本用途和运行步骤也有简要说明，能帮助快速定位修改点。

1. EfficientNet 迁移学习为什么值得放在自己的数据集上试一次

EfficientNet 的核心不是把网络做得更深或更宽，而是用复合缩放（compound scaling）同时调整深度、宽度和输入分辨率，在同等 FLOPs 下拿到比 ResNet 更低的错误率。放到实际项目里，这意味着用同样的显存预算，你可以跑一个容量更大的模型，或者用更小的模型达到接近的效果。这次要拆的演示工程，是基于 EfficientNet-Pytorch 实现的一套最小可运行的分类训练流程：你只要把图片按类目文件夹整理好，改几行参数，就能跑通从训练到测试的完整流程。它特别适合两类人：一类是刚开始做图像分类、想找一个可靠 baseline 的工程师；另一类是已经把 ResNet 跑熟，想看看 EfficientNet 在同等数据下是否真的更划算的人。下面我会从源码结构、数据组织、训练脚本参数、排错和模型导出五个层面把这个工程拆透。

2. 先看 EfficientNet-Pytorch 的源码结构与数据集契约

2.1 efficientnet 包里的 model.py 和 utils.py 各管什么

解压 efficientnet-pytorch-master 后，根目录下会看到 efficientnet 子文件夹，里面是 efficientnet_model.py、model.py、utils.py 这些核心文件。model.py 对外暴露的是 EfficientNet.from_pretrained 和 EfficientNet.from_name ，它负责把模型名字（如 'efficientnet-b0'）映射到具体的网络结构，并处理预训练权重的下载和加载。utils.py 则提供了一些工具函数，比如全局池化、批归一化的参数调整，还有一个 get_model_params 用来读取不同版本的缩放系数。实际写代码时你通常不会直接碰这两个文件，但理解它们的边界很重要：model.py 决定网络长什么样，utils.py 决定训练时哪些层要解冻、哪些层保持原样。如果你的自定义数据集很小，建议在调用 from_pretrained 时传 advprop=False ，因为 AdvProp 的权重期望输入做的是 ImageNet 的标准化，而不是普通 ToTensor 的归一化。

2.2 dataset 目录的层级必须严格按类目文件夹来

演示工程对数据集的要求非常具体：在 dataset 目录下建 train 和 test 两个子目录，每个子目录里再按类别建数字命名的文件夹，比如 1、2、3。图片就放在这些数字文件夹里，文件夹名会被当作类别的 label。这里的"数字命名"不是强制要求，但脚本里通常直接用 ImageFolder 读取，它会按文件夹名的字典序生成 class_to_idx 映射。你后续要在测试代码里还原标签名，就要保存这个映射表，否则预测结果就是一堆数字。

import os
from shutil import copy2
from sklearn.model_selection import train_test_split

# 假设原始数据在一个叫 raw 的目录下，每个子类一个文件夹
root = 'raw'
train_root = 'dataset/train'
test_root = 'dataset/test'

for cls in os.listdir(root):
    cls_path = os.path.join(root, cls)
    imgs = os.listdir(cls_path)
    train_list, test_list = train_test_split(imgs, test_size=0.2, random_state=42)
    for split_root, split_imgs in [(train_root, train_list), (test_root, test_list)]:
        target_dir = os.path.join(split_root, cls)
        os.makedirs(target_dir, exist_ok=True)
        for img in split_imgs:
            copy2(os.path.join(cls_path, img), os.path.join(target_dir, img))

这个脚本做的事情是：读原始目录下每个类别的图片列表，用 train_test_split 按 8:2 划分，然后复制到 dataset/train/类别号 和 dataset/test/类别号 下。 random_state=42 保证重复执行时划分结果一致，方便对比实验。如果你有现成的 train/val 目录，就跳过这一步，但要注意类别编号必须一致，训练和测试集里同一个类不能用不同编号。

2.3 ImageFolder 的隐式标签顺序与文件名无关

PyTorch 的 torchvision.datasets.ImageFolder 在构造时会扫描根目录下的子文件夹，然后按字母序对类别排序，再生成从 0 开始的索引。也就是说，如果你的文件夹叫 1、2、10，那么排序后是 1、10、2，类别 ID 就不是视觉上的顺序。这在推理时会造成困惑。常见做法是在训练完成后把 dataset.classes 保存成 json，推理时用同一个映射来翻译结果。

import json
from torchvision import datasets

train_set = datasets.ImageFolder('dataset/train')
with open('class_names.json', 'w') as f:
    json.dump(train_set.classes, f)
print(train_set.class_to_idx)

这段代码会把类别名列表写到 class_names.json 中。 train_set.class_to_idx 是一个字典，键是文件夹名，值是整数 label。后续测试脚本里可以通过 json.load 读回来，再结合模型输出的 argmax 索引查出真实类别名。很多新手在这里栽跟头，认为第 10 个文件夹的 label 就是 10，实际上可能是 0 也可能是 2，取决于字典序。

在动手训练之前，建议先用上面的脚本把类别映射固定下来，然后检查每个类别的图片数量。EfficientNet 的输入分辨率最低是 224（b0），如果你的图片大小差异极大，最好在 DataLoader 里统一做 resize 和中心裁剪，否则会引入不必要的噪声。

3. 训练脚本的改动点：从 line 13-22 到 line 169

3.1 三个最重要的参数：模型名、num_classes、数据路径

演示脚本里，line 13-22 就是为你准备的"配置文件"。改动时不建议直接改文件系统路径，而是把路径集中在这几行，后续换数据集只需要更新这里。

# efficientnet_sample.py 第 13-22 行，按自己的数据集修改
model_name = 'efficientnet-b0'      # 可换 b1-b7, 模型越大越吃显存
num_classes = 20                    # 自己的分类数量
train_dir = 'dataset/train'
test_dir = 'dataset/test'
batch_size = 32
num_epochs = 30
lr = 0.01
weight_decay = 1e-5
# 是否从 ImageNet 权重开始
pretrained = True

model_name 决定网络结构，b0 是基础版本，输入 224x224，参数约 5.3M，在单张消费级显卡上训练 20 类数据基本没有压力。如果你觉得精度不够，优先换 b1 或 b2，而不是直接上 b7，因为 b7 在相同 epoch 下收敛更慢。 num_classes 必须和 ImageFolder 扫描到的子文件夹数量一致，否则最后一层全连接的输出维度和标签对不上，训练时会报维度不匹配的错误。 lr=0.01 这个值对于从头训练的模型是合理的，但如果你从预训练权重开始，建议改成 0.001，后面会解释原因。

3.2 预训练权重的下载开关与断点续训

line 169 控制的是是否自动下载预训练模型。EfficientNet-Pytorch 官方实现里， from_pretrained 会从指定的 URL 下载权重到 eff_weights 文件夹。如果你在离线环境，就要提前把权重文件放进去，并把这里的逻辑改为读取本地权重。演示工程在 eff_weights 下放 pre-trained 模型也对应这个流程。

if pretrained:
    model = EfficientNet.from_pretrained(model_name, num_classes=num_classes, weights_path='eff_weights/efficientnet-b0-355c32eb.pth')
else:
    model = EfficientNet.from_name(model_name, num_classes=num_classes)

from_pretrained 内部会先把 ImageNet 分类层替换成你自己定义的 num_classes 的输出，然后加载其他部分的参数。 weights_path 可以指定本地文件，避免每次运行都去下载。如果自定义数据集的图片域和 ImageNet 差异很大（比如红外图、医学切片），我一般会保留预训练权重，但把 model._fc 层替换成 Dropout + Linear，并且前几个 epoch 冻结 model._blocks ，只训练最后的全连接，后面再解冻。演示工程里没有做这个分阶段操作，你可以在拿到 baseline 以后自己扩展。

3.3 训练主循环里的优化器和学习率策略

训练循环本身不复杂，但有几个参数值得注意。常见做法是使用 torch.optim.SGD 加动量，配合 CosineAnnealingLR 让学习率周期性下降。如果数据集只有几千张，训练十几个 epoch 就会过拟合，所以验证集上的 loss 反而更重要。

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=weight_decay)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)

for epoch in range(num_epochs):
    model.train()
    for inputs, labels in train_loader:
        inputs, labels = inputs.cuda(), labels.cuda()
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
    scheduler.step()

这里的 momentum=0.9 是分类任务默认值，能有效抑制梯度摆动。 CosineAnnealingLR 的 T_max 设为 num_epochs ，表示整个训练周期内学习率从初始值余弦下降到接近 0。如果你用了预训练权重，建议 lr 从 0.001 起步，并且把 weight_decay 设为 1e-4 而不是 1e-5 ，因为大模型在自定义数据上更容易过拟合，权重衰减更重一点能提升验证精度。

3.4 混合精度训练能省多少显存

如果你的显卡支持 CUDA 11 以上，可以在训练循环里加上 GradScaler 做混合精度，代码改动不多，但显存占用能减少约 40%。EfficientNet 的 batch norm 对精度比较敏感，所以只在卷积和全连接层开启 FP16 即可。

scaler = torch.cuda.amp.GradScaler()
for inputs, labels in train_loader:
    inputs, labels = inputs.cuda(), labels.cuda()
    optimizer.zero_grad()
    with torch.cuda.amp.autocast():
        outputs = model(inputs)
        loss = criterion(outputs, labels)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

autocast 只对内部算子做低精度计算，不需要手动改模型。 GradScaler 会在 loss 放大后再反向传播，防止梯度下溢。如果你用的是 10 系列显卡（Pascal 架构），这套 API 兼容性不佳，建议直接跑 FP32。通过这个改动，b4 模型也能在 8G 显存上塞进 batch_size 16，这在演示工程原始配置里是做不到的。

4. 训练过程中的验证、观察点与常见失败模式

4.1 每轮结束跑一次测试集，别只看训练 loss

演示工程在训练完成后统一测试，但这样遇到问题只能等全部跑完。建议在每个 epoch 结束后立即对测试集做一次前向推理，记录 top-1 准确率。代码可以复用测试逻辑，但要注意把模型切到 eval 模式，并关闭梯度计算。

def evaluate(model, test_loader):
    model.eval()
    correct = 0.0
    total = 0.0
    with torch.no_grad():
        for inputs, labels in test_loader:
            inputs, labels = inputs.cuda(), labels.cuda()
            outputs = model(inputs)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    return correct / total

torch.max 返回每一行最大值的值和索引， predicted 就是模型预测的类别 ID。 model.eval() 尤其重要，因为 batch norm 和 dropout 在训练和推理时的行为不同，忘记切模式会导致验证准确率虚高或忽高忽低。每一次验证消耗的时间只相当于一次前向，对于几百张测试集来说可以忽略不计。

4.2 Loss 不降的排查路径

最常见的情况是学习率过大，模型在最优解附近震荡。如果训练 loss 在某个值附近反复横跳，可以把 lr 降到当前值的 1/10，再把 momentum 从 0.9 提到 0.95。另一个高频错误是预训练模型加载失败，而服务器又恰好允许网络访问失败时静默跳过，导致模型从随机权重开始训练，收敛极慢。判断方法是在第一个 batch 前打印模型参数的 L2 范数，如果太小说明权重被错误初始化了。

还有一个隐蔽问题是类别不平衡。演示工程的数据集在 ImageFolder 下没有做任何采样控制，当某个类只有 10 张、另一个类有 1000 张时，模型会偏向大类别。需要用 WeightedRandomSampler 重采样。

from torch.utils.data import sampler, DataLoader

targets = train_set.targets
class_counts = torch.bincount(torch.tensor(targets))
weights = 1.0 / class_counts.float()
sample_weights = weights[targets]
train_sampler = sampler.WeightedRandomSampler(sample_weights, len(sample_weights), replacement=True)
train_loader = DataLoader(train_set, batch_size=batch_size, sampler=train_sampler)

这里 weights 是每个类的样本权重， WeightedRandomSampler 会根据它放大样本量少的类别的出现概率。 replacement=True 允许同一个样本被重复采样，让模型在小类别上见到更多次。这个改动不会改变数据集文件的物理分布，只是 DataLoader 的采样策略。

4.3 显存溢出时不是只有减小 batch_size 一条路

如果 batch_size 调到 8 还是 OOM，可以考虑把输入分辨率从 224 降到 192 或 160。EfficientNet 是复合缩放的产物，输入分辨率对显存的影响是平方级。另一个技巧是开启 torch.backends.cudnn.benchmark = True ，它能让 PyTorch 在卷积算法搜索时选择最快实现，减少显存碎片。还有一种情况是你的 GPU 卡在 CPU 到 GPU 拷贝上， DataLoader 的 num_workers 设成 4 或 8 能隐藏这部分延迟，但 Windows 下 num_workers 不要超过 0，否则多进程会产生奇怪占用。

4.4 训练和推理的语义差异决定了测试准确率低于训练

很多第一次跑这个 demo 的人会困惑：训练集准确率 99%，测试集只有 85%。原因不止过拟合，还有 RandomResizedCrop 和 RandomHorizontalFlip 只在训练 pipeline 里，测试集只有 Resize 和 CenterCrop 。这是图像分类的标准做法，但你要确保测试集预处理函数和训练集除了数据增强外完全一致，比如 ToTensor 之后的 Normalize 必须相同，否则输入分布直接被改变。

5. 把训练好的 EfficientNet 导出成 TorchScript 并做 CPU 推理

5.1 导出时要把预处理也包进去

训练完以后，模型还依赖 Python 端的预处理逻辑。部署到生产环境时，最好先用 TorchScript 把模型和预处理一起打包，这样 Java/Go 服务就不用复刻 PyTorch 的归一化细节。常见做法是先定义一个包装模块，把 Resize、Normalize 放到前处理里。

import torch
from torchvision import transforms
from efficientnet_pytorch import EfficientNet

class WrappedEfficientNet(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
    def forward(self, x):
        # x 是 PIL Image 或 numpy 数组
        x = self.transform(x)
        x = x.unsqueeze(0)
        return self.model(x)

model = EfficientNet.from_name('efficientnet-b0', num_classes=20)
state_dict = torch.load('best_model.pth', map_location='cpu')
model.load_state_dict(state_dict)
wrapped = WrappedEfficientNet(model)
wrapped.eval()
traced = torch.jit.script(wrapped)
traced.save('efficientnet_b0.pt')

torch.jit.script 会跟踪模块的前向过程，把 Python 代码固化成 TorchScript 图。注意这里用了 script 而不是 trace ，因为 transform 里有条件分支和复合操作， trace 会丢失控制流。如果你的部署环境没有 PIL，就把 Resize 改成 Tensor 上的插值，但这需要做一次 transform 的等价替换。

5.2 CPU 推理时的 batch 和线程设置

导出的模型可以直接用 PyTorch 的 C++ 库加载，但如果你还在 Python 环境，就需要注意线程数设置。推理时把 torch.set_num_threads(4) 设为物理核数的一半，往往比默认的自动检测更快。batch_size 在推理时没必要保持训练时的值，CPU 上一次只推一张通常延迟最低，因为你避免了等待最后一批补零。

import torch

torch.set_num_threads(4)
model = torch.jit.load('efficientnet_b0.pt')
model.eval()

# 对单张图片做预测
with torch.no_grad():
    output = model(img)
prob = torch.softmax(output, dim=1)
confidence, idx = torch.topk(prob, k=3)

torch.topk 返回概率最高的前三个类别和对应置信度。这样就能在日志里输出 "class: 5, conf: 0.92" 这样的结构化结果。如果你的服务器对时延更敏感，可以考虑把模型转成 ONNX，再用 ONNX Runtime 的 intra_op_num_threads 调优，这一步能再快 20% 到 30%。转 ONNX 时要把 wrapped 模块中的 Resize 替换成一个固定尺寸的插值 Tensor 操作，因为很多推理引擎不支持 PIL 类型的输入。

现在你已经有一个可以离线加载的 efficientnet_b0.pt ，写一个基础的 Flask 服务接收图片，调用这个模型返回 top-1 分类结果，大约也就 50 行代码。整个流程从数据集组织、训练参数调整到部署落地的路径都清晰了，下一步要做的，就是用你自己的数据替换 demo 里的 dataset 目录。

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