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

JJJJ_iii2025-10-10 18:36

文章目录

视频链接
【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文件，这个文件只负责一件事：定义神经网络的结构。这是一种良好的工程实践，它让模型结构与训练逻辑分离，使代码更清晰。

python 复制代码 
# 文件: model.py

import torch
from torch import nn

class 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中定义的模型，并执行完整的训练和测试流程。

python 复制代码 
# 文件: 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 = data
            outputs = 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将测试损失和准确率记录到TensorBoard
    writer.add_scalar("test_loss", total_test_loss, i) # x轴使用轮数i
    writer.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个样本的预测类别。
相关推荐
科技小E5 小时前
流媒体视频技术在明厨亮灶场景中的深度应用
人工智能
geneculture5 小时前
融智学院十大学部知识架构示范样板
人工智能·数据挖掘·信息科学·哲学与科学统一性·信息融智学
无风听海6 小时前
神经网络之交叉熵与 Softmax 的梯度计算
人工智能·深度学习·神经网络
算家计算6 小时前
AI树洞现象：是社交降级，还是我们都在失去温度？
人工智能
java1234_小锋6 小时前
TensorFlow2 Python深度学习 - TensorFlow2框架入门 - 神经网络基础原理
python·深度学习·tensorflow·tensorflow2
JJJJ_iii6 小时前
【深度学习03】神经网络基本骨架、卷积、池化、非线性激活、线性层、搭建网络
网络·人工智能·pytorch·笔记·python·深度学习·神经网络
sensen_kiss6 小时前
INT301 Bio-computation 生物计算（神经网络）Pt.1 导论与Hebb学习规则
人工智能·神经网络·学习
玉石观沧海6 小时前
高压变频器故障代码解析F67 F68
运维·经验分享·笔记·分布式·深度学习
mwq301236 小时前
GPT系列模型演进：从GPT-1到GPT-4o的技术突破与差异解析
人工智能
热门推荐
01两千字总结：Codex 国内如何安装和使用的教程，以及如何设置中文回答02GitHub 镜像站点03UV安装并设置国内源0446个Nano-banana 精选提示词，持续更新中05Linux下V2Ray安装配置指南06智能库存管理的需求预测模型：从业务痛点到落地代码的完整实践07GitLab 零基础入门指南：从安装到项目管理全流程08Cursor Plan Mode：AI 终于知道先想后做了09一文了解国产算子编程语言 TileLang，TileLang 对国产开源生态的影响与启示10jdk21下载、安装（Windows、Linux、macOS）