【深度学习05】PyTorch:完整的模型训练套路

---

视频链接
【PyTorch深度学习快速入门教程（绝对通俗易懂！）【小土堆】】p27-29

---

完整的模型训练套路

本章的目标是将前面所有独立的知识点——数据加载、网络搭建、损失函数、优化器——全部整合起来，形成一个标准、规范且可复用的神经网络训练与测试流程。我们将从零开始，通过编写model.py和train.py两个核心文件，构建一个完整的项目。

总体思路

一个标准的模型训练脚本，其逻辑流程是固定且清晰的，可以分为以下八个核心步骤：

1. 准备数据集：从硬盘加载或从网络下载数据集，并切分为训练集和测试集。使用DataLoader进行封装，以实现批量（batch）加载。
2. 搭建网络模型：在一个独立的model.py文件中清晰地定义神经网络的结构。
3. 创建模型实例：在主训练脚本train.py中，实例化我们定义好的网络模型。
4. 定义损失函数和优化器：根据任务类型（如分类、回归）选择合适的损失函数和优化算法。
5. 设置训练循环：代码的主体是一个双层循环。外层循环控制训练的总轮数（epoch），内层循环负责遍历数据集中的每一个批次。
6. 核心训练步骤：在内层循环中，严格执行训练的“四步曲”：
   • 前向传播：将数据输入模型，得到预测结果。
   • 计算损失：用预测结果和真实标签计算损失值。
   • 反向传播：调用loss.backward()计算梯度。
   • 更新参数：调用optimizer.step()更新模型权重。
7. 添加测试步骤：在每一轮（epoch）训练结束后，使用独立的测试集来评估模型的性能。这可以帮助我们监控模型的泛化能力，判断是否发生过拟合。
8. 保存模型与可视化：在训练过程中，定期保存模型的检查点（checkpoint），并使用TensorBoard等工具记录损失、准确率等关键指标的变化，实现训练过程的可视化。

我们将严格遵循此思路，构建我们的代码。

---

1. 搭建神经网络 (model.py)

我们首先创建model.py文件，这个文件只负责一件事：定义神经网络的结构。这是一种良好的工程实践，它让模型结构与训练逻辑分离，使代码更清晰。

# 文件: model.pyimport torch
from torch import nnclass Tudui(nn.Module):"""一个针对CIFAR-10数据集(3x32x32)的卷积神经网络模型。该结构参考了经典的CIFAR-10模型设计。"""def __init__(self):super(Tudui, self).__init__()self.model = nn.Sequential(# 第一个卷积层# 输入形状: [batch_size, 3, 32, 32]nn.Conv2d(in_channels=3, out_channels=32, kernel_size=5, stride=1, padding=2),# 经过卷积后，形状变为: [batch_size, 32, 32, 32] (因为 stride=1, padding=2 保持了尺寸)# 第一个最大池化层nn.MaxPool2d(kernel_size=2),# 经过池化后，形状变为: [batch_size, 32, 16, 16] (高度和宽度减半)# 第二个卷积层nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5, stride=1, padding=2),# 形状仍为: [batch_size, 32, 16, 16]# 第二个最大池化层nn.MaxPool2d(kernel_size=2),# 形状变为: [batch_size, 32, 8, 8]# 第三个卷积层nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, stride=1, padding=2),# 形状变为: [batch_size, 64, 8, 8]# 第三个最大池化层nn.MaxPool2d(kernel_size=2),# 形状变为: [batch_size, 64, 4, 4]# 压平层，为全连接层做准备nn.Flatten(),# 形状变为: [batch_size, 64 * 4 * 4], 即 [batch_size, 1024]# 第一个全连接层nn.Linear(in_features=1024, out_features=64),# 形状变为: [batch_size, 64]# 第二个全连接层 (输出层)nn.Linear(in_features=64, out_features=10)# 最终输出形状: [batch_size, 10])def forward(self, x):"""定义数据的前向传播过程"""x = self.model(x)return x# --- 用于验证模型正确性的代码 ---
if __name__ == '__main__':tudui = Tudui()# 创建一个假的输入张量来测试网络# torch.ones创建一个全为1的张量# (64, 3, 32, 32) 表示一个批次包含64张图片，每张图片3个通道，高和宽都是32像素input_tensor = torch.ones((64, 3, 32, 32))output_tensor = tudui(input_tensor)# 打印输出的形状，以验证网络结构是否按预期工作print(output_tensor.shape) # 期望输出: torch.Size([64, 10])

代码参数超详细讲解 (model.py)

• nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding):

  • in_channels=3: 输入的通道数。对于CIFAR-10这种RGB彩色图像，通道数是3（红、绿、蓝）。
  • out_channels=32: 输出的通道数。这个数值也代表了卷积核（或称滤波器）的数量。这里我们用了32个卷积核，所以会产生32张特征图（feature map）。
  • kernel_size=5: 卷积核的大小。这里是5x5的卷积核。
  • stride=1: 步长，即卷积核每次在图像上滑动的距离。1表示一次移动一个像素。
  • padding=2: 填充。在图像的边界周围添加额外的像素（这里是2圈）。公式 Output_size = (Input_size - Kernel_size + 2*Padding) / Stride + 1，代入数值 (32 - 5 + 2*2)/1 + 1 = 32。设置padding=2的目的是为了让卷积后的特征图尺寸与输入保持不变。

