pytorch train 模板

发表于 分类于 ， 阅读次数： 本文字数： 9.2k 阅读时长 ≈ 17 分钟

本文提供了一个完整、通用的 PyTorch 模型训练流程模板。它涵盖了从数据准备、模型构建、训练循环、测试评估到结果可视化的每一个关键步骤，是快速搭建深度学习项目的实用骨架。

🚀 整体流程一览

一个标准的 PyTorch 训练任务通常遵循以下流程。我们将通过一个具体的 CIFAR-10 图像分类任务来逐步解析代码。



graph TD
    subgraph "准备阶段"
        A["准备数据集 (CIFAR-10)"] --> B["创建 DataLoader"];
        C["定义网络模型 (Tudui)"] --> D["定义损失函数 (CrossEntropy)"];
        D --> E["定义优化器 (SGD)"];
    end

    subgraph "循环训练与评估 (Epochs)"
        F["开始 Epoch 循环"];
        F --> G["**训练模式** `model.train()`"];
        G --> H["遍历 Train DataLoader"];
        H --> I["前向传播 `output = model(input)`"];
        I --> J["计算损失 `loss = loss_fn(output, target)`"];
        J --> K["反向传播 `loss.backward()`"];
        K --> L["更新参数 `optimizer.step()`"];
        L --> M["记录训练 Loss (TensorBoard)"];
        M --> N["**评估模式** `model.eval()`"];
        N --> O["遍历 Test DataLoader (with `torch.no_grad()`)"];
        O --> P["计算测试 Loss 和 Accuracy"];
        P --> Q["记录测试结果 (TensorBoard)"];
        Q --> R["保存模型 `torch.save()`"];
        R --> F;
    end
    
    B --> H;
    E --> L;
    C --> I;
    B --> O

    A & C & E --> F

📝 代码详解

下面我们来分步解析实现这个流程的 Python 代码。

1. 导入库并准备数据

第一步是导入所有必要的库，并使用 torchvision.datasets 加载 CIFAR-10 数据集。

  • torchvision: 提供了常用的数据集、模型和图像转换。
  • Tensorboard SummaryWriter: 用于可视化训练过程。
  • DataLoader: 用于批量加载数据。
  • transforms.ToTensor: 将 PIL Image 或 NumPy ndarray 转换为 FloatTensor，并将像素值从 [0, 255] 缩放到 [0.0, 1.0]。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
import torchvision
from torch.utils.tensorboard import SummaryWriter
from model import *  # 模型定义在 model.py 中
import torch
from torch import nn
from torch.utils.data import DataLoader

# 准备数据集
train_data = torchvision.datasets.CIFAR10(root="../data", train=True, transform=torchvision.transforms.ToTensor(),
                                          download=True)
test_data = torchvision.datasets.CIFAR10(root="../data", train=False, transform=torchvision.transforms.ToTensor(),
                                         download=True)

# length 长度
train_data_size = len(train_data)
test_data_size = len(test_data)
print("训练数据集的长度为: {}".format(train_data_size))
print("测试数据集的长度为: {}".format(test_data_size))

# 利用 DataLoader 来加载数据集
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)

2. 定义核心组件

接下来，我们创建模型实例、定义损失函数和优化器。

  • 模型 (Tudui): 这是一个自定义的神经网络模型，其具体结构定义在 model.py 文件中。 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    import torch.nn as nn

    class Tudui(nn.Module):
        def __init__(self):
            super(Tudui, self).__init__()
            self.model = nn.Sequential(
                nn.Conv2d(3, 32, 5, 1, 2),
                nn.MaxPool2d(2),
                nn.Conv2d(32, 32, 5, 1, 2),
                nn.MaxPool2d(2),
                nn.Conv2d(32, 64, 5, 1, 2),
                nn.MaxPool2d(2),
                nn.Flatten(),
                nn.Linear(64*4*4, 64),
                nn.Linear(64, 10)
            )

        def forward(self, x):
            x = self.model(x)
            return x
  • 损失函数 (nn.CrossEntropyLoss): 交叉熵损失函数，常用于多分类任务。它内部已经包含了 Softmax 操作。
  • 优化器 (torch.optim.SGD): 随机梯度下降（Stochastic Gradient Descent）优化器，用于根据损失函数的梯度来更新模型的参数（权重和偏置）。lr 是学习率（learning rate）。
