PyTorch 模型训练完整工作流程详解

定义模型类

通过子类化 nn.Module 创建 PyTorch 模型的基本构建模块。对于 nn.Module 子类的对象，必须定义 forward() 方法。

import torch.nn as nn

class LinearModel(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        # 定义线性层，等价于 y = wx + b
        self.linear = nn.Linear(input_dim, output_dim)

    def forward(self, x):
        return self.linear(x)

# 实例化模型
model = LinearModel(input_dim=1, output_dim=1)

# 查看模型参数
for name, param in model.named_parameters():
    print(f"{name}: {param.shape}, requires_grad={param.requires_grad}")

由于我们使用了默认初始化，weights 和 bias 都是随机张量。

3. 使用推理模式进行预测

我们将测试集 X_test 的数据通过 forward() 方法得到模型的预测结果 y_preds。

with torch.inference_mode():
    y_preds = model(x_test)

print("初始预测值:", y_preds[:5])
print("真实值:", y_test[:5])

在旧版本的 PyTorch 中，我们使用的是 torch.no_grad()。无论是哪种，我们都抑制了 PyTorch 自动计算梯度的功能，这有助于加速 forward 流程。由于我们模型参数都是随机初始化的，而且模型又没有经过训练，所以现在模型的预测性能显然很糟糕。

4. 训练模型

训练模型的过程其实就是不断更新我们之前由默认初始化或 nn.Parameter 初始化的模型参数 weights 和 bias。

损失函数和优化器

损失函数（loss function） 负责衡量模型的预测值（y_preds）和真实值（y_tests）之间的差异，通常越小越好。torch.nn 模块内置了很多损失函数。一般对于回归任务，使用平均绝对误差（torch.nn.L1Loss()）；对于二分类任务，使用二元交叉熵损失（torch.nn.BCELoss()）。

优化器（optimizer） 负责告诉模型如何调整内部参数以最小化损失函数。torch.optim 模块实现了许多优化函数。常见的优化函数有随机梯度下降（torch.optim.SGD()）和 Adam（torch.optim.Adam()）。

我们需要根据不同的任务选择合适的损失函数和优化器。由于我们现在是在预测数值，所以我们选择 torch.nn.L1Loss()。

# 定义损失函数
loss_fn = nn.L1Loss()

# 定义优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

平均绝对误差测量两点（预测和标签）之间的绝对差，然后取所有样本的平均值。

我们将选择 torch.optim.SGD(params, lr)：

• params 就是我们要更新的模型参数，在这里就是 weights 和 bias。
• lr 是学习率（learning rate）的缩写。它表示参数更新的幅度，lr 越大，则参数更新幅度越大，越可能错过最佳值；lr 越小，则参数更新幅度越小，模型训练速度就越慢。

创建训练循环

现在我们可以创建训练循环（training loop）和测试循环（testing loop）。

模型在训练循环中遍历训练集中的样本，学习样本的特征和标签之间的关系。

典型的训练循环如下所示，首先我们需要设定模型训练的循环次数 epochs。然后在每个 epoch（也就是每次循环中）：

1. 通过模型的 forward() 方法计算模型对于训练集的预测 y_pred。
2. 根据损失函数计算模型预测值 y_pred 和真实值 y_true 之间的损失 loss。
3. 清零优化器的梯度，否则梯度会在每次循环中累积。
4. 利用计算损失相对于模型参数的梯度：loss.backward()。
5. 使用随机梯度下降算法用梯度更新模型参数：optimizer.step()。

epochs = 100

for epoch in range(epochs):
    # 1. 前向传播
    y_pred = model(x_train)
    
    # 2. 计算损失
    loss = loss_fn(y_pred, y_train)
    
    # 3. 梯度清零
    optimizer.zero_grad()
    
    # 4. 反向传播
    loss.backward()
    
    # 5. 更新参数
    optimizer.step()
    
    if (epoch + 1) % 10 == 0:
        print(f"Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}")

模型在测试循环中遍历测试集中的样本，衡量模型学习到关系在未知数据集中的表现。

测试循环就简单了，还是在每个 epoch（也就是每次循环中）：

1. 通过 model.eval() 说明现在是在评估状态。
2. 在 torch.inference_mode() 中：
3. 利用 forward() 方法计算模型在测试集上预测 test_pred。
4. 计算 test_pred 和真实值 y_test 之间的损失：test_loss。

隔几个 epoch（一般是 10）输出一下平均损失。

5. 用训练好的模型进行预测

模型训练好之后就可以用来预测了，有的地方也称为推理（inference）。但是在此之前，我们需要做好以下三件事：

• 设置评估模式：model.eval()。
• 在 with torch.inference_mode(): … 上下文管理器中进行预测。
• 模型和数据必须要在同一个设备上。

model.eval()
with torch.inference_mode():
    final_preds = model(x_test)

print("最终预测值:", final_preds[:5].detach())
print("真实值:", y_test[:5])

看起来结果已经是相当接近了。

6. 保存和加载 PyTorch 模型

有三种主要的方法可以保存和加载 PyTorch 模型：

• torch.save：使用 Python 的 pickle 工具将对象序列化后保存到磁盘。模型、张量和其他 Python 对象都可以用这种方法保存。
• torch.load：将 torch.save 保存的文件反序列化之后加载到内存，我们还可以选择加载到 GPU 还是 CPU。
• torch.nn.Module.load_state_dict：加载保存在磁盘上的模型参数字典（用 model.state_dict() 方法可以获得）。

不过根据 pickle 的官方文档介绍，pickle 是不安全的，所以你必须要相信你要加载的对象是安全的。

保存和加载模型的推荐方式是通过模型的 state_dict() 方法。

保存模型

torch.save 的 obj 参数是模型的参数字典、f 参数是保存的路径。一般我们将 PyTorch 模型保存为以 .pt 或 .pth 这样的后缀。

不过以这种方式保存的模型仅仅是保存了模型训练之后的参数字典，而不是整个模型。所以我们首先需要初始化一个模型实例，然后用保存在磁盘上的模型参数字典更新这个初始化的模型参数。

# 保存模型参数
torch.save(model.state_dict(), 'linear_model.pth')
print("模型已保存")

加载模型

看到上面的提示说明加载成功。

加载成功之后，我们再次用这个新加载的模型做预测，并将预测结果和之前训练的模型进行比较。

# 重新实例化模型
new_model = LinearModel(input_dim=1, output_dim=1)

# 加载参数
new_model.load_state_dict(torch.load('linear_model.pth'))
new_model.eval()

# 验证一致性
with torch.inference_mode():
    loaded_preds = new_model(x_test)

print("加载后预测值:", loaded_preds[:5].detach())
print("原预测值:", final_preds[:5].detach())

结果显然是一致的。

根据官方文档，保存整个模型的缺点是序列化数据绑定到特定的类以及保存模型时使用的确切目录结构。这样做的原因是因为 pickle 不保存模型类本身。相反，它保存包含该类的文件的路径，该路径在加载时使用。因此，在其他项目中使用或重构后，您的代码可能会以各种方式损坏。

7. 完整示例整合

下面将前面所有的步骤整合到一个完整的脚本中，并支持设备（CPU/GPU）自动切换。

import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.model_selection import train_test_split

# 1. 确定当前可用的设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# 2. 创建数据集
torch.manual_seed(42)
x = torch.rand(50, 1).to(device)
y = x * 3.0 + 2.0 + torch.randn(50, 1).to(device) * 0.1

# 3. 划分数据集
x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=42)

