完整训练流程#

优化器:用梯度更新参数中我们掌握了如何更新参数。现在的问题是:如何把所有组件整合成一个完整的训练系统?

神经网络训练基础中,我们学习了训练的核心概念——Epoch、Batch、过拟合与欠拟合、正则化技巧等。但这些知识零散,需要串联成可运行的代码。

本节就是答案:我们将构建一个完整的 MNIST 训练系统,涵盖数据加载、模型训练、验证评估、模型保存等全部环节。

训练流程概览#

flowchart LR A[数据准备] --> B[模型定义] B --> C[训练循环] C --> D[验证评估] D --> E[模型保存] E --> F{继续训练?} F -->|是| C F -->|否| G[结束]

核心流程(对应 神经网络训练基础):

  1. 数据准备:加载数据,划分训练/验证集,创建 DataLoader

  2. 模型定义:搭建网络架构

  3. 训练循环:前向传播 → 计算损失 → 反向传播 → 更新参数

  4. 验证评估:评估模型性能,检测过拟合

  5. 模型保存:保存最佳模型

数据准备#

数据集与 DataLoader#

神经网络训练基础中提到:Batch Size 是每次参数更新使用的样本数。PyTorch 用 DataLoader 实现这个机制。

import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

# 数据预处理
# transforms.Compose 将多个预处理步骤串联
# MNIST 数据原始范围 [0, 255],需要归一化到 [0, 1]
transform = transforms.Compose([
    transforms.ToTensor(),           # PIL Image → Tensor,自动除以 255
    transforms.Normalize((0.1307,), (0.3081,))  # 标准化:(x - mean) / std
])

# 加载 MNIST 数据集
# train=True 表示训练集,train=False 表示测试集
train_dataset = torchvision.datasets.MNIST(
    root='./data',      # 数据存储路径
    train=True,         # 训练集
    download=True,      # 自动下载
    transform=transform # 应用预处理
)

test_dataset = torchvision.datasets.MNIST(
    root='./data',
    train=False,        # 测试集
    download=True,
    transform=transform
)

# 创建 DataLoader
# batch_size:每批样本数(对应 neural-training-basics 中的 Batch Size)
# shuffle=True:每 epoch 打乱数据顺序,防止模型记住顺序
# num_workers:多进程加载数据,加速训练
train_loader = DataLoader(
    train_dataset, 
    batch_size=64,      # MNIST 有 60,000 样本,每 epoch 约 938 个 batch
    shuffle=True,       # 打乱顺序
    num_workers=2       # 2 个子进程加载数据
)

test_loader = DataLoader(
    test_dataset,
    batch_size=64,
    shuffle=False,      # 测试集不需要打乱
    num_workers=2
)

print(f"训练集大小: {len(train_dataset)}")  # 60,000
print(f"测试集大小: {len(test_dataset)}")   # 10,000
print(f"每 epoch 迭代次数: {len(train_loader)}")  # 938

数据增强(可选)#

神经网络训练基础中的数据增强可以增加训练样本多样性:

