4. 模型训练实战

本节把前面准备好的数据真正“喂”进网络，完成一个可复现的训练闭环：定义作物病害分类模型、编写训练循环、用学习率调度控制优化节奏、按验证集指标保存最优权重，并在训练结束后用验证集评估模型。示例以“番茄叶片三分类（健康、早疫病、晚疫病）”为假设场景，输入为 224×224 的 RGB 叶片图像，训练集约 4800 张、验证集约 1200 张，类别分布大致均衡。这些数字属于教学示例，实际项目应按自己的数据集替换。

一、模型结构：从骨干网络到分类头

农业病害图像的特点是病斑区域小、纹理细、背景（土壤、支架、光照）干扰大，因此常用迁移学习：加载在 ImageNet 上预训练的 ResNet18 或 MobileNetV3，替换最后的全连接层为与类别数一致的输出。下面用 ResNet18 构建分类器，并冻结骨干的前若干层以减少小数据集上的过拟合。

import torch
import torch.nn as nn
from torchvision import models

def build_model(num_classes=3, freeze_backbone=True):
    model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
    if freeze_backbone:
        for name, p in model.named_parameters():
            if not name.startswith("fc"):
                p.requires_grad = False
    in_features = model.fc.in_features
    model.fc = nn.Sequential(
        nn.Dropout(0.3),
        nn.Linear(in_features, num_classes)
    )
    return model

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = build_model(num_classes=3).to(device)

冻结骨干后，只有分类头参与反向传播，第一轮训练相当于“线性探针”，收敛快且不易破坏预训练特征。待验证集准确率稳定后，可解冻最后一个 stage 做小学习率微调，这是农业小样本场景中较稳妥的两段式策略。

二、训练循环：前向、损失、反向、更新

训练循环的核心是四个动作的顺序与作用域：optimizer.zero_grad() 清空上一轮梯度，前向得到 logits，交叉熵计算损失，loss.backward() 累积梯度，optimizer.step() 更新参数。同时必须显式切换 model.train() 与 model.eval()，因为 Dropout 和 BatchNorm 在两种模式下的行为不同，漏掉 eval() 会让验证指标虚高或抖动。

import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR

criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()),
                        lr=3e-4, weight_decay=1e-4)
scheduler = CosineAnnealingLR(optimizer, T_max=15, eta_min=1e-6)

def train_one_epoch(model, loader):
    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()
        logits = model(images)
        loss = criterion(logits, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * labels.size(0)
        correct += (logits.argmax(1) == labels).sum().item()
        total += labels.size(0)
    return total_loss / total, correct / total

这里用 loss.item() * labels.size(0) 累加再除以样本总数，得到的是按样本加权的平均损失，比逐 batch 求平均再平均更准确，尤其在最后一个 batch 不满时差别明显。准确率同理按样本累计。

三、学习率调度：让优化节奏随训练阶段变化

学习率是训练中最敏感的超参数。固定学习率在初期可能震荡，在后期又难以精细收敛。余弦退火让学习率从初始值平滑下降到接近零，适合总轮数已知的短周期训练；若训练轮数不确定，可用 ReduceLROnPlateau，在验证损失连续若干轮不下降时衰减。注意调度器的调用位置：按 epoch 更新的调度器应在每个 epoch 结束后调用 scheduler.step()，而不是每个 batch。

def evaluate(model, loader):
    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)
            logits = model(images)
            loss = criterion(logits, labels)
            total_loss += loss.item() * labels.size(0)
            correct += (logits.argmax(1) == labels).sum().item()
            total += labels.size(0)
    return total_loss / total, correct / total

验证阶段必须包在 torch.no_grad() 中，既关闭梯度计算节省显存，也避免误把验证数据用于参数更新。评估函数与训练函数结构对称，便于对比两者的损失与准确率。

四、保存最优权重与验证集评估

训练过程中不能只看最后一轮的权重，而应以验证集准确率为准保存“最优检查点”。下面的主循环把训练损失、验证损失、验证准确率和当前学习率记录下来，并在验证准确率提升时覆盖保存 best_crop_disease.pth。保存内容建议包含 state_dict、类别名映射和输入尺寸，便于部署时对齐预处理。

best_acc = 0.0
history = []
for epoch in range(15):
    tr_loss, tr_acc = train_one_epoch(model, train_loader)
    val_loss, val_acc = evaluate(model, val_loader)
    lr_now = optimizer.param_groups[0]["lr"]
    scheduler.step()
    history.append((epoch + 1, tr_loss, val_loss, val_acc, lr_now))
    print(f"Epoch {epoch+1:02d} | train_loss={tr_loss:.4f} "
          f"train_acc={tr_acc:.3f} | val_loss={val_loss:.4f} "
          f"val_acc={val_acc:.3f} | lr={lr_now:.2e}")
    if val_acc > best_acc:
        best_acc = val_acc
        torch.save({
            "state_dict": model.state_dict(),
            "classes": ["健康", "早疫病", "晚疫病"],
            "input_size": 224,
            "val_acc": best_acc
        }, "best_crop_disease.pth")
print("最优验证准确率:", round(best_acc, 4))

如何判断训练是否健康？可观察三条曲线：训练损失持续下降而验证损失开始上升，说明过拟合，应加强数据增强、提高 Dropout 或提前停止；训练与验证损失都居高不下，说明欠拟合或学习率过大，可解冻更多层或降低初始学习率；验证准确率在若干轮内反复横跳，往往是学习率偏大或 batch 过小。示例中若第 8 轮验证准确率达到 0.94 后不再提升，而训练准确率已到 0.99，就应保留第 8 轮权重作为最终模型。

训练闭环：数据 → 前向 → 损失 → 反向 → 更新 → 验证 → 保存 训练批次 224×224 叶片图像 ResNet18 前向 冻结骨干 + 分类头 交叉熵损失 logits vs 标签 反向传播 backward 累积梯度 AdamW 更新 lr=3e-4 余弦退火调度 每 epoch step 验证集评估 no_grad + eval() 最优权重保存 val_acc 提升时覆盖 下一 epoch 结果判读要点 训练损失↓ 验证损失↑ → 过拟合：加强增强 / 提高 Dropout / 提前停止 两者都高 → 欠拟合：解冻更多层 / 降低初始学习率 验证准确率反复横跳 → 学习率偏大或 batch 过小

训练结束后，用保存的最优权重重新加载并做一次完整验证，确认指标与训练日志一致，同时输出混淆矩阵，重点看“早疫病”与“晚疫病”是否互相误判——这两类病斑形态接近，是农业分类任务中最常见的混淆对。若某一类召回率明显偏低，应回到数据层检查该类样本数量与标注质量，而不是盲目增加训练轮数。至此，一个可复现、可评估、可保存的作物病害分类器训练流程就完成了，导出的权重文件可直接用于后续的 ONNX 转换与部署。

← 上一章 📖 返回目录 下一章 →