# 4. 建立模型
class LinearModel(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.linear = nn.Linear(input_dim, output_dim)
    def forward(self, x):
        return self.linear(x)

model = LinearModel(input_dim=1, output_dim=1).to(device)

# 5. 定义损失函数和优化器
loss_fn = nn.L1Loss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 6. 训练循环
epochs = 1000
for epoch in range(epochs):
    model.train()
    y_pred = model(x_train)
    loss = loss_fn(y_pred, y_train)
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if (epoch + 1) % 100 == 0:
        print(f"Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}")

# 7. 测试与预测
model.eval()
with torch.inference_mode():
    test_pred = model(x_test)
    print(f"Test Loss: {loss_fn(test_pred, y_test).item():.4f}")

# 8. 保存模型
torch.save(model.state_dict(), 'final_model.pth')
print("Training and Saving Complete.")

当我们将训练 epochs 从 100 增加到 1000 之后，模型学到的参数很逼近真实参数值了。现在红色点几乎和蓝色点重叠了，非常好！

总结

本文详细拆解了 PyTorch 模型训练的标准工作流。从数据预处理、模型定义、损失计算、优化器配置到训练循环的实现，再到最终的模型持久化。掌握这一流程是深入深度学习开发的基础。开发者可根据具体任务调整网络结构、损失函数及超参数，以适应更复杂的场景。