• nn.MaxPool2d(kernel_size):

  • kernel_size=2: 池化窗口的大小。这里是2x2的窗口。它会从输入的2x2区域中取出最大的那个值作为输出，从而将特征图的高度和宽度都缩小一半（例如，从32x32缩小到16x16）。

• nn.Flatten():

  • 这是一个没有参数的层。它的唯一作用就是将输入的多维张量“压平”成一个一维向量。例如，一个[64, 64, 4, 4]的张量（batch_size, channels, height, width）会被转换成[64, 1024]的张量（batch_size, features），其中 1024 = 64 * 4 * 4。

• nn.Linear(in_features, out_features):

  • in_features=1024: 输入特征的维度（神经元数量）。这个数值必须与前一层Flatten的输出维度完全匹配。
  • out_features=10: 输出特征的维度。在最后一层，这个数值必须等于我们任务的类别总数。CIFAR-10有10个类别，所以这里是10。

---

2. 完整的训练与测试脚本 (train.py)

这个文件是项目的核心，它将调用model.py中定义的模型，并执行完整的训练和测试流程。

# 文件: train.py# 导入所有需要的库
import torch
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter# 从我们自己写的 model.py 文件中，导入我们定义的 Tudui 模型类
from model import Tudui# -------------------- 1. 准备数据集 --------------------
# 使用 torchvision.datasets.CIFAR10 加载CIFAR-10的训练数据集
train_data = torchvision.datasets.CIFAR10(root="./data",        # 数据集下载后存放的根目录train=True,           # 指定这是训练集 (如果为False，则表示加载测试集)transform=torchvision.transforms.ToTensor(), # 创建一个转换，将图像数据转换为PyTorch张量，并自动将像素值从[0, 255]归一化到[0.0, 1.0]download=True         # 如果在'root'目录下找不到数据集，则自动从网上下载
)
# 加载CIFAR-10的测试数据集
test_data = torchvision.datasets.CIFAR10(root="./data",train=False,          # 指定这是测试集transform=torchvision.transforms.ToTensor(),download=True
)# 获取训练集和测试集的大小，用于后续计算（如准确率）
train_data_size = len(train_data)
test_data_size = len(test_data)
# 使用f-string打印信息，更直观
print(f"训练数据集的长度为: {train_data_size}") # 输出: 50000
print(f"测试数据集的长度为: {test_data_size}")   # 输出: 10000# 使用DataLoader将数据集封装成可迭代对象，实现批量加载
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)# -------------------- 2. 创建网络模型实例 --------------------
# 实例化我们从model.py中导入的Tudui模型
tudui = Tudui()# -------------------- 3. 定义损失函数 --------------------
# 使用交叉熵损失函数，它在多分类问题中非常常用
# 它内部已经包含了Softmax操作，所以我们的模型输出层不需要加Softmax
loss_fn = nn.CrossEntropyLoss()# -------------------- 4. 定义优化器 --------------------
# 定义学习率，这是训练中最重要的超参数之一
learning_rate = 0.01 
# 创建一个SGD(随机梯度下降)优化器
# 第一个参数 tudui.parameters() 是告诉优化器，模型中所有需要更新的参数都在这里
# 第二个参数 lr=learning_rate 是设置学习率
optimizer = torch.optim.SGD(tudui.parameters(), lr=learning_rate)# -------------------- 5. 设置训练网络的一些超参数 --------------------
total_train_step = 0 # 定义一个计数器，记录总的训练步数（一个batch算一步）
total_test_step = 0  # 定义一个计数器，记录总的测试轮数（一个epoch算一轮）
epoch = 10           # 定义训练的总轮数# -------------------- 6. 添加TensorBoard用于可视化 --------------------
# 创建一个SummaryWriter实例，它会将日志数据写入到'./logs_train'文件夹
writer = SummaryWriter("./logs_train")# -------------------- 7. 开始训练循环 --------------------
# 外层循环控制训练的轮数(epoch)
for i in range(epoch):print(f"-------- 第 {i+1} 轮训练开始 --------")# --- 训练步骤 ---# 调用 tudui.train()，将模型设置为训练模式。# 这对于包含Dropout或BatchNorm层的模型是必需的，以确保它们在训练时正常工作。tudui.train()# 内层循环遍历训练数据加载器，每次取出一个批次(batch)的数据for data in train_dataloader:# 从data中解包出图像数据(imgs)和对应的标签(targets)imgs, targets = data# 1. 前向传播：将图像数据输入到模型中，得到预测输出outputs = tudui(imgs)# 2. 计算损失：使用损失函数比较预测输出和真实标签，得到损失值loss = loss_fn(outputs, targets)# 3. 反向传播：这是PyTorch自动求导的核心# 首先，清除上一轮计算残留的梯度optimizer.zero_grad()# 然后，调用loss.backward()计算当前损失相对于模型所有参数的梯度loss.backward()# 最后，调用optimizer.step()，优化器会根据计算出的梯度来更新模型的参数optimizer.step()# 更新总训练步数total_train_step += 1# 为了避免打印过于频繁，我们设置每训练100步打印一次信息if total_train_step % 100 == 0:# loss是一个张量，loss.item()可以从中获取其数值print(f"训练步数: {total_train_step}, Loss: {loss.item()}")# 使用writer将训练损失记录到TensorBoard，方便可视化writer.add_scalar("train_loss", loss.item(), total_train_step)# --- 测试步骤 ---# 在每轮训练结束后，进行一次测试来评估模型的性能# 调用 tudui.eval()，将模型设置为评估模式。# 这会关闭Dropout层，并让BatchNorm层使用全局统计数据，确保测试结果的确定性。tudui.eval()total_test_loss = 0  # 初始化测试集上的总损失total_accuracy = 0   # 初始化测试集上的总正确数# 使用 with torch.no_grad(): 块，暂时禁用所有梯度计算。# 这可以节省内存并加快计算速度，因为在测试时我们不需要进行反向传播。with torch.no_grad():# 遍历测试数据加载器for data in test_dataloader:imgs, targets = dataoutputs = tudui(imgs)loss = loss_fn(outputs, targets)# 累加每个批次的损失total_test_loss += loss.item()# 计算这个批次的正确预测数# outputs.argmax(1) 会返回在维度1上最大值的索引，即模型预测的类别# (outputs.argmax(1) == targets) 会得到一个布尔张量，预测正确的位置为True# .sum() 会将所有True（计为1）加起来，得到正确预测的数量accuracy = (outputs.argmax(1) == targets).sum()# 累加每个批次的正确数total_accuracy += accuracy# 打印本轮训练结束后，在整个测试集上的性能指标print(f"本轮训练结束，在测试集上的总Loss为: {total_test_loss}")print(f"本轮训练结束，在测试集上的总正确率为: {total_accuracy / test_data_size}")# 使用writer将测试损失和准确率记录到TensorBoardwriter.add_scalar("test_loss", total_test_loss, i) # x轴使用轮数iwriter.add_scalar("test_accuracy", total_accuracy / test_data_size, i)# 保存当前轮次的模型状态torch.save(tudui, f"tudui_epoch_{i}.pth")print(f"模型 tudui_epoch_{i}.pth 已保存")# 所有训练轮数结束后，关闭SummaryWriter
writer.close()

