完整训练流程#
优化器:用梯度更新参数中我们掌握了如何更新参数。现在的问题是:如何把所有组件整合成一个完整的训练系统?
在 神经网络训练基础中,我们学习了训练的核心概念——Epoch、Batch、过拟合与欠拟合、正则化技巧等。但这些知识零散,需要串联成可运行的代码。
本节就是答案:我们将构建一个完整的 MNIST 训练系统,涵盖数据加载、模型训练、验证评估、模型保存等全部环节。
训练流程概览#
核心流程(对应 神经网络训练基础):
数据准备:加载数据,划分训练/验证集,创建 DataLoader
模型定义:搭建网络架构
训练循环:前向传播 → 计算损失 → 反向传播 → 更新参数
验证评估:评估模型性能,检测过拟合
模型保存:保存最佳模型
数据准备#
数据集与 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%}")
总结#
训练流程回顾#
完整训练脚本模板#
# 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 框架把这些重复劳动封装成了可复用模块。你刚学到的每个手写环节,在框架中都有对应的工程化模块:
手动实现(你刚写的) |
框架模块 |
功能 |
|---|---|---|
|
|
数据集封装,自动下载与预处理 |
|
|
|
|
|
|
手动保存 checkpoint |
|
|
手动记录实验日志 |
|
|
手写 |
|
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 等工具可视化训练过程,让"黑盒"训练变得透明可控。
从"能训练"到"会调试",让我们成为真正的深度学习工程师!
贡献者与修订历史
查看详细修订记录
-
0cdb1e42026-04-29 - Heyan Zhu: feat: add model-serving chapter and update related content -
b20ef3e2026-04-28 - Heyan Zhu: docs: update pytorch practice section with detailed explanations and code examples -
0c291d72025-12-10 - Heyan Zhu: docs: restructure course materials and add new content