使用 PyTorch 进行模型训练
你可以观看下方视频或 YouTube 上的对应内容跟随学习。
引言
在之前的视频中，我们已经讨论并演示了：
- 使用
torch.nn模块的神经网络层和函数构建模型 - 自动梯度计算的原理（这是基于梯度的模型训练的核心）
- 使用 TensorBoard 可视化训练进度和其他过程
在本视频中，我们将为你新增一些实用工具：
- 熟悉
Dataset和DataLoader抽象类，以及它们如何简化训练循环中的数据喂入流程 - 讲解特定的损失函数及其适用场景
- 学习 PyTorch 优化器（Optimizer）------ 它们能根据损失函数的计算结果调整模型权重
- 最后，将所有这些组件整合，完整演示 PyTorch 训练循环的运行过程
Dataset 和 DataLoader
Dataset 和 DataLoader 类封装了从存储介质读取数据，并以批次形式提供给训练循环的全过程：
Dataset负责读取和处理单个数据样本DataLoader从Dataset中抽取样本（自动抽取或通过自定义采样器）、组合成批次，并返回给训练循环使用。DataLoader适用于所有类型的数据集，与数据类型无关
本教程中，我们将使用 TorchVision 提供的 Fashion-MNIST 数据集。通过 torchvision.transforms.Normalize() 实现图像数据的零均值化和归一化，并下载训练集和验证集。
python
import torch
import torchvision
import torchvision.transforms as transforms
# PyTorch TensorBoard 支持
from torch.utils.tensorboard import SummaryWriter
from datetime import datetime
# 数据预处理管道
transform = transforms.Compose(
[transforms.ToTensor(), # 将PIL图像转为张量
transforms.Normalize((0.5,), (0.5,))]) # 归一化：均值0.5，标准差0.5
# 创建训练集和验证集，若本地不存在则自动下载
training_set = torchvision.datasets.FashionMNIST('./data', train=True, transform=transform, download=True)
validation_set = torchvision.datasets.FashionMNIST('./data', train=False, transform=transform, download=True)
# 创建数据加载器：训练集打乱数据，验证集不打乱
training_loader = torch.utils.data.DataLoader(training_set, batch_size=4, shuffle=True)
validation_loader = torch.utils.data.DataLoader(validation_set, batch_size=4, shuffle=False)
# 类别标签
classes = ('T恤/上衣', '裤子', '套头衫', '连衣裙', '外套',
'凉鞋', '衬衫', '运动鞋', '包', '短靴')
# 打印数据集大小
print('训练集包含 {} 个样本'.format(len(training_set)))
print('验证集包含 {} 个样本'.format(len(validation_set)))数据下载输出示例
shell
0%| | 0.00/26.4M [00:00<?, ?B/s]
0%| | 65.5k/26.4M [00:00<01:12, 365kB/s]
1%| | 164k/26.4M [00:00<00:55, 472kB/s]
2%|▏ | 655k/26.4M [00:00<00:17, 1.50MB/s]
10%|▉ | 2.62M/26.4M [00:00<00:04, 5.22MB/s]
31%|███▏ | 8.29M/26.4M [00:00<00:01, 14.7MB/s]
54%|█████▍ | 14.4M/26.4M [00:01<00:00, 21.1MB/s]
77%|███████▋ | 20.2M/26.4M [00:01<00:00, 24.9MB/s]
98%|█████████▊| 26.0M/26.4M [00:01<00:00, 27.2MB/s]
100%|██████████| 26.4M/26.4M [00:01<00:00, 18.3MB/s]
0%| | 0.00/29.5k [00:00<?, ?B/s]
100%|██████████| 29.5k/29.5k [00:00<00:00, 328kB/s]
0%| | 0.00/4.42M [00:00<?, ?B/s]
1%|▏ | 65.5k/4.42M [00:00<00:12, 360kB/s]
5%|▌ | 229k/4.42M [00:00<00:06, 678kB/s]
21%|██ | 918k/4.42M [00:00<00:01, 2.09MB/s]
83%|████████▎ | 3.67M/4.42M [00:00<00:00, 7.23MB/s]
100%|██████████| 4.42M/4.42M [00:00<00:00, 6.06MB/s]
0%| | 0.00/5.15k [00:00<?, ?B/s]
100%|██████████| 5.15k/5.15k [00:00<00:00, 58.0MB/s]
训练集包含 60000 个样本
验证集包含 10000 个样本和往常一样，我们先可视化数据以验证数据加载的正确性：
python
import matplotlib.pyplot as plt
import numpy as np
# 辅助函数：内嵌显示图像
def matplotlib_imshow(img, one_channel=False):
if one_channel:
img = img.mean(dim=0) # 单通道图像处理
img = img / 2 + 0.5 # 反归一化（恢复原始像素值范围）
npimg = img.numpy() # 转为numpy数组
if one_channel:
plt.imshow(npimg, cmap="Greys") # 灰度图显示
else:
plt.imshow(np.transpose(npimg, (1, 2, 0))) # 调整维度顺序为(H,W,C)
# 迭代数据加载器
dataiter = iter(training_loader)
images, labels = next(dataiter)
# 创建图像网格并显示
img_grid = torchvision.utils.make_grid(images)
matplotlib_imshow(img_grid, one_channel=True)
# 打印对应标签
print(' '.join(classes[labels[j]] for j in range(4)))图像可视化输出示例
凉鞋 套头衫 短靴 凉鞋模型定义
本示例中使用的模型是 LeNet-5 的变体------如果你看过本系列之前的视频，应该对这个结构很熟悉。
python
import torch.nn as nn
import torch.nn.functional as F
# PyTorch 模型需继承 torch.nn.Module
class GarmentClassifier(nn.Module):
def __init__(self):
super(GarmentClassifier, self).__init__()
# 卷积层1：输入通道1（灰度图），输出通道6，卷积核5x5
self.conv1 = nn.Conv2d(1, 6, 5)
# 池化层：2x2最大池化，步长2
self.pool = nn.MaxPool2d(2, 2)
# 卷积层2：输入通道6，输出通道16，卷积核5x5
self.conv2 = nn.Conv2d(6, 16, 5)
# 全连接层1：输入维度16*4*4，输出维度120
self.fc1 = nn.Linear(16 * 4 * 4, 120)
# 全连接层2：输入维度120，输出维度84
self.fc2 = nn.Linear(120, 84)
# 全连接层3：输入维度84，输出维度10（对应10个类别）
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
# 前向传播路径
x = self.pool(F.relu(self.conv1(x))) # 卷积1 → ReLU → 池化
x = self.pool(F.relu(self.conv2(x))) # 卷积2 → ReLU → 池化
x = x.view(-1, 16 * 4 * 4) # 展平张量（batch_size, 16*4*4）
x = F.relu(self.fc1(x)) # 全连接1 → ReLU
x = F.relu(self.fc2(x)) # 全连接2 → ReLU
x = self.fc3(x) # 全连接3（无激活，输出原始得分）
return x
# 实例化模型
model = GarmentClassifier()损失函数
本示例中我们使用交叉熵损失（CrossEntropyLoss）。为了演示，我们创建虚拟的输出和标签批次，传入损失函数并查看计算结果。
python
loss_fn = torch.nn.CrossEntropyLoss()
# 注意：损失函数期望输入是批次数据，因此我们创建4个样本的批次
# 虚拟输出：代表模型对每个输入在10个类别上的置信度
dummy_outputs = torch.rand(4, 10)
# 虚拟标签：代表每个输入对应的正确类别
dummy_labels = torch.tensor([1, 5, 3, 7])
print(dummy_outputs)
print(dummy_labels)
# 计算损失
loss = loss_fn(dummy_outputs, dummy_labels)
print('该批次的总损失值：{}'.format(loss.item()))损失函数输出示例
css
tensor([[0.8112, 0.4959, 0.7542, 0.3991, 0.2028, 0.4509, 0.1630, 0.8028, 0.2457, 0.2236],
[0.5571, 0.2879, 0.1011, 0.6718, 0.4791, 0.4399, 0.9477, 0.7674, 0.4289, 0.0269],
[0.2238, 0.1125, 0.2443, 0.2570, 0.5993, 0.0243, 0.3440, 0.2244, 0.6830, 0.9687],
[0.6951, 0.7913, 0.3395, 0.5821, 0.1017, 0.6777, 0.8524, 0.4503, 0.5550, 0.3649]])
tensor([1, 5, 3, 7])
该批次的总损失值：2.3832545280456543优化器
本示例中我们使用简单的带动量的随机梯度下降（SGD）优化器。
尝试调整优化器参数会很有启发：
- 学习率（lr）：决定优化器每次更新权重的步长。不同的学习率会如何影响训练的准确率和收敛速度？
- 动量（momentum）：使优化器在多步更新中向梯度最大的方向靠拢。调整这个值会带来什么变化？
- 尝试其他优化算法（如平均SGD、Adagrad、Adam），结果会有何不同？
python
# 优化器来自 torch.optim 包
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)训练循环
下面定义一个执行单个训练轮次（epoch） 的函数。它遍历 DataLoader 中的数据，在每次循环中执行以下操作：
- 从
DataLoader获取一批训练数据 - 清零优化器的梯度
- 执行推理：模型根据输入批次生成预测结果
- 计算预测结果与数据集标签之间的损失
- 反向传播计算权重的梯度
- 让优化器执行一次权重更新（即根据当前批次的梯度调整模型权重）
函数会每 1000 个批次打印一次损失值，并返回最后 1000 个批次的平均损失（用于与验证损失对比）。
python
def train_one_epoch(epoch_index, tb_writer):
running_loss = 0.
last_loss = 0.
# 使用 enumerate(training_loader) 而非 iter(training_loader)
# 这样可以跟踪批次索引，方便在轮次内进行进度报告
for i, data in enumerate(training_loader):
# 每个数据样本包含输入和标签
inputs, labels = data
# 重要：每个批次都要清零梯度！
optimizer.zero_grad()
# 前向传播：模型预测
outputs = model(inputs)
# 计算损失并反向传播
loss = loss_fn(outputs, labels)
loss.backward()
# 更新权重
optimizer.step()
# 累计损失并报告
running_loss += loss.item()
if i % 1000 == 999:
last_loss = running_loss / 1000 # 计算每批次平均损失
print(' 批次 {} 损失值：{}'.format(i + 1, last_loss))
# 记录到 TensorBoard
tb_x = epoch_index * len(training_loader) + i + 1
tb_writer.add_scalar('Loss/train', last_loss, tb_x)
running_loss = 0.
return last_loss轮次级操作（Per-Epoch Activity）
每个训练轮次结束后，我们需要执行两项关键操作：
- 验证：使用未参与训练的数据计算相对损失并报告
- 保存模型副本
这里我们将使用 TensorBoard 进行结果可视化。需要在命令行启动 TensorBoard，并在新的浏览器标签页中打开。
python
# 单独初始化 TensorBoard 写入器，方便后续继续训练
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
writer = SummaryWriter('runs/fashion_trainer_{}'.format(timestamp))
epoch_number = 0
# 训练轮次总数
EPOCHS = 5
# 记录最佳验证损失（初始值设为很大的数）
best_vloss = 1_000_000.
for epoch in range(EPOCHS):
print('第 {} 轮训练开始：'.format(epoch_number + 1))
# 开启梯度跟踪，执行训练
model.train(True)
avg_loss = train_one_epoch(epoch_number, writer)
# 验证阶段
running_vloss = 0.0
# 设置模型为评估模式：关闭 dropout，使用批次归一化的全局统计量
model.eval()
# 禁用梯度计算，减少内存消耗
with torch.no_grad():
for i, vdata in enumerate(validation_loader):
vinputs, vlabels = vdata
voutputs = model(vinputs)
vloss = loss_fn(voutputs, vlabels)
running_vloss += vloss
# 计算验证集平均损失
avg_vloss = running_vloss / (i + 1)
print('损失值：训练 {} 验证 {}'.format(avg_loss, avg_vloss))
# 将训练和验证的批次平均损失记录到 TensorBoard
writer.add_scalars('训练 vs. 验证损失',
{ '训练损失' : avg_loss, '验证损失' : avg_vloss },
epoch_number + 1)
writer.flush()
# 跟踪最佳性能，并保存模型状态
if avg_vloss < best_vloss:
best_vloss = avg_vloss
model_path = 'model_{}_{}'.format(timestamp, epoch_number)
torch.save(model.state_dict(), model_path)
epoch_number += 1训练过程输出示例
yaml
第 1 轮训练开始：
批次 1000 损失值：1.9134942837655544
批次 2000 损失值：0.9125308708772063
批次 3000 损失值：0.7421692375075072
批次 4000 损失值：0.6706725437843707
批次 5000 损失值：0.6199451773476321
批次 6000 损失值：0.5670827368514146
批次 7000 损失值：0.5232817540006945
批次 8000 损失值：0.5094636065706145
批次 9000 损失值：0.48667522480562914
批次 10000 损失值：0.46624169800852544
批次 11000 损失值：0.4639233185091289
批次 12000 损失值：0.4594155636130599
批次 13000 损失值：0.4203426251715282
批次 14000 损失值：0.4268976293912856
批次 15000 损失值：0.4139624496238539
损失值：训练 0.4139624496238539 验证 0.4292227625846863
第 2 轮训练开始：
批次 1000 损失值：0.4067946516573429
批次 2000 损失值：0.39065092354803344
批次 3000 损失值：0.38280759067872716
批次 4000 损失值：0.38190718797489537
批次 5000 损失值：0.3618355706045404
批次 6000 损失值：0.38087676811014537
批次 7000 损失值：0.38793834817246536
批次 8000 损失值：0.3708066731508006
批次 9000 损失值：0.37955640835637316
批次 10000 损失值：0.3687095919056446
批次 11000 损失值：0.3686677168744791
批次 12000 损失值：0.3549465914624743
批次 13000 损失值：0.33877516484307124
批次 14000 损失值：0.3504257514170604
批次 15000 损失值：0.3441471123093143
损失值：训练 0.3441471123093143 验证 0.3788158595561981
第 3 轮训练开始：
批次 1000 损失值：0.317229493965453
批次 2000 损失值：0.3474526008111425
批次 3000 损失值：0.3238617765499948
批次 4000 损失值：0.319467453146819
批次 5000 损失值：0.32635068734473316
批次 6000 损失值：0.3054303097274678
批次 7000 损失值：0.3287870975938349
批次 8000 损失值：0.33125486329917475
批次 9000 损失值：0.32507053730732877
批次 10000 损失值：0.3185750473064691
批次 11000 损失值：0.34855063943209824
批次 12000 损失值：0.3281271118210134
批次 13000 损失值：0.3221102311969153
批次 14000 损失值：0.3190491863854768
批次 15000 损失值：0.30807110028983153
损失值：训练 0.30807110028983153 验证 0.3304081857204437
第 4 轮训练开始：
批次 1000 损失值：0.2946426784224022
批次 2000 损失值：0.287669164258632
批次 3000 损失值：0.3078085952404508
批次 4000 损失值：0.29501632773253367
批次 5000 损失值：0.29912974609442244
批次 6000 损失值：0.3017247214637937
批次 7000 损失值：0.2840568649524648
批次 8000 损失值：0.29276608184371433
批次 9000 损失值：0.3079204446072108
批次 10000 损失值：0.2956557339110841
批次 11000 损失值：0.2810574857474421
批次 12000 损失值：0.30411081384155114
批次 13000 损失值：0.30809544625972196
批次 14000 损失值：0.30132432371701906
批次 15000 损失值：0.3057659530805977
损失值：训练 0.3057659530805977 验证 0.3213053345680237
第 5 轮训练开始：
批次 1000 损失值：0.26890792168607003
批次 2000 损失值：0.27582182378186915
批次 3000 损失值：0.29604251561401546
批次 4000 损失值：0.2800273624816691
批次 5000 损失值：0.2783352249941127
批次 6000 损失值：0.26813480685117974
批次 7000 损失值：0.29592685499209254
批次 8000 损失值：0.302574578157466
批次 9000 损失值：0.2913306496424302
批次 10000 损失值：0.27227645949238105
批次 11000 损失值：0.2863563198988322
批次 12000 损失值：0.2805033384739072
批次 13000 损失值：0.27869366186248956
批次 14000 损失值：0.2735917175477589
批次 15000 损失值：0.27029165163794916
损失值：训练 0.27029165163794916 验证 0.3117460012435913加载保存的模型
要加载已保存的模型，执行以下代码：
python
# 重新实例化模型
saved_model = GarmentClassifier()
# 加载权重（PATH 替换为实际的模型文件路径）
saved_model.load_state_dict(torch.load(PATH))加载完成后，模型即可用于后续操作------继续训练、推理或分析。
注意：如果你的模型构造函数参数会影响模型结构，加载时需要传入与保存时完全相同的参数，确保模型结构一致。
其他资源
- PyTorch 官方文档：数据工具（包括 Dataset 和 DataLoader）：pytorch.org
- GPU 训练中固定内存（pinned memory）的使用说明
- TorchVision、TorchText、TorchAudio 内置数据集文档
- PyTorch 损失函数文档
torch.optim包文档（包含优化器和学习率调度器等工具）- 模型保存与加载详细教程
- PyTorch 官网教程区：包含各类训练任务的教程（如不同领域的分类任务、生成对抗网络、强化学习等）
脚本总运行时间
2 分 57.614 秒总结
- PyTorch 训练循环核心组件包括：
Dataset/DataLoader（数据加载）、模型（网络结构）、损失函数（误差计算）、优化器（权重更新）； - 训练过程需区分
train()和eval()模式，验证阶段需禁用梯度计算以节省内存； - 关键实践：每个批次清零梯度、监控训练/验证损失、保存最佳模型权重。