3. 基于迁移学习的作物叶片病害分类模型训练

任务定位与迁移学习策略

本案例在已清洗、类别平衡的作物叶片病害数据集基础上，构建一个可区分多种病害类别的分类模型。核心做法是：以 ImageNet 预训练的卷积网络为骨干，冻结特征提取层先训练分类头，再解冻顶层做小学习率微调，配合面向田间的数据增强，最后用验证集评估并输出模型。相比从零训练，迁移学习能在样本量有限的农业数据集上更快收敛并降低过拟合风险。

下面给出一个自洽、可运行的 PyTorch 实现。数据目录采用 ImageFolder 标准结构（每个病害类别一个子文件夹），示例类别为 健康/早疫病/晚疫病/叶斑病，实际使用时替换为你的类别名即可。

关键实现代码

代码包含数据增强、两阶段训练、验证评估与模型保存，可直接运行。示例使用 CPU/GPU 自适应，类别数由数据目录自动推断。

import copy
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms, models

DATA_DIR = "leaf_dataset"   # 结构: leaf_dataset/train/类别名/*.jpg, leaf_dataset/val/类别名/*.jpg
BATCH = 32
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 1. 数据增强：训练集强增强，验证集仅缩放+归一化
train_tf = transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.7, 1.0)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomVerticalFlip(),
    transforms.RandomRotation(20),
    transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
val_tf = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

train_ds = datasets.ImageFolder(f"{DATA_DIR}/train", transform=train_tf)
val_ds = datasets.ImageFolder(f"{DATA_DIR}/val", transform=val_tf)
train_loader = DataLoader(train_ds, batch_size=BATCH, shuffle=True, num_workers=2)
val_loader = DataLoader(val_ds, batch_size=BATCH, shuffle=False, num_workers=2)
num_classes = len(train_ds.classes)
print("类别:", train_ds.classes, "| 训练样本:", len(train_ds), "| 验证样本:", len(val_ds))

# 2. 构建迁移学习模型
model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
for p in model.parameters():
    p.requires_grad = False                      # 先冻结全部骨干
model.fc = nn.Sequential(nn.Dropout(0.3), nn.Linear(model.fc.in_features, num_classes))
model = model.to(DEVICE)

criterion = nn.CrossEntropyLoss()

def run_epoch(loader, train=True):
    model.train() if train else model.eval()
    total_loss, correct, total = 0.0, 0, 0
    with torch.set_grad_enabled(train):
        for x, y in loader:
            x, y = x.to(DEVICE), y.to(DEVICE)
            out = model(x)
            loss = criterion(out, y)
            if train:
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
            total_loss += loss.item() * x.size(0)
            correct += (out.argmax(1) == y).sum().item()
            total += x.size(0)
    return total_loss / total, correct / total

# 3. 阶段一：仅训练分类头
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3)
best_acc, best_state = 0.0, None
for epoch in range(5):
    tr_loss, tr_acc = run_epoch(train_loader, True)
    va_loss, va_acc = run_epoch(val_loader, False)
    print(f"[阶段一] epoch {epoch+1}: train_acc={tr_acc:.3f} val_acc={va_acc:.3f}")
    if va_acc > best_acc:
        best_acc, best_state = va_acc, copy.deepcopy(model.state_dict())

# 4. 阶段二：解冻 layer4 与分类头，小学习率微调
for p in model.layer4.parameters():
    p.requires_grad = True
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-4)
for epoch in range(10):
    tr_loss, tr_acc = run_epoch(train_loader, True)
    va_loss, va_acc = run_epoch(val_loader, False)
    print(f"[阶段二] epoch {epoch+1}: train_acc={tr_acc:.3f} val_acc={va_acc:.3f}")
    if va_acc > best_acc:
        best_acc, best_state = va_acc, copy.deepcopy(model.state_dict())

# 5. 保存最优模型与类别映射
model.load_state_dict(best_state)
torch.save({"state_dict": best_state,
            "class_to_idx": train_ds.class_to_idx,
            "num_classes": num_classes}, "best_leaf_model.pth")
print(f"最优验证准确率: {best_acc:.3f}，模型已保存")

参数选择与结果判读

迁移学习的效果高度依赖几个关键设置，下表给出本案例的推荐取值与调整依据：