# 训练时增强,测试时不增强
train_transform = transforms.Compose([
    transforms.RandomRotation(10),      # 随机旋转 ±10 度
    transforms.RandomAffine(degrees=0, translate=(0.1, 0.1)),  # 随机平移
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

test_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

模型定义#

CNN 架构实现#

对应 卷积神经网络中的卷积网络设计:

import torch.nn as nn
import torch.nn.functional as F

class MNISTNet(nn.Module):
    """
    MNIST 分类网络
    输入: [batch, 1, 28, 28]
    输出: [batch, 10] (10 个数字类别的 logits)
    """
    def __init__(self):
        super(MNISTNet, self).__init__()
        
        # C1: 卷积层,1→32 通道,3×3 卷积核,输出 28×28
        # 参数量: 32×1×3×3 + 32 = 320
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        
        # C2: 卷积层,32→64 通道,3×3 卷积核,输出 28×28
        # 参数量: 64×32×3×3 + 64 = 18,496
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        
        # 池化层: 2×2,输出 14×14 → 7×7
        self.pool = nn.MaxPool2d(2)
        
        # 全连接层1: 64×7×7=3136 → 128
        # 参数量: 3136×128 + 128 = 401,536
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        
        # 全连接层2: 128 → 10(类别数)
        # 参数量: 128×10 + 10 = 1,290
        self.fc2 = nn.Linear(128, 10)
        
        # Dropout: 防止过拟合(见 neural-training-basics)
        self.dropout = nn.Dropout(0.25)  # 25% 的 dropout
        
    def forward(self, x):
        # 卷积块1: [batch, 1, 28, 28] → [batch, 32, 28, 28] → [batch, 32, 14, 14]
        x = self.conv1(x)
        x = F.relu(x)           # 激活函数
        x = self.pool(x)        # 降采样
        
        # 卷积块2: [batch, 32, 14, 14] → [batch, 64, 14, 14] → [batch, 64, 7, 7]
        x = self.conv2(x)
        x = F.relu(x)
        x = self.pool(x)
        
        # 展平: [batch, 64, 7, 7] → [batch, 3136]
        x = x.view(x.size(0), -1)
        
        # 全连接 + Dropout
        x = self.dropout(x)
        x = self.fc1(x)
        x = F.relu(x)
        
        x = self.dropout(x)
        x = self.fc2(x)  # 输出 logits(未归一化)
        
        return x

总参数量:320 + 18,496 + 401,536 + 1,290 = 421,642

训练函数#

单 epoch 训练#

对应 神经网络训练基础中的训练循环:

def train_epoch(model, device, train_loader, optimizer, criterion, epoch):
    """
    训练一个 epoch
    
    Args:
        model: 神经网络模型
        device: 计算设备(CPU/GPU)
        train_loader: 训练数据加载器
        optimizer: 优化器
        criterion: 损失函数
        epoch: 当前 epoch 数
    
    Returns:
        avg_loss: 平均损失
        accuracy: 准确率
    """
    model.train()  # 设置为训练模式(启用 Dropout、BatchNorm 等)
    
    train_loss = 0.0
    correct = 0
    total = 0
    
    for batch_idx, (data, target) in enumerate(train_loader):
        # 1. 数据转移到设备
        data, target = data.to(device), target.to(device)
        
        # 2. 清零梯度(关键!见 auto-grad)
        optimizer.zero_grad()
        
        # 3. 前向传播
        output = model(data)  # [batch, 10]
        
        # 4. 计算损失
        # CrossEntropyLoss 内部已经包含 Softmax
        loss = criterion(output, target)
        
        # 5. 反向传播
        loss.backward()
        
        # 6. 更新参数
        optimizer.step()
        
        # 7. 统计信息
        train_loss += loss.item()
        _, predicted = output.max(1)  # 取最大值的索引作为预测类别
        total += target.size(0)
        correct += predicted.eq(target).sum().item()
        
        # 8. 进度显示
        if batch_idx % 100 == 0:
            print(f'Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} '
                  f'({100. * batch_idx / len(train_loader):.0f}%)]\t'
                  f'Loss: {loss.item():.6f}')
    
    # 计算平均指标
    avg_loss = train_loss / len(train_loader)
    accuracy = 100. * correct / total
    
    return avg_loss, accuracy

训练模式的重要性

  • model.train():启用 Dropout(随机丢弃神经元)、BatchNorm 使用 batch 统计量

  • model.eval():关闭 Dropout,BatchNorm 使用全局统计量

验证函数#

评估模型性能#

def validate(model, device, val_loader, criterion):
    """
    验证模型性能
    
    Args:
        model: 神经网络模型
        device: 计算设备
        val_loader: 验证数据加载器
        criterion: 损失函数
    
    Returns:
        avg_loss: 平均验证损失
        accuracy: 验证准确率
    """
    model.eval()  # 设置为评估模式
    
    val_loss = 0.0
    correct = 0
    total = 0
    
    # 禁用梯度计算(节省内存,加速推理)
    with torch.no_grad():
        for data, target in val_loader:
            data, target = data.to(device), target.to(device)
            
            # 前向传播
            output = model(data)
            loss = criterion(output, target)
            
            # 统计
            val_loss += loss.item()
            _, predicted = output.max(1)
            total += target.size(0)
            correct += predicted.eq(target).sum().item()
    
    avg_loss = val_loss / len(val_loader)
    accuracy = 100. * correct / total
    
    return avg_loss, accuracy

主训练循环#

完整训练流程#

import torch
import torch.optim as optim
import matplotlib.pyplot as plt

def main():
    # ========== 1. 设备设置 ==========
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"使用设备: {device}")
    
    # ========== 2. 数据准备 ==========
    # ...(前面的数据加载代码)...
    
    # ========== 3. 模型、损失、优化器 ==========
    model = MNISTNet().to(device)
    
    # 损失函数:交叉熵(见 neural-training-basics)
    criterion = nn.CrossEntropyLoss()
    
    # 优化器:Adam(常用选择)
    # weight_decay 对应 neural-training-basics 中的 L2 正则化
    optimizer = optim.Adam(
        model.parameters(), 
        lr=0.001,
        weight_decay=1e-4  # L2 正则化系数 λ
    )
    
    # 学习率调度器:每 5 个 epoch 学习率乘以 0.1
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
    
    # ========== 4. 训练历史记录 ==========
    # 用于检测过拟合(见 neural-training-basics)
    history = {
        'train_loss': [], 'train_acc': [],
        'val_loss': [], 'val_acc': []
    }
    
    best_acc = 0.0
    
    # ========== 5. 训练循环 ==========
    num_epochs = 10
    
    for epoch in range(1, num_epochs + 1):
        print(f"\n{'='*50}")
        print(f"Epoch {epoch}/{num_epochs}")
        print(f"{'='*50}")
        
        # 训练
        train_loss, train_acc = train_epoch(
            model, device, train_loader, optimizer, criterion, epoch
        )
        
        # 验证
        val_loss, val_acc = validate(model, device, test_loader, criterion)
        
        # 更新学习率
        scheduler.step()
        current_lr = optimizer.param_groups[0]['lr']
        
        # 记录历史
        history['train_loss'].append(train_loss)
        history['train_acc'].append(train_acc)
        history['val_loss'].append(val_loss)
        history['val_acc'].append(val_acc)
        
        # 打印结果
        print(f'\n训练集 - Loss: {train_loss:.4f}, Acc: {train_acc:.2f}%')
        print(f'验证集 - Loss: {val_loss:.4f}, Acc: {val_acc:.2f}%')
        print(f'学习率: {current_lr:.6f}')
        
        # 保存最佳模型
        if val_acc > best_acc:
            best_acc = val_acc
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'best_acc': best_acc,
            }, 'best_model.pth')
            print(f'保存最佳模型,准确率: {best_acc:.2f}%')
    
    return history

