3行代码搞定自定义训练：EfficientNet-PyTorch从CSV到评估全流程-优快云博客

3行代码搞定自定义训练：EfficientNet-PyTorch从CSV到评估全流程

【免费下载链接】EfficientNet-PyTorch A PyTorch implementation of EfficientNet and EfficientNetV2 (coming soon!) 项目地址: https://gitcode.com/gh_mirrors/ef/EfficientNet-PyTorch

你还在为自定义数据集训练 EfficientNet 模型烦恼吗？标注好的 CSV 文件不知如何导入？训练过程中参数调优毫无头绪？本文将带你用极简代码实现从 CSV 数据加载到模型评估的全流程，即使是机器学习新手也能快速上手。读完本文你将掌握：数据格式转换技巧、迁移学习参数设置、训练可视化方法以及模型性能评估完整流程。

项目准备与环境配置

首先通过以下命令克隆项目仓库并安装依赖：

git clone https://gitcode.com/gh_mirrors/ef/EfficientNet-PyTorch
cd EfficientNet-PyTorch
pip install -r requirements.txt

项目核心模块包括：

数据集准备与格式转换

CSV 文件规范

自定义数据集需遵循以下 CSV 格式（以花卉分类为例）：

image_path,label
data/roses/1.jpg,0
data/tulips/2.jpg,1
data/daisy/3.jpg,2

其中 image_path 为图片相对路径，label 为整数类别标签。项目示例中的标签映射文件可参考 examples/simple/labels_map.txt。

数据加载器实现

创建 custom_dataset.py 文件，实现 CSV 数据加载功能：

import pandas as pd
from PIL import Image
from torch.utils.data import Dataset

class CSVDataset(Dataset):
    def __init__(self, csv_path, transform=None):
        self.data = pd.read_csv(csv_path)
        self.transform = transform
        
    def __len__(self):
        return len(self.data)
        
    def __getitem__(self, idx):
        img_path = self.data.iloc[idx]['image_path']
        image = Image.open(img_path).convert('RGB')
        label = int(self.data.iloc[idx]['label'])
        
        if self.transform:
            image = self.transform(image)
            
        return image, label

模型加载与迁移学习设置

加载预训练模型

使用一行代码加载 EfficientNet-B0 预训练模型：

from efficientnet_pytorch import EfficientNet
model = EfficientNet.from_pretrained('efficientnet-b0', num_classes=3)

其中 num_classes 参数需设置为你的分类类别数。模型结构定义详见 efficientnet_pytorch/model.py 中的 EfficientNet 类。

冻结与微调策略

采用分层学习率策略进行迁移学习：

# 冻结特征提取层
for param in model._conv_stem.parameters():
    param.requires_grad = False
    
# 仅训练分类头
optimizer = torch.optim.Adam(model._fc.parameters(), lr=1e-4)

完整训练配置可参考 examples/imagenet/main.py 中的参数设置。

训练流程实现

数据预处理管道

from torchvision import transforms

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])
])

训练循环代码

from torch.utils.data import DataLoader

train_dataset = CSVDataset('train.csv', transform=transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

for epoch in range(10):
    model.train()
    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
    print(f'Epoch {epoch}, Loss: {loss.item()}')

模型评估与可视化

评估指标计算

model.eval()
correct = 0
total = 0
with torch.no_grad():
    for images, labels in test_loader:
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accuracy: {100 * correct / total}%')

混淆矩阵可视化

import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix
import seaborn as sns

# 生成混淆矩阵代码省略
plt.figure(figsize=(10,8))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
plt.savefig('confusion_matrix.png')

训练完成后可参考 examples/simple/check.ipynb 进行模型验证，示例图片如下：

常见问题解决方案

内存溢出：降低 examples/imagenet/main.py 中的 batch_size 参数
过拟合：增加 dropout 率，修改 efficientnet_pytorch/model.py 中的 _dropout 层
收敛缓慢：调整学习率，参考 adjust_learning_rate 函数实现

总结与扩展

本文展示了如何用最少代码实现基于 CSV 数据集的 EfficientNet 训练流程。通过合理的迁移学习策略，即使在小数据集上也能获得良好性能。进阶用户可尝试：

模型压缩：修改模型宽度/深度系数
数据增强：扩展 transforms 实现更多 augmentation
多GPU训练：参考 examples/imagenet/main.py 中的分布式训练代码

收藏本文，下次训练 EfficientNet 无需从头开始！如有疑问，可查阅项目官方文档 README.md 或提交 issue。

【免费下载链接】EfficientNet-PyTorch A PyTorch implementation of EfficientNet and EfficientNetV2 (coming soon!) 项目地址: https://gitcode.com/gh_mirrors/ef/EfficientNet-PyTorch

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考