12.5 训练循环实现 - AI应用开发

🎯 学习目标

掌握完整的模型训练流程
学会使用损失函数和优化器
了解模型保存与加载方法
掌握训练过程的监控与调试

训练循环概述

训练循环是深度学习的核心流程，包括前向传播、计算损失、反向传播和参数更新四个步骤。一个完整的训练流程还包括验证评估、模型保存、早停策略等。

🔄 训练流程四步骤

前向传播

计算损失

反向传播

参数更新

📊 损失函数选择

损失函数	适用场景	PyTorch实现
交叉熵损失	多分类问题	nn.CrossEntropyLoss()
二元交叉熵	二分类问题	nn.BCELoss() / nn.BCEWithLogitsLoss()
均方误差	回归问题	nn.MSELoss()
平均绝对误差	回归问题（鲁棒）	nn.L1Loss()
负对数似然	多分类（配合LogSoftmax）	nn.NLLLoss()

⚙️ 优化器选择

import torch.optim as optim

# SGD（随机梯度下降）
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# Adam（自适应学习率）
optimizer = optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.999))

# AdamW（带权重衰减的Adam）
optimizer = optim.AdamW(model.parameters(), lr=0.001, weight_decay=0.01)

# RMSprop
optimizer = optim.RMSprop(model.parameters(), lr=0.01)

# 学习率调度器
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
      

💻 完整训练循环

import torch
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm

def train_model(model, train_loader, val_loader, epochs, device):
    """完整训练函数"""

    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)

    model = model.to(device)
    best_val_acc = 0.0

    for epoch in range(epochs):
        # ========== 训练阶段 ==========
        model.train()
        train_loss = 0.0
        train_correct = 0
        train_total = 0

        for batch_idx, (data, targets) in enumerate(tqdm(train_loader, desc=f'Epoch {epoch+1}')):
            data, targets = data.to(device), targets.to(device)

            # 1. 清零梯度
            optimizer.zero_grad()

            # 2. 前向传播
            outputs = model(data)
            loss = criterion(outputs, targets)

            # 3. 反向传播
            loss.backward()

            # 4. 参数更新
            optimizer.step()

            # 统计
            train_loss += loss.item()
            _, predicted = outputs.max(1)
            train_total += targets.size(0)
            train_correct += predicted.eq(targets).sum().item()

        scheduler.step()

        # ========== 验证阶段 ==========
        model.eval()
        val_loss = 0.0
        val_correct = 0
        val_total = 0

        with torch.no_grad():
            for data, targets in val_loader:
                data, targets = data.to(device), targets.to(device)
                outputs = model(data)
                loss = criterion(outputs, targets)

                val_loss += loss.item()
                _, predicted = outputs.max(1)
                val_total += targets.size(0)
                val_correct += predicted.eq(targets).sum().item()

        # 打印结果
        print(f'Epoch [{epoch+1}/{epochs}]')
        print(f'Train Loss: {train_loss/len(train_loader):.4f}, '
              f'Train Acc: {100.*train_correct/train_total:.2f}%')
        print(f'Val Loss: {val_loss/len(val_loader):.4f}, '
              f'Val Acc: {100.*val_correct/val_total:.2f}%')

        # 保存最佳模型
        val_acc = val_correct / val_total
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save(model.state_dict(), 'best_model.pth')

    print(f'Best Validation Accuracy: {best_val_acc*100:.2f}%')
    return model
      

💾 模型保存与加载

# 保存模型参数（推荐）
torch.save(model.state_dict(), 'model_weights.pth')

# 加载模型参数
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()  # 切换到评估模式

# 保存完整模型（包含结构）
torch.save(model, 'complete_model.pth')

# 加载完整模型
model = torch.load('complete_model.pth')

# 保存训练检查点（断点续训）
checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
}
torch.save(checkpoint, 'checkpoint.pth')

# 加载检查点
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
start_epoch = checkpoint['epoch']
      

⚖️ train()与eval()模式区别

模式	Dropout	BatchNorm	梯度计算
train()	启用（随机丢弃）	使用批次统计量更新	启用
eval()	禁用（保留所有神经元）	使用固定的均值和方差	可启用（配合no_grad禁用）

💡

最佳实践

训练时务必调用model.train()，验证/测试时调用model.eval()。这会影响Dropout和BatchNorm的行为。推理时建议同时使用with torch.no_grad()来节省内存。

图：训练循环是深度学习的核心流程

📝 本节小结

✅

• 训练循环包括：前向传播、计算损失、反向传播、参数更新
• 选择合适的损失函数：分类用CrossEntropy，回归用MSE
• Adam是最常用的优化器，适合大多数场景
• model.train()和model.eval()切换训练/评估模式
• 使用torch.save/load进行模型保存和加载