1
2
3
4
5
6
7
8
9
# 创建网络模型
tudui = Tudui()

# 损失函数
loss_fn = nn.CrossEntropyLoss()

# 优化器
learning_rate = 1e-2
optimizer = torch.optim.SGD(tudui.parameters(), lr=learning_rate)

3. 设置训练参数与 TensorBoard

在开始训练之前，我们需要初始化一些计数器和超参数，并设置 TensorBoard 的 SummaryWriter 来记录训练数据。

  • total_train_step / total_test_step: 记录训练和测试的步数，作为 TensorBoard 的 x 轴。
  • epoch: 训练的总轮数。一轮（epoch）指完整地遍历一次整个训练数据集。
  • SummaryWriter: 创建一个写入器实例，它会将事件日志写入指定的目录（../logs_train）。
1
2
3
4
5
6
7
8
9
10
# 设置训练网络的一些参数
# 记录训练的次数
total_train_step = 0
# 记录测试的次数
total_test_step = 0
# 训练的轮数
epoch = 10

# 添加tensorboard
writer = SummaryWriter("../logs_train")

4. 训练与评估循环

这是整个脚本的核心。外层是 epoch 循环，内层包含一个训练步骤和一个评估步骤。

训练步骤
  • tudui.train(): 将模型设置为训练模式。这会启用 Dropout 和 BatchNorm 等层。
  • optimizer.zero_grad(): 在计算新梯度之前，清除上一轮的旧梯度。
  • loss.backward(): 计算损失函数相对于模型参数的梯度（反向传播）。
  • optimizer.step(): 根据计算出的梯度更新模型参数。
  • writer.add_scalar(): 每 100 步记录一次训练损失。
评估步骤
  • tudui.eval(): 将模型设置为评估模式。这会禁用 Dropout 和 BatchNorm 等层，确保评估结果的确定性。
  • with torch.no_grad(): 在此代码块内，所有计算都不会追踪梯度，从而节省内存和计算资源。
  • 计算准确率: (outputs.argmax(1) == targets).sum() 计算一个批次中预测正确的样本数量。argmax(1) 找到概率最高的类别的索引。
  • 记录结果: 在一轮 epoch 结束后，计算并打印在整个测试集上的总损失和正确率，并用 TensorBoard 记录。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
for i in range(epoch):
    print("-------第 {} 轮训练开始-------".format(i+1))

    # 训练步骤开始
    tudui.train()
    for data in train_dataloader:
        imgs, targets = data
        outputs = tudui(imgs)
        loss = loss_fn(outputs, targets)

        # 优化器优化模型
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_train_step = total_train_step + 1
        if total_train_step % 100 == 0:
            print("训练次数: {}, Loss: {}".format(total_train_step, loss.item()))
            writer.add_scalar("train_loss", loss.item(), total_train_step)

    # 测试步骤开始
    tudui.eval()
    total_test_loss = 0
    total_accuracy = 0
    with torch.no_grad():
        for data in test_dataloader:
            imgs, targets = data
            outputs = tudui(imgs)
            loss = loss_fn(outputs, targets)
            total_test_loss = total_test_loss + loss.item()
            accuracy = (outputs.argmax(1) == targets).sum()
            total_accuracy = total_accuracy + accuracy

    print("整体测试集上的Loss: {}".format(total_test_loss))
    print("整体测试集上的正确率: {}".format(total_accuracy/test_data_size))
    writer.add_scalar("test_loss", total_test_loss, total_test_step)
    writer.add_scalar("test_accuracy", total_accuracy/test_data_size, total_test_step)
    total_test_step = total_test_step + 1

    torch.save(tudui, "tudui_{}.pth".format(i))
    # torch.save(tudui.state_dict(), "tudui_{}.pth".format(i))
    print("模型已保存")

