PyTorch实战：EfficientNet自定义数据集训练全流程

1. 项目概述

在计算机视觉领域，EfficientNet系列模型因其卓越的性能和效率平衡而广受欢迎。但很多开发者在使用自定义数据集训练时常常遇到各种问题。本文将分享我使用PyTorch框架训练EfficientNet模型处理自定义数据集的完整流程，包含从数据准备到模型部署的全套解决方案。

2. 核心需求解析

2.1 为什么选择EfficientNet

EfficientNet通过复合缩放方法统一调整网络宽度、深度和分辨率，在ImageNet上达到84.4% top-1准确率的同时，参数数量比ResNet-50少8.4倍。对于资源受限的应用场景，B0-B7不同规模的变体可以灵活选择。

2.2 自定义数据集的挑战

与标准数据集不同，自定义数据集通常面临：

• 样本量不足（可能只有几百张）
• 类别不平衡
• 标注质量参差不齐
• 图像尺寸和比例不统一

3. 环境准备与数据预处理

3.1 开发环境配置

推荐使用Python 3.8+和PyTorch 1.10+环境：

bash复制conda create -n efficientnet python=3.8
conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch
pip install efficientnet_pytorch

3.2 数据组织规范

采用标准ImageFolder目录结构：

code复制dataset/
    train/
        class1/
            img1.jpg
            img2.jpg
        class2/
            img1.jpg
    val/
        ...

3.3 数据增强策略

针对小样本数据集的关键增强技巧：

python复制from torchvision import transforms

train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

4. 模型训练实战

4.1 模型初始化

使用预训练权重进行迁移学习：

python复制from efficientnet_pytorch import EfficientNet

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

4.2 训练超参数设置

经验验证的推荐配置：

python复制optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-5)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)
criterion = torch.nn.CrossEntropyLoss(label_smoothing=0.1)

4.3 训练循环实现

关键训练代码结构：

python复制for epoch in range(epochs):
    model.train()
    for inputs, labels in train_loader:
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    
    scheduler.step()

5. 模型优化技巧

5.1 混合精度训练

显著减少显存占用：

python复制scaler = torch.cuda.amp.GradScaler()

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

5.2 类别不平衡处理

采用加权采样策略：

python复制class_weights = 1. / torch.tensor(class_counts)
samples_weights = class_weights[labels]
sampler = WeightedRandomSampler(samples_weights, len(samples_weights))

6. 模型评估与部署

6.1 评估指标设计

除准确率外应关注：

• 混淆矩阵
• 各类别Precision/Recall
• F1 Score

6.2 模型导出方案

导出为ONNX格式：

python复制dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(model, dummy_input, "model.onnx")

7. 常见问题解决

7.1 过拟合应对策略

• 增加Dropout层（p=0.2-0.5）
• 使用Early Stopping
• 尝试更强的数据增强

7.2 显存不足解决方案

• 减小batch size（最低可到8）
• 使用梯度累积
• 尝试模型蒸馏（如B0→B1）

8. 进阶优化方向

8.1 知识蒸馏应用

使用大模型指导小模型训练：

python复制teacher_model = EfficientNet.from_pretrained('efficientnet-b7')
student_model = EfficientNet.from_name('efficientnet-b0')

# 蒸馏损失计算
loss = 0.7*KLDivLoss(student_logits, teacher_logits) + 0.3*CrossEntropyLoss(student_logits, labels)

8.2 模型量化部署

动态量化实现：

python复制quantized_model = torch.quantization.quantize_dynamic(
    model, {torch.nn.Linear}, dtype=torch.qint8
)