可视化训练过程#

绘制损失和准确率曲线#

对应 神经网络训练基础中的过拟合检测:

def plot_history(history):
    """
    可视化训练历史
    用于检测过拟合:训练损失↓但验证损失↑时,说明过拟合
    """
    fig, axes = plt.subplots(1, 2, figsize=(12, 4))
    
    # 损失曲线
    axes[0].plot(history['train_loss'], label='Train Loss')
    axes[0].plot(history['val_loss'], label='Val Loss')
    axes[0].set_xlabel('Epoch')
    axes[0].set_ylabel('Loss')
    axes[0].set_title('Loss Curve')
    axes[0].legend()
    axes[0].grid(True)
    
    # 准确率曲线
    axes[1].plot(history['train_acc'], label='Train Acc')
    axes[1].plot(history['val_acc'], label='Val Acc')
    axes[1].set_xlabel('Epoch')
    axes[1].set_ylabel('Accuracy (%)')
    axes[1].set_title('Accuracy Curve')
    axes[1].legend()
    axes[1].grid(True)
    
    plt.tight_layout()
    plt.savefig('training_history.png')
    plt.show()

# 使用
# history = main()
# plot_history(history)

曲线解读

  • 理想情况:两条曲线同步下降/上升,差距小

  • 过拟合迹象:训练准确率持续上升,验证准确率停滞或下降

  • 欠拟合迹象:两条曲线都停滞在较低水平

模型加载与推理#

加载保存的模型#

