PyTorch 基本工作流

2025-02-08

Reference: PyTorch Workflow Fundamentals

该页面由 Jupyter Notebook 生成，原文件于 Github

In [1]:

# 先导入包
import torch
from torch import nn
import matplotlib.pyplot as plt

torch.__version__

Out[1]:

'2.5.1+cu124'

准备数据集

数据可以是很多东西，如一个表格、任何类型的图像、视频、歌曲或播客等音频文件，蛋白质结构，文本等等。

机器学习是一个由两部分组成：

把你的数据转换成数字表示。
选择或构建一个模型来尽可能地学习数据的表征。

获得数据之后，需要将数据划分为训练集、验证集和测试集。

类型	目的	占比	使用情况
训练集	模型从这些数据里面学习（比如学习的课程材料）	~60~80%	必须有
验证集	模型会根据这些数据进行调整（比如期末考试前的练习）	~10~20%	不必有
测试集	模型根据这些数据进行评估，以测试它学到了什么（比如期末考试）	~10~20%	必须有

In [2]:

# 创建一个 y = weight * x + bias 的数据集
weight = 0.7
bias = 0.3

X = torch.arange(0, 1, 0.02).unsqueeze(dim=1)
y = weight * X + bias

# 划分训练集和测试集
train_split = int(0.8 * len(X))
X_train, y_train = X[:train_split], y[:train_split]
X_test, y_test = X[train_split:], y[train_split:]

def plot_predictions(train_data=X_train,
                     train_labels=y_train,
                     test_data=X_test,
                     test_labels=y_test,
                     predictions=None):
    plt.figure(figsize=(5, 3))
    plt.scatter(train_data, train_labels, c="b", s=4, label="Training data")
    plt.scatter(test_data, test_labels, c="g", s=4, label="Testing data")

    if predictions is not None:
        plt.scatter(test_data, predictions, c="r", s=4, label="Prediction")

    plt.legend(prop={"size": 8})

plot_predictions()

No description has been provided for this image

构建模型

PyTorch 模型构建要点

PyTorch 有四个基本模块，可以用它来创建神经网络：

模块	作用
`torch.nn`	包含计算图的所有构建块
`torch.nn.Parameter`	存储可用于 `nn.Module` 的张量。如果 `requires_grad=True` 则自动计算梯度（用于通过梯度下降更新模型参数），这通常被称为“autograd”
`torch.nn.Module`	所有神经网络模块的基类，神经网络的所有构建块都是子类。在PyTorch中构建一个神经网络，模型应该继承 `nn.Module`，需要实现 `forward()` 方法
`torch.optim`	包含各种优化算法（这些算法告诉存储在 `nn.Parameter` 中的模型参数。如何最好地改变，以改善梯度下降，从而减少损失)
`def forward()`	所有的 `nn.Module` 子类都需要一个 `forward()` 方法，定义传递给特定 `nn.Module` 的数据进行的计算（例如线性回归公式）

简而言之：

nn.Module 包括大的构建块，如神经网络中的层；
nn.Parameter 包括小的参数，比如权重和偏置，众多参数构成 nn.Module；
forward() 定义了在 nn.Module 中对输入的计算；
torch.optim 包含如何改进 nn.Parameter中的参数的算法，以更好地表征数据

一个简单的神经网络例子：

In [3]:

class LinearRegressionModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.weights = nn.Parameter(torch.randn(1, dtype=torch.float),
                                    requires_grad=True)
        self.bias = nn.Parameter(torch.randn(1, dtype=torch.float),
                                 requires_grad=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.weights * x + self.bias

检视模型中的内容

使用 model.parameters() 检查参数：

In [4]:

torch.manual_seed(42)

model_0 = LinearRegressionModel()

list(model_0.parameters())

Out[4]:

[Parameter containing:
 tensor([0.3367], requires_grad=True),
 Parameter containing:
 tensor([0.1288], requires_grad=True)]

使用 model.state_dict() 获得模型状态（包含什么）：

In [5]:

model_0.state_dict()

Out[5]:

OrderedDict([('weights', tensor([0.3367])), ('bias', tensor([0.1288]))])

使用 `torch.inference_mode()` 进行预测

torch.inference_mode() 关闭了一些东西，比如梯度跟踪（训练所必需的，但不是推理所必需的），所以 forward 传递更快。

In [6]:

with torch.inference_mode():
    y_preds = model_0(X_test)

plot_predictions(predictions=y_preds)

由图可知模型预测距离真实值仍有一段距离，所以需要训练模型以达到更好的效果。

训练模型

建立损失函数和优化器

为了让模型自己更新参数，需要添加损失函数和优化器。

功能	作用	PyTorch的形式	常用值
损失函数	测量模型的预测（ `y_preds` ）与真值标签（ `y_test` ）相比的错误程度。越低越好	`torch.nn` 中有很多内置的损失函数	回归问题的平均绝对误差（MAE，`torch.nn.L1Loss()`）；二元分类问题的二元交叉熵（ `torch.nn.BCELoss()` ）
优化器	告诉模型如何更新其内部参数以最大程度地降低损失	`torch.optim` 中的各种优化函数实现	随机梯度下降（ `torch.optim.SGD()` ）；Adam 优化器（ `torch.optim.Adam()` ）

为模型添加损失函数和优化器：

训练循环

训练过程有以下步骤：

序号	步骤	做法	代码
1	前向传播	模型将所有训练数据遍历一次，执行其 `forward()` 函数计算	`model(x_train)`
2	计算损失	将模型的输出（预测）与真实数据进行评估，计算损失值	`loss=loss_fn(y_pred, y_train)`
3	梯度归零	优化器的梯度被设置为零（默认情况下它们是累积的），因此它们可以为特定的训练步骤重新计算	`optimizer.zero_grad()`
4	反向传播损失	计算每个要更新的模型参数的损失梯度（每个参数 `requires_grad=True`）	`loss.backward()`
5	更新参数	使用 `requires_grad=True` 更新损耗梯度的参数	`optimizer.step()`

测试循环

测试过程有以下步骤：

序号	步骤	做法	代码
1	前向传播	模型将所有训练数据遍历一次，执行其 `forward()` 函数计算	`model(x_train)`
2	计算损失	将模型的输出（预测）与真实数据进行评估，计算损失值	`loss=loss_fn(y_pred, y_train)`
3	计算评估指标（可选）	计算其他评估指标，例如测试集上的准确性	自定义函数

训练并测试代码

In [7]:

torch.manual_seed(42)

device = "cuda" if torch.cuda.is_available() else "cpu"
X_train = X_train.to(device)
y_train = y_train.to(device)
X_test = X_test.to(device)
y_test = y_test.to(device)
model_0 = model_0.to(device)

epochs = 200

train_loss_values = []
test_loss_values = []
epoch_cnt = []

loss_fn = nn.L1Loss() # 损失函数
optimizer = torch.optim.SGD(params=model_0.parameters(),
                           lr=0.01) # 优化器

for epoch in range(epochs):
    model_0.train() # 设置为训练模式

    # 1. 前向传播
    y_pred = model_0(X_train)
    # 2. 计算损失
    loss = loss_fn(y_pred, y_train)

    # 3. 梯度归零、反向传播、更新参数
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    model_0.eval() # 设置为评估模式
    with torch.inference_mode():
        test_pred = model_0(X_test)
        test_loss = loss_fn(test_pred, y_test.type(torch.float))

        if epoch % 20 == 0:
            epoch_cnt.append(epoch)
            train_loss_values.append(loss.cpu().detach().numpy())
            test_loss_values.append(test_loss.cpu().detach().numpy())
            print(f"{epoch}/{epochs} | MAE 训练损失: {loss} | MAE 测试损失: {test_loss}")

0/200 | MAE 训练损失: 0.31288135051727295 | MAE 测试损失: 0.48106518387794495
20/200 | MAE 训练损失: 0.08908725529909134 | MAE 测试损失: 0.21729658544063568
40/200 | MAE 训练损失: 0.04543796926736832 | MAE 测试损失: 0.11360953003168106
60/200 | MAE 训练损失: 0.03818932920694351 | MAE 测试损失: 0.08886633068323135
80/200 | MAE 训练损失: 0.03132382780313492 | MAE 测试损失: 0.07232122868299484
100/200 | MAE 训练损失: 0.024458957836031914 | MAE 测试损失: 0.05646304413676262
120/200 | MAE 训练损失: 0.01758546754717827 | MAE 测试损失: 0.04060482606291771
140/200 | MAE 训练损失: 0.010716589167714119 | MAE 测试损失: 0.024059748277068138
160/200 | MAE 训练损失: 0.003851776709780097 | MAE 测试损失: 0.008201557211577892
180/200 | MAE 训练损失: 0.008932482451200485 | MAE 测试损失: 0.005023092031478882

损失随着时间的推移而下降，绘图：

In [8]:

plt.figure(figsize=(5, 3))
plt.plot(epoch_cnt, train_loss_values, label="Train Loss")
plt.plot(epoch_cnt, test_loss_values, label="Test Loss")
plt.title("Loss Curves")
plt.ylabel("Loss")
plt.xlabel("Epochs")
plt.legend()

Out[8]:

<matplotlib.legend.Legend at 0x19c1693b990>

输出一下训练得到的 weight 和 bias：

In [9]:

print(model_0.state_dict())
print(f"{weight}, {bias}")

OrderedDict([('weights', tensor([0.6990], device='cuda:0')), ('bias', tensor([0.3093], device='cuda:0'))])
0.7, 0.3

由此可知训练得到的参数与实际的参数已经相差很小了。

使用训练后的PyTorch模型进行推理

在使用 PyTorch 模型进行预测（也称为执行推理）时，需要记住三件事：

将模型设置为评估模式（ model.eval() ）。
使用推理模式上下文管理器（使用 torch.inference_mode() ）进行预测。
所有的预测都应该在同一设备上进行（仅在GPU上的数据和模型或仅在CPU上的数据和模型）。

前两项确保关闭 PyTorch 在训练期间在幕后使用的所有有用的计算和设置，但这些计算和设置对于推理是不必要的（这导致更快的计算）。第三个确保你不会遇到跨设备错误。

最后查看整体的分布：

In [10]:

model_0.eval()

with torch.inference_mode():
  y_preds = model_0(X_test)

plot_predictions(predictions=y_preds.cpu())

保存和加载模型

比如在服务器上训练模型后，需要转移到本地或其他地方进行使用。

方法	作用
`torch.save`	使用 Python 的 `pickle` 实用程序将序列化的对象保存到磁盘。模型、张量和其他各种Python对象（如字典）都可以使用 `torch.save` 保存
`torch.load`	使用 `pickle` 的 unpickling 特性来反序列化并将文件（如模型，张量或字典）加载到内存中。也可以设置加载对象到哪个设备（CPU， GPU等）
`torch.nn.Module.load_state_dict`	使用已保存的 `state_dict()` 对象加载模型的参数字典（ `model.state_dict()` ）

保存模型的参数字典

In [11]:

torch.save(obj=model_0.state_dict(),
          f="model_0.pth")

加载参数字典到模型中

In [12]:

loaded_model_0 = LinearRegressionModel()

loaded_model_0.load_state_dict(torch.load(f="model_0.pth", weights_only=True))

# 测试看看预测结果是否相等
loaded_model_0.to(device)
loaded_model_0.eval()

with torch.inference_mode():
    loaded_model_preds = loaded_model_0(X_test)

loaded_model_preds == y_preds

Out[12]:

tensor([[True],
        [True],
        [True],
        [True],
        [True],
        [True],
        [True],
        [True],
        [True],
        [True]], device='cuda:0')

Tensor 张量

PyTorch 分类模型

Style Transfer By Fingsinz