代码参数超详细讲解 (train.py)

• torchvision.datasets.CIFAR10(...):

  • root="./data": 指定一个文件夹路径，PyTorch会把数据集下载并解压到这里。
  • train=True / False: 布尔值，True表示加载训练集，False表示加载测试集。
  • transform=torchvision.transforms.ToTensor(): 对加载的图像进行预处理。ToTensor做了两件核心的事：1. 将PIL Image格式的图像或Numpy数组转换为torch.FloatTensor。2. 将图像的像素值从[0, 255]的整数范围，缩放到[0.0, 1.0]的浮点数范围。归一化是神经网络训练中至关重要的一步，能让模型收敛得更快、更稳定。
  • download=True: 如果root路径下找不到数据集，就自动从网上下载。

• DataLoader(dataset, batch_size):

  • dataset: 要加载的数据集对象，即上面创建的train_data或test_data。
  • batch_size=64: 批处理大小。表示每次从数据集中取出64个样本打包成一个批次。这是训练效率和内存消耗之间的一个权衡。

• tudui.train() 和 tudui.eval():

  • train(): 将模型切换到训练模式。这会启用Dropout层和BatchNorm层的训练行为（例如，BatchNorm会计算并更新每个批次的均值和方差）。
  • eval(): 将模型切换到评估/测试模式。这会禁用Dropout层，并让BatchNorm层使用在整个训练集上学习到的固定的均值和方差，确保测试结果是确定性的。在训练和测试之间切换这两种模式是绝对必要的，否则会导致结果不一致或错误。

• with torch.no_grad()::

  • 这是一个上下文管理器，它告诉PyTorch在这个代码块内部的所有计算都不需要计算梯度。在测试阶段，我们只关心模型的前向传播结果，不需要进行反向传播，所以禁用梯度可以：1. 节省大量内存；2. 显著加快计算速度。

• outputs.argmax(1):

  • outputs的形状是[64, 10]，代表64个样本，每个样本对应10个类别的原始得分（logits）。
  • argmax(1)沿着维度1（即类别维度）查找最大值的索引。例如，如果一个样本的得分是[0.1, 2.5, 0.3, ...]，argmax会返回索引1，代表模型预测这个样本属于第1类。
  • 所以 outputs.argmax(1) 会返回一个形状为[64]的张量，包含了对这个批次中64个样本的预测类别。