def load_and_predict(model_path, image):
    """
    加载模型并进行预测
    
    Args:
        model_path: 模型文件路径
        image: 输入图像 tensor,形状 [1, 1, 28, 28]
    
    Returns:
        prediction: 预测的类别(0-9)
        probability: 预测概率
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # 创建模型
    model = MNISTNet().to(device)
    
    # 加载权重
    checkpoint = torch.load(model_path)
    model.load_state_dict(checkpoint['model_state_dict'])
    
    # 推理
    model.eval()
    with torch.no_grad():
        image = image.to(device)
        output = model(image)
        
        # Softmax 获取概率
        probabilities = F.softmax(output, dim=1)
        
        # 预测结果
        prediction = output.argmax(dim=1).item()
        confidence = probabilities.max().item()
    
    return prediction, confidence

# 使用示例
# image = test_dataset[0][0].unsqueeze(0)  # 取第一张测试图片,加 batch 维度
# pred, conf = load_and_predict('best_model.pth', image)
# print(f"预测结果: {pred}, 置信度: {conf:.2%}")

总结#

训练流程回顾#

步骤

关键代码

对应理论

数据加载

DataLoader

神经网络训练基础 中的 Batch

模型定义

nn.Module

卷积神经网络

损失函数

CrossEntropyLoss

损失函数

优化器

optim.Adam

梯度下降与优化算法

反向传播

loss.backward()

反向传播算法

正则化

weight_decay, Dropout

神经网络训练基础

验证评估

model.eval()

检测过拟合

学习率调度

lr_scheduler

训练后期微调

完整训练脚本模板#

# main.py - 可运行的完整训练脚本
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

def main():
    # 配置
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    batch_size = 64
    epochs = 10
    lr = 0.001
    
    # 数据
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    train_loader = torch.utils.data.DataLoader(
        datasets.MNIST('./data', train=True, download=True, transform=transform),
        batch_size=batch_size, shuffle=True
    )
    test_loader = torch.utils.data.DataLoader(
        datasets.MNIST('./data', train=False, transform=transform),
        batch_size=batch_size
    )
    
    # 模型
    model = MNISTNet().to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=1e-4)
    
    # 训练
    for epoch in range(1, epochs + 1):
        train_epoch(model, device, train_loader, optimizer, criterion, epoch)
        validate(model, device, test_loader, criterion)
    
    # 保存
    torch.save(model.state_dict(), 'mnist_model.pth')

if __name__ == '__main__':
    main()

从手写到工程框架#

上面这 200 行代码包含了训练 MNIST 分类器的全部要素。但在真实项目中,每次从头写训练循环、手动管理 checkpoint、手工对比实验是不现实的。

社团的 mnist-helloworld 框架把这些重复劳动封装成了可复用模块。你刚学到的每个手写环节,在框架中都有对应的工程化模块:

手动实现(你刚写的)

框架模块

功能

DataLoader + transform

src/datasets/mnist.py

数据集封装,自动下载与预处理

class MNISTNet(nn.Module)

src/models/

BaseModel 抽象,统一接口

for epoch 训练循环

src/training/trainer.py

Trainer 类,封装完整训练逻辑

手动保存 checkpoint

src/training/checkpoint.py

CheckpointManager,支持断点续训

手动记录实验日志

src/training/experiment.py

ExperimentManager,YOLO 风格 runs/expN/

手写 @dataclass Config

src/config/config.py

YAML + CLI 双配置模式

# 一行命令跑通默认训练(MNIST + MyNet)
python train.py

# 更换模型和数据集
python train.py --model lenet --dataset mnist --training.epochs 30

# 使用配置文件
python train.py --config my_experiment.yaml

框架采用 YOLO 风格自动管理实验目录,每次运行创建 runs/exp1/runs/exp2/ … 自动保存配置、checkpoint 和训练曲线,不再担心覆盖实验结果。

手写是为了理解原理,框架是为了提升效率——先理解,再用工具,这是正确的学习路径。

完整的框架使用指南见使用训练框架,包括模型管理:从 nn.Module 到 BaseModel数据集管理:从 DataLoader 到 BaseDataset实验管理:从手动记录到自动追踪配置系统:从硬编码到配置文件等全部功能详解。

下一步#

掌握了完整训练流程后,下一节 调试与可视化技巧 我们将学习如何调试训练过程中的常见问题,以及如何使用 TensorBoard 等工具可视化训练过程,让"黑盒"训练变得透明可控。

从"能训练"到"会调试",让我们成为真正的深度学习工程师!

贡献者与修订历史

查看详细修订记录
  • 0cdb1e4 2026-04-29 - Heyan Zhu: feat: add model-serving chapter and update related content
  • b20ef3e 2026-04-28 - Heyan Zhu: docs: update pytorch practice section with detailed explanations and code examples
  • 0c291d7 2025-12-10 - Heyan Zhu: docs: restructure course materials and add new content