writer.close()

💡 模型的两种保存方式

PyTorch 提供了两种主要的方式来保存模型，代码中展示了这两种方式：

  1. 保存整个模型对象 (torch.save(model, PATH))
    • 做法: torch.save(tudui, "tudui_complete.pth")
    • 优点: 非常简单，加载时也只需一行 model = torch.load("tudui_complete.pth")。
    • 缺点: 这种方式使用了 Python 的 pickle 来序列化整个对象，它将模型的类定义、目录结构等都绑定在了一起。如果你的项目代码发生了变化（例如，重命名了模型类或存放模型的文件），加载时就可能会失败。可移植性较差。
  2. 仅保存模型参数 (torch.save(model.state_dict(), PATH)) - (推荐)
    • 做法: torch.save(tudui.state_dict(), "tudui_params.pth")
    • 优点: 只保存模型的“状态字典”(state_dict)，即所有的权重和偏置。它不依赖于具体的类结构，非常轻量和灵活，是官方推荐的、更具 Python 风格的方式。
    • 缺点: 加载时需要先创建模型实例，然后再将参数加载进去： 
      1
      2
      3
      4
      # 需要先创建模型实例
      model = Tudui() 
      # 再加载参数
      model.load_state_dict(torch.load("tudui_params.pth"))

结论：为了代码的健壮性和可移植性，强烈推荐使用第二种方法，即只保存和加载模型的 state_dict。

为何需要 train() 和 eval()模式?

模型的某些层（如 Dropout 和 BatchNorm）在训练和评估时的行为是不同的。 - 训练时 (train()模式): Dropout会随机失活一些神经元来防止过拟合；BatchNorm 会计算当前批次的均值和方差来归一化数据。 - 评估时 (eval()模式): Dropout会被禁用；BatchNorm 会使用在整个训练过程中学习到的全局均值和方差。

切换模式能确保模型在正确的状态下进行训练和推理。

⚡️ 迁移到 GPU 训练

如果你的机器上有 NVIDIA 显卡并正确安装了 CUDA，那么将训练迁移到 GPU 上可以极大地加速计算。PyTorch 让这个过程变得非常简单，核心只有三步：定义设备、模型上设备、数据上设备。

1. 定义设备

在代码的准备阶段，首先定义一个 device 对象。下面的代码会自动检测 CUDA 是否可用，如果可用则选择 GPU，否则回退到 CPU。

1
2
3
# 定义训练的设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("将使用 {} 设备进行训练".format(device))

2. 将模型和损失函数移动到设备

创建完网络模型和损失函数后，调用 .to(device) 方法将它们发送到指定的设备上。

1
2
3
4
5
6
7
8
9
10
11
# 创建网络模型
tudui = Tudui()
tudui.to(device) # 发送到设备

# 损失函数
loss_fn = nn.CrossEntropyLoss()
loss_fn.to(device) # 发送到设备

# 优化器 (无需移动)
learning_rate = 1e-2
optimizer = torch.optim.SGD(tudui.parameters(), lr=learning_rate)

提示：优化器在初始化时已经注册了模型的参数。当模型被 .to(device) 移动后，优化器会自动感知到参数在 GPU 上，因此优化器本身不需要被移动。

3. 将数据移动到设备

最后，也是最关键的一步，是在训练和测试的循环中，将每一批次的数据也发送到同一个设备上。

1
2
3
4
5
6
7
8
9
10
# 在训练/测试循环中
for data in train_dataloader:
    imgs, targets = data
    imgs = imgs.to(device) # 将图像数据发送到设备
    targets = targets.to(device) # 将标签数据发送到设备

    # 现在模型和数据都在同一个设备上，可以进行计算
    outputs = tudui(imgs)
    loss = loss_fn(outputs, targets)
    ...

完成这三步修改后，你的训练代码就能在 GPU 上飞速运行了！

🔬 模型验证

训练完成后，我们可以加载保存好的模型，用它来对全新的、从未见过的数据进行预测。下面是一个完整的验证（或称推理，Inference）流程示例。

1. 准备环境和数据

首先，我们需要导入相关库，并准备一张待预测的图片。请注意，这张图片需要经过和训练时完全相同的预处理（transform）。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
import torch
from PIL import Image
from torchvision import transforms
# 确保 Tudui 类的定义可见
from model import Tudui 

# CIFAR-10 类别
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

# 加载图片
image_path = "imgs/dog.png" # 请替换为你的图片路径
image = Image.open(image_path)
# PNG 图片可能有4个通道 (RGBA)，需要转换为3通道 (RGB)
image = image.convert('RGB') 

# 和训练时一样的 Transform
transform = transforms.Compose([transforms.Resize((32, 32)),
                                transforms.ToTensor()])
image = transform(image)
print(f"图片预处理后形状: {image.shape}")

# 调整形状以匹配模型输入
# 模型需要一个4D张量: [batch_size, channels, height, width]
image = torch.reshape(image, (1, 3, 32, 32))
print(f"调整后送入模型的形状: {image.shape}")

2. 加载模型并进行预测

接下来，我们加载已经训练好的模型文件，并进行预测。

  • torch.load(): 加载模型文件。
  • map_location: 这是一个非常重要的参数！如果你的模型是在 GPU 上训练保存的（.pth 文件中包含了 GPU 张量），而你现在想在只有 CPU 的机器上加载它，就必须指定 map_location=torch.device('cpu')，它会智能地将所有张量重新映射到 CPU 上。
  • model.eval(): 切换到评估模式，这和训练时的验证步骤一样重要。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
# 加载模型
# 假设模型是在 GPU 保存的，我们现在用 CPU 加载
# 注意：如果使用 torch.save(model)，加载时不需要 Tudui 类的定义
# 但如果使用 torch.save(model.state_dict())，则需要先 model=Tudui()
model = torch.load("tudui_29_gpu.pth", map_location=torch.device('cpu'))
print(model)

# 模型评估与预测
model.eval()
with torch.no_grad():
    output = model(image)

print(f"模型输出 (logits): \n{output}")

# 获取预测结果
predicted_class_index = output.argmax(1)
print(f"预测结果的索引: {predicted_class_index.item()}")
print(f"预测的类别是: {classes[predicted_class_index.item()]}")

通过这个流程，你就可以用训练好的模型来处理任何新的输入数据了。

📈 启动 TensorBoard

训练开始后，你可以在终端中运行以下命令来启动 TensorBoard 服务：

1
tensorboard --logdir=../logs_train

然后打开浏览器访问 http://localhost:6006 即可看到 train_loss, test_loss 和 test_accuracy 的变化曲线。

📝 小结

这个模板展示了 PyTorch 训练一个模型的标准范式：

  1. 数据加载：使用 Dataset 和 DataLoader。
  2. 模型构建：定义网络、损失函数和优化器。
  3. 循环迭代：通过 for 循环进行多轮 (epoch) 训练。
  4. 训练模式：在每轮 epoch 开始时调用 model.train()，进行前向传播、计算损失、反向传播和参数更新。
  5. 评估模式：调用 model.eval() 和 with torch.no_grad()，在测试集上评估模型性能。
  6. 可视化与保存：使用 TensorBoard 监控过程，使用 torch.save 保存模型。

参考自 B站up 我是土堆，链接：https://www.bilibili.com/video/BV1hE411t7RN?spm_id_from=333.788.videopod.episodes&vd_source=a1e48cea1c1e1104579f0aefa6e00490&p=1