PyTorch 工作流程

关注

原创最新推荐文章于 2026-08-01 08:03:26 发布 1.1k 阅读 AI 写同款· GEO 优化 ›

【TorchScript】PyTorch模型转换Libtorch模型任务简介：使用PyTorch训练的模型只能在Python环境中使用，在自动驾驶场景中，模型推理过程通常是在硬件设备上进行。TorchScript可以将PyTorch训练的模型转换为C++环境支持的模型，推理速度比Python环境更快。本文对整体转换流程做一个简单的记录，后续需要补充TorchScript的支持的各种语法规则以及注意点。............... 阅读详情

前言

机器学习和深度学习的本质是从过去获取一些数据，构建算法（如神经网络）来发现其中的模式，并使用发现的模式来预测未来。有很多方法可以做到这一点，并且一直在发现许多新方法。

本节内容

我们将介绍一个标准的 PyTorch 工作流程（它可以根据需要进行更改）。

现在，我们将使用此工作流程来预测一条简单的直线。

具体如下:

名称	内容
1. 准备数据	数据几乎可以是任何东西，在本文中，我们将创建一条简单的直线。
2. 建立模型	创建一个模型来学习数据中的模式，将选择损失函数、优化器来构建训练过程。
3. 将模型拟合到数据（训练）	已经有了数据和模型，现在让模型尝试在（训练）数据中找到模式。
4. 做出预测和评估模型（推理）	模型在数据中找到了模式，将其预测与实际（测试）数据进行比较。
5. 保存和加载模型	当想在其他地方使用模型，或者稍后再回来使用它时需要保存和加载模型。
6. 完整过程	将以上所有内容结合起来。

本文来源

GitHub 上提供了本节的英文版：https://github.com/mrdbourke/pytorch-deep-learning

如果遇到麻烦，可以在讨论页面上提问：https://github.com/mrdbourke/pytorch-deep-learning/discussions

还有 PyTorch 开发者论坛：https://discuss.pytorch.org/)

Kaggle 平台的练习地址

在原文的基础上汉化，复制链接，点击即可练习，https://www.kaggle.com/zymzym/pytorch-workflow-01

1. 数据（准备和加载）

导入此模块所需的内容：torch, torch.nn (nn 代表神经网络，这个包用在 PyTorch 中创建神经网络) 和 matplotlib。

import torch
from torch import nn 
import matplotlib.pyplot as plt

# 检查 PyTorch 版本
torch.__version__

'1.11.0'

机器学习中的“数据”几乎可以是能想象到的任何东西，例如数字表格（如大型 Excel 电子表格）、任何类型的图像、视频（YouTube 有大量数据！）、歌曲或播客等音频文件、蛋白质结构、文本等。

机器学习可以看成两部分：

将数据（无论是什么）转换为数字（一种表示形式）。

选择或构建模型来学习表示。

但是，如果没有数据怎么办？就是现在我们所处的状态，但可以创造一些，将数据创建为一条直线。

将使用线性回归来创建数据，线性回归里的已知参数，是模型需要学习到的，将使用 PyTorch 构建模型，使用梯度下降来估计这些参数。

如果上面的术语现在没有多大意义，请不要担心，我们会看到它们的实际应用。

# 创建*已知*参数
weight = 0.7
bias = 0.3

# 创建数据
start = 0
end = 1
step = 0.02
X = torch.arange(start, end, step).unsqueeze(dim=1)
y = weight * X + bias

X[:10], y[:10]

(tensor([[0.0000],
         [0.0200],
         [0.0400],
         [0.0600],
         [0.0800],
         [0.1000],
         [0.1200],
         [0.1400],
         [0.1600],
         [0.1800]]),
 tensor([[0.3000],
         [0.3140],
         [0.3280],
         [0.3420],
         [0.3560],
         [0.3700],
         [0.3840],
         [0.3980],
         [0.4120],
         [0.4260]]))

现在将着手构建一个可以学习 X (features 称为特征) 和 y (labels 称为标签) 之间关系的模型。

将数据拆分为训练集和测试集

我们有一些数据。

但在构建模型之前，需要将其拆分。

机器学习项目中最重要的步骤之一是创建训练和测试集（以及在需要时创建验证集）。

数据集的每个拆分都有一个特定的目的：

名称	目的	总数据量占比	使用频率
训练集	该模型从这些数据中学习（比如在学期中学习的课程材料）。	~60-80%	总是
验证集	该模型根据这些数据进行调整（就像在期末考试前参加的模拟考试）。	~10-20%	经常
测试集	该模型根据这些数据进行评估，以测试它学到了什么（比如你在学期末参加的期末考试）。	~10-20%	总是

现在，我们将只使用训练和测试集，这意味着将有一个数据集供模型学习和评估。

我们可以通过拆分张量 X 和 y 来创建它们。

注意: 在处理真实世界的数据时，此步骤通常在项目开始时就完成（测试集应始终与所有其他数据分开）。我们希望模型在训练数据上学习，然后在测试数据上对其进行评估，了解它对未见示例的泛化能力（generalizes ）。

# 创建训练/测试拆分，80% 的数据用于训练集，20% 用于测试
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:]

len(X_train), len(y_train), len(X_test), len(y_test)

(40, 40, 10, 10)

我们有 40 个训练样本 (X_train & y_train) 和 10 个测试样本 (X_test & y_test)。

模型将尝试学习 X_train 和 y_train 之间的关系，然后在 X_test 和 y_test上测试模型学到的内容。

但现在数据只是页面上的数字。

让我们创建一个函数来可视化它。

def plot_predictions(train_data=X_train, 
                     train_labels=y_train, 
                     test_data=X_test, 
                     test_labels=y_test, 
                     predictions=None):
  """
     绘制训练数据、测试数据。
  """
  plt.figure(figsize=(10, 7))

  # 用蓝色绘制训练数据
  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="Predictions")

  # 显示图例
  plt.legend(prop={"size": 14});

plot_predictions();

现在，我们的数据不再只是页面上的数字，而是一条直线。

2. 建立模型

现在已经有了一些数据，让我们建立一个模型来使用蓝点来预测绿点。

我们先编写代码，然后再解释它。

使用 PyTorch 手动写一个标准线性回归模型。

# 创建一个线性回归模型类
class LinearRegressionModel(nn.Module): 
    # PyTorch中几乎所有的东西都是一个nn.Module （把它想象成神经网络乐高积木块）
    def __init__(self):
        super().__init__() 
        self.weights = nn.Parameter(torch.randn(1, # 从随机权重开始（这将随着模型的学习而调整）
                                                dtype=torch.float), # PyTorch 默认 float32
                                   requires_grad=True) # 可以用梯度下降更新这个值

        self.bias = nn.Parameter(torch.randn(1, # 从随机偏差开始（这将随着模型的学习而调整）
                                            dtype=torch.float), # PyTorch 默认 float32 
                                requires_grad=True) # 可以用梯度下降更新这个值

    # 前向定义模型中的计算                                      
    def forward(self, x: torch.Tensor) -> torch.Tensor: 
        # “x”是输入数据（例如训练/测试特征）
        return self.weights * x + self.bias # 这是线性回归公式 (y = m*x + b)

接下来，我们细致了解一下以上代码。

资源: 我们将使用 Python 类来创建神经网络。如果不熟悉 Python 类，建议多读几遍 Real Python 的面向对象编程指南。 [Real Python's Object Orientating programming in Python 3 guide]:( https://realpython.com/python3-object-oriented-programming/)

PyTorch 模型构建要领

PyTorch 有四个基本模块，可以使用它们来创建任何类型的神经网络。

它们是 torch.nn、torch.optim、torch.utils.data.Dataset 和 torch.utils.data.DataLoader。

现在，我们将重点关注前两个，稍后再讨论其他两个。

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

如果以上听起来很复杂，那么可以这样想，PyTorch 神经网络中几乎所有的东西都来自 torch.nn。

nn.Module 包含较大的构建块（层）

nn.Parameter 包含较小的参数，如权重和偏差（将它们放在一起构成 nn.Module

forward() 告诉较大的块如何对 nn.Module 内的输入（充满数据的张量）进行计算

torch.optim 包含有关如何改进内部参数 nn.Parameter 以更好地表示输入数据的优化方法

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

资源: 在 PyTorch Cheat Sheet 中查看更多这些基本模块及其用例。[PyTorch Cheat Sheet]:( https://pytorch.org/tutorials/beginner/ptcheat.html).

检查 PyTorch 模型的内容

让我们用创建的类创建一个模型实例，并使用检查它的参数 .parameters()。

# 设置手动种子，因为 nn.Parameter 是随机初始化的
torch.manual_seed(42)

# 创建模型的实例（这是包含 nn.Parameter(s) 的 nn.Module 的子类）
model_0 = LinearRegressionModel()

# 检查 nn.Module 子类中的 nn.Parameter(s) 
list(model_0.parameters())

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

我们还可以使用 .state_dict() 获取模型的状态（模型包含的内容）。

# 列出命名参数
model_0.state_dict()

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

请注意 model_0.state_dict() 中的 weights 和 bias 是如何作为随机浮点张量出现的？

这是因为在上面使用 torch.randn()。

本质上，我们希望从随机参数开始，让模型将它们更新为最适合数据的参数 (在创建直线数据时设置的 weight 和 bias) 。

因为模型从随机值开始，所以现在它的预测能力很差。

使用 torch.inference_mode() 进行预测

要检查训练的好坏，可以将测试数据 X_test 传递给模型，看看它的预测有多接近 y_test。

当我们将数据传递给模型时，它将通过模型的 forward() 方法，使用我们定义的计算产生结果。

# 模型预测
with torch.inference_mode(): 
    y_preds = model_0(X_test)

# 注意：在旧的 PyTorch 代码中你可能还会看到 torch.no_grad() 
# with torch.no_grad():
#   y_preds = model_0(X_test)

可能注意到使用 torch.inference_mode() 作为上下文管理器 context manager (也就是 with torch.inference_mode(): ) 来进行预测。

顾名思义，torch.inference_mode() 在模型进行推理（进行预测）时使用。

torch.inference_mode() 关闭一堆东西（比如梯度跟踪，这是训练所必需的而不是推理所必需的）以使前向传递（forward-passes）（通过该 forward() 方法的数据）更快。

注意: 在较旧的 PyTorch 代码中，可能还会看到 torch.no_grad() 用于推理的情况。虽然 torch.inference_mode() 和 torch.no_grad() 做类似的事情，但 torch.inference_mode() 更新可能更快并且更受欢迎。

做了一些预测，看看它们是什么样的。

# 检查预测
print(f"测试样本数: {len(X_test)}") 
print(f"预测数量: {len(y_preds)}")
print(f"预测值:\n{y_preds}")

测试样本数: 10
预测数量: 10
预测值:
tensor([[0.3982],
        [0.4049],
        [0.4116],
        [0.4184],
        [0.4251],
        [0.4318],
        [0.4386],
        [0.4453],
        [0.4520],
        [0.4588]])

请注意每个测试样本有一个预测值。

我们使用的数据类型，一条直线，一个X 值映射到一个 y 值。

然而，机器学习模型非常灵活。您可以将 100 个 X 值映射到一个、两个、三个或 10 个 y 值。这完全取决于解决的问题。

我们的预测仍然是页面上的数字，用上面创建的函数 plot_predictions() 将它们可视化。

plot_predictions(predictions=y_preds)

y_test - y_preds

tensor([[0.4618],
        [0.4691],
        [0.4764],
        [0.4836],
        [0.4909],
        [0.4982],
        [0.5054],
        [0.5127],
        [0.5200],
        [0.5272]])

哇！这些预测看起来很糟糕......

不过，还记得我们的模型只是使用随机参数值进行预测时。

它甚至没有拟合蓝点来尝试预测绿点。

是时候改变它了。

3. 训练模型

现在我们的模型正在使用随机参数进行预测，它基本上是猜测（随机）。

为了解决这个问题，我们可以更新它的内部参数 (将 *参数 parameters * 称为模式 patterns)，使用 nn.Parameter() 和 torch.randn() 随机初始化 weights 和 bias ，能更好的拟合数据。

我们刚开始设置了weight=0.7 和 bias=0.3 生成了数据，但很多时候不知道模型的理想参数是什么。若是编写代码以查看模型是否可以自己弄清楚它们会更有趣。

在 PyTorch 中创建损失函数和优化器

为了让模型自行更新其参数，需要添加更多内容。

一个损失函数 loss function 和一个优化器 optimizer。

它们具体是:

函数	作用	它在 PyTorch 中的什么位置？	举例
Loss function	衡量模型预测 (e.g. y_preds) 与真实标签 (e.g. y_test)的不同程度，越低越好。	torch.nn	回归问题的平均绝对误差 (MAEMean absolute error (MAE)）(torch.nn.L1Loss())。二元分类问题的二元交叉熵（Binary cross entropy） (torch.nn.BCELoss()).
Optimizer	模型如何更新其内部参数以最好地降低损失。	torch.optim	随机梯度下降（Stochastic gradient descent） (torch.optim.SGD()). Adam optimizer (torch.optim.Adam()).

让我们创建一个损失函数和一个优化器，改进模型效果。

一般根据正在处理的问题类型，选择合适的损失函数和优化器。

当然，有一些众所周知的常用值，例如 SGD（随机梯度下降）或 Adam 优化器。以及用于回归问题（预测数字）的 MAE（平均绝对误差）或用于分类问题（预测一件事或另一件事）的二元交叉熵损失函数。

对于我们的问题，因为要预测一个数字，所以使用 torch.nn.L1Loss()) 作为损失函数。

平均绝对误差（ (MAE, in PyTorch: torch.nn.L1Loss) 测量两点（预测和标签）之间的绝对差异，然后取所有示例的平均值。

我们将使用 SGD，torch.optim.SGD(params, lr)，其中:

params 是要优化的模型参数（例如之前随机设置的 weights 和 bias ）。

lr 是学习率 learning rate 通过优化器更新参数，数值的高意味着优化器将尝试更大的更新（这些有时可能太大，优化器将无法工作），数值设置的低意味着优化器将尝试更小的更新（这些有时可能太小，优化器将花费太长时间才能找到理想值）。学习率被认为是一个超参数 hyperparameter 因为它是由机器学习工程师设置的）。学习率的常见起始值是 0.01, 0.001, 0.0001，但是，这些值也可以随时间调整 learning rate scheduling).

哇，这么多，让我们在代码中看看吧。

# 创建损失函数
loss_fn = nn.L1Loss() # MAE 损失与 L1Loss 相同

# 创建优化器

optimizer = torch.optim.SGD(params=model_0.parameters(), # 要优化的模型参数
                            lr=0.01) 
# 学习率是优化器应该在每一步改变多少参数， 设置高会不太稳定，设置低了可能需要很长时间

在 PyTorch 中创建循环优化

哇哦！现在我们有了一个损失函数和一个优化器，现在是时候创建一个训练循环 (或测试循环).

在训练循环中模型遍历训练数据并学习特征 features 和标签labels之间的关系。

在测试循环中模型遍历测试数据并评估模型在训练数据上学习到的模式有多好（模型在训练期间永远不会看到测试数据）。

在每一个“循环”中，我们希望我模型查看每个数据集中的每个样本。

Daniel Bourke 将循环的过程写了一首歌，PyTorch 优化循环歌曲（非官方）如下，以一种有趣的方式来记住 PyTorch 训练（和测试）循环中的步骤。

PyTorch 训练循环

对于训练循环，我们将构建以下步骤：

数字	步骤名称	它有什么作用?	代码示例
1	前向传播	该模型一次遍历所有训练数据，执行其forward() 函数计算。	model(x_train)
2	计算损失	将模型的输出（预测）与真实标签进行比较并进行评估以查看它们的错误程度。	loss = loss_fn(y_pred, y_train)
3	梯度清零	优化器梯度设置为零（默认情况下它们是累积的），因此可以为特定的训练步骤重新计算它们。	optimizer.zero_grad()
4	对损失进行反向传播	更新每个模型参数的损失梯度 (每个参数带有requires_grad=True)。这被称为反向传播 backpropagation。	loss.backward()
5	更新优化器的参数 (使用梯度下降 gradient descent)	设置 requires_grad=True，更新损失梯度的参数以改进它们。	optimizer.step()

注意: 以上只是如何对步骤进行排序的一个示例。凭借经验，读者会发现制作 PyTorch 训练循环非常灵活。
以上是一个很好的默认顺序，但读者可能会看到略有不同的顺序。一些经验法则：
在执行反向传播 ( loss.backward()) 之前计算损失函数 ( loss = ...)。
在更新优化器 ( optimizer.step()) 之前将梯度清零 ( optimizer.zero_grad())。
在更新优化器 ( optimizer.step()) 之前执行反向传播 ( loss.backward())。

PyTorch 测试循环

测试循环（评估模型），典型的步骤包括：

数字	步骤名称	它有什么作用?	代码示例
1	前向传播	该模型一次遍历所有测试数据，执行其 forward() 函数计算。	model(x_test)
2	计算损失	将模型的输出（预测）与真实标签进行比较并进行评估以查看它们的错误程度。	loss = loss_fn(y_pred, y_test)
3	计算评估指标（可选）	除了损失值，可能还想计算其他评估指标，例如测试集的准确性。	选择合适问题的指标，可自定义函数

请注意，测试循环不包含执行反向传播(loss.backward()) 和更新优化器 (optimizer.step())，这是因为模型中的参数在测试期间没有被更改，它们已经被计算出来。对于测试，只是通过模型的前向传递计算输出值。

将以上所有内容放在一起并训练模型 100 个 epochs （前向传递数据），在每 10 个 epochs 进行评估。

torch.manual_seed(42)

# 设置 epoch 的数量（模型训练数据的次数）
epochs = 100

# 创建空损失列表以跟踪值
train_loss_values = []
test_loss_values = []
epoch_count = []

for epoch in range(epochs):
    ### Training

    # 将模型置于训练模式（这是模型的默认状态）
    model_0.train()

    # 1. 使用 forward() 方法前向传播训练数据 
    y_pred = model_0(X_train)
    # print(y_pred)

    # 2. 计算损失（模型预测与真实标签有何不同）
    loss = loss_fn(y_pred, y_train)

    # 3. 优化器梯度清零
    optimizer.zero_grad()

    # 4. 对损失进行反向传播
    loss.backward()

    # 5. 执行一次优化步骤，通过梯度下降方法来更新参数的值
    optimizer.step()

    ### Testing

    # 将模型置于评估模式
    model_0.eval()

    with torch.inference_mode():
      # 1. 前向传递测试数据
      test_pred = model_0(X_test)

      # 2. 计算测试数据的损失
# 预测来自 torch.float 数据类型，所以需要用相同类型的张量进行比较
      test_loss = loss_fn(test_pred, y_test.type(torch.float)) 
      # 打印出发生了什么
      if epoch % 10 == 0:
            epoch_count.append(epoch)
            train_loss_values.append(loss.detach().numpy())
            test_loss_values.append(test_loss.detach().numpy())
            print(f"Epoch: {epoch} | MAE Train Loss: {loss} | MAE Test Loss: {test_loss} ")

Epoch: 0 | MAE Train Loss: 0.31288138031959534 | MAE Test Loss: 0.48106518387794495 
Epoch: 10 | MAE Train Loss: 0.1976713240146637 | MAE Test Loss: 0.3463551998138428 
Epoch: 20 | MAE Train Loss: 0.08908725529909134 | MAE Test Loss: 0.2172965705394745 
Epoch: 30 | MAE Train Loss: 0.053148530423641205 | MAE Test Loss: 0.14464019238948822 
Epoch: 40 | MAE Train Loss: 0.04543796926736832 | MAE Test Loss: 0.11360955238342285 
Epoch: 50 | MAE Train Loss: 0.04167862981557846 | MAE Test Loss: 0.09919948875904083 
Epoch: 60 | MAE Train Loss: 0.03818932920694351 | MAE Test Loss: 0.08886633813381195 
Epoch: 70 | MAE Train Loss: 0.03476089984178543 | MAE Test Loss: 0.0805937796831131 
Epoch: 80 | MAE Train Loss: 0.03132382780313492 | MAE Test Loss: 0.07232122868299484 
Epoch: 90 | MAE Train Loss: 0.02788739837706089 | MAE Test Loss: 0.06473556160926819

以上训练结果看起来损失在每个时期都在下降，让我们绘制它。

# 绘制损失曲线
plt.plot(epoch_count, train_loss_values, label="Train loss")
plt.plot(epoch_count, test_loss_values, label="Test loss")
plt.title("Training and test loss curves")
plt.ylabel("Loss")
plt.xlabel("Epochs")
plt.legend();

损失曲线 loss curves 显示损失随时间下降。请记住，损失是衡量模型错误程度的指标，因此越低越好。

但是为什么损失下降了呢？

由于损失函数和优化器，模型的内部参数 (weights 和 bias) 得到了更新，以更好地反映数据中的潜在模式。

让我们检查一下模型 .state_dict() 看看模型与权重和偏差设置的原始值有多接近。

# 找到模型的学习参数
print("该模型学习了以下权重和偏差值：")
print(model_0.state_dict())
print("\n权重和偏差的原始值为")
print(f"weights: {weight}, bias: {bias}")

该模型学习了以下权重和偏差值：
OrderedDict([('weights', tensor([0.5784])), ('bias', tensor([0.3513]))])

权重和偏差的原始值为
weights: 0.7, bias: 0.3

哇！多么酷啊？

我们的模型非常接近初始值 weight 和 bias (如果我们训练它更长时间，它可能会更接近)。

它可能永远不会完美地猜出它们（尤其是在使用更复杂的数据集时），但这没关系，通常你可以用一个近似值来做非常酷的事情。

这就是机器学习和深度学习的全部思想，有一些理想值描述了数据，可以训练一个模型以编程方式计算出它们，而不是手动计算出它们。

4. 使用训练好的 PyTorch 模型进行预测（推理）

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

将模型设置为评估模式 (model.eval())。

使用推理模式上下文管理器进行预测 (with torch.inference_mode(): ...)。

所有预测都应使用同一设备上进行 (例如，仅 GPU 上的数据和模型或仅 CPU 上的数据和模型)。

前两项确保 PyTorch 关闭对于推理不是必需的计算和设置（对于训练必需），这会加快计算速度，第三个确保不会遇到跨设备错误。

# 1. 将模型设置为评估模式
model_0.eval()

# 2. 设置推理模式上下文管理器
with torch.inference_mode():
  # 3. 确保在同一设备上使用模型和数据进行计算
  # 在本例中，我们还没有设置与设备无关的代码，
  # 所以我们的数据和模型默认情况下在 CPU 上。 
  # model_0.to(device)
  # X_test = X_test.to(device)
  y_preds = model_0(X_test)
y_preds

tensor([[0.8141],
        [0.8256],
        [0.8372],
        [0.8488],
        [0.8603],
        [0.8719],
        [0.8835],
        [0.8950],
        [0.9066],
        [0.9182]])

好的！模型做出了一些预测，现在它们看起来如何？

plot_predictions(predictions=y_preds)

哇哦！那些红点看起来比以前更近了！

让我们开始在 PyTorch 中保存并重新加载模型。

5. 保存和加载 PyTorch 模型

如果读者已经训练了 PyTorch 模型，很可能想要保存它并将其导出到某个地方。

比如，读者可能会在 Kaggle\Google Colab 或带有 GPU 的本地机器上对其进行训练，将其导出到其他人可以使用的某种应用程序中，或者可能想保存在模型上的进度并稍后返回并加载它。

对于在 PyTorch 中保存和加载模型，应该了解三种主要方法（以下所有方法均摘自 PyTorch 保存和加载模型的指南）：

PyTorch 方法	它有什么作用？
torch.save	使用 Python's pickle 的实用程序将序列化对象保存到磁盘。模型、张量和各种其他 Python 对象（如字典）可以使用 torch.save。
torch.load	使用 pickle 的 unpickling 功能反序列化并加载 pickled Python 对象文件（如模型、张量或字典）到内存中。还可以设置将对象加载到哪个设备（CPU、GPU 等）。
torch.nn.Module.load_state_dict	使用state_dict() 对象加载模型的参数字典 (model.state_dict())

保存 PyTorch 模型的state_dict()

保存和加载推理模型（进行预测）的推荐方法是 state_dict().

步骤如下:

使用 Python 的 pathlib 模块创建一个目录（命名为 models）来保存模型.

创建一个文件路径来保存模型。

使用 torch.save(obj, f) 调用模型，其中 obj 是保存模型的状态（如只保存模型的参数）， f 模型文件保存的地址。

注意: PyTorch 保存的模型或对象通常以 .pt 或 .pth 结尾，例如 saved_model_01.pth。

from pathlib import Path

# 1. 创建模型目录 
MODEL_PATH = Path("models")
MODEL_PATH.mkdir(parents=True, exist_ok=True)

# 2. 创建模型保存路径 
MODEL_NAME = "01_pytorch_workflow_model_0.pth"
MODEL_SAVE_PATH = MODEL_PATH / MODEL_NAME

# 3. 保存模型状态 
print(f"Saving model to: {MODEL_SAVE_PATH}")
torch.save(obj=model_0.state_dict(), # state_dict() 只保存模型学习到的参数
           f=MODEL_SAVE_PATH)

Saving model to: models/01_pytorch_workflow_model_0.pth

# 检查保存的文件路径
!ls -l models/01_pytorch_workflow_model_0.pth

-rw-r--r-- 1 root root 1063 Jan 13 07:03 models/01_pytorch_workflow_model_0.pth

加载保存的 PyTorch 模型 state_dict()

由于我们现在有一个保存的模型 state_dict() ，模型地址是 models/01_pytorch_workflow_model_0.pth 。我们可以使用 torch.nn.Module.load_state_dict(torch.load(f)) 加载它，其中 f 是保存模型的地址state_dict()。

为什么 torch.load() 在 torch.nn.Module.load_state_dict()里?

因为我们只保存了模型 state_dict()，它是学习参数的字典而不是整个模型，所以我们首先必须使用torch.load() 加载state_dict()，然后将 state_dict() 传递给模型的新实例（它是 nn.Module 的子类）。

为什么不保存整个模型？

保存整个模型相比保存模型的 state_dict() 更直观，但是这种方法（保存整个模型）的缺点是序列化数据绑定到特定类以及保存模型时需使用确切目录结构，因此，你的代码在其他项目中使用时或在重构后可能会以各种方式中断。

因此，我们使用灵活的方法来保存和加载 state_dict()，这也是模型参数的字典。

让我们通过创建LinearRegressionModel() 的另一个实例来测试，它是 torch.nn.Module 的子类，具有内置方法load_state_dit()。

# 实例化模型的一个新实例（这将使用随机权重实例化）
loaded_model_0 = LinearRegressionModel()

# 加载我们保存的模型的 state_dict（这将使用训练好的权重更新模型的新实例）
loaded_model_0.load_state_dict(torch.load(f=MODEL_SAVE_PATH))

<All keys matched successfully>

现在要测试加载的模型，让我们用它对测试数据进行推理（做出预测）。

还记得使用 PyTorch 模型执行推理的规则吗？

# 1. 将加载的模型置于评估模式
loaded_model_0.eval()

# 2. 使用推理模式上下文管理器
with torch.inference_mode():
    # 使用加载的模型对测试数据执行正向传递
    loaded_model_preds = loaded_model_0(X_test)

现在我们已经对加载的模型进行了一些预测，让我们看看它们是否与之前的预测相同。

# 将之前的模型预测与加载的模型预测进行比较（这些应该是相同的）
y_preds == loaded_model_preds

tensor([[True],
        [True],
        [True],
        [True],
        [True],
        [True],
        [True],
        [True],
        [True],
        [True]])

看起来加载的模型预测与之前的模型预测相同。这表明我们的模型正在按预期保存和加载。

注意: 还有更多方法可以保存和加载 PyTorch 模型，有关更多信息，请参阅 PyTorch 指南以保存和加载模型。 [PyTorch guide for saving and loading models]：( https://pytorch.org/tutorials/beginner/saving_loading_models.html#saving-and-loading-models) for more.

6. 把以上过程放在一起

我们将从导入所需的标准库开始。

# Import PyTorch and matplotlib
import torch
from torch import nn 
# nn contains all of PyTorch's building blocks for neural networks
import matplotlib.pyplot as plt

# Check PyTorch version
torch.__version__

'1.11.0'

现在，根据实际情况设置设备 device="cuda" 或者 device="cpu"。

# Setup device agnostic code
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

Using device: cuda

如果有 GPU，上面的内容应该已经打印出来了:

Using device: cuda

否则，将使用 CPU 进行下列计算。这对于小型数据集还好，但是对于大型数据集则需要更长的时间。

6.1 数据

让我们像以前一样创建一些数据，

# Create weight and bias
weight = 0.7
bias = 0.3

# Create range values
start = 0
end = 1
step = 0.02

# Create X and y (features and labels)
X = torch.arange(start, end, step).unsqueeze(dim=1) # without unsqueeze, errors will happen later on (shapes within linear layers)
y = weight * X + bias 
X[:10], y[:10]

(tensor([[0.0000],
         [0.0200],
         [0.0400],
         [0.0600],
         [0.0800],
         [0.1000],
         [0.1200],
         [0.1400],
         [0.1600],
         [0.1800]]),
 tensor([[0.3000],
         [0.3140],
         [0.3280],
         [0.3420],
         [0.3560],
         [0.3700],
         [0.3840],
         [0.3980],
         [0.4120],
         [0.4260]]))

现在有了一些数据，把它分成 80% 的训练数据和 20% 的测试数据。

# Split data
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:]

len(X_train), len(y_train), len(X_test), len(y_test)

(40, 40, 10, 10)

6.2 构建一个线性模型

已经得到了一些数据，现在是时候建立一个模型。

将创建与前面相同的模型样式，这次不是使用 nn.Parameter() 手动定义模型的权重和偏差参数。将使用 nn.Linear(in_features, out_features) 。

其中 in_features 表示输入数据具有的维数和out_features 表示希望输出到的维数。

在本次例子中，输入和输出数据的维数都为 1，因为每个输入 (X) 对应每个标签 (y)。

本次使用 nn.Linear 创建一个线性回归模型，之前是用 nn.Parameter 手工设计。说明 torch.nn 包含许多神经网络层。

# Subclass nn.Module to make our model
class LinearRegressionModelV2(nn.Module):
    def __init__(self):
        super().__init__()
        # Use nn.Linear() for creating the model parameters
        self.linear_layer = nn.Linear(in_features=1, 
                                      out_features=1)
    
    # Define the forward computation (input data x flows through nn.Linear())
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.linear_layer(x)

# Set the manual seed when creating the model (this isn't always need but is used for demonstrative purposes, try commenting it out and seeing what happens)
torch.manual_seed(42)
model_1 = LinearRegressionModelV2()
model_1, model_1.state_dict()

(LinearRegressionModelV2(
   (linear_layer): Linear(in_features=1, out_features=1, bias=True)
 ),
 OrderedDict([('linear_layer.weight', tensor([[0.7645]])),
              ('linear_layer.bias', tensor([0.8300]))]))

注意 model_1.state_dict() 的输出，这个 nn.Linear() 层创建了随机的 weight 和 bias 参数。

现在把模型放在 GPU 上。可以使用.to(device) 更改 PyTorch 对象所在的设备。

首先检查模型的当前设备。

# Check model device
next(model_1.parameters()).device

device(type='cpu')

这个模型默认在中央处理器上，把它改为 GPU (如果可用的话)。

# Set model to GPU if it's availalble, otherwise it'll default to CPU
model_1.to(device) # the device variable was set above to be "cuda" if available or "cpu" if not
next(model_1.parameters()).device

device(type='cuda', index=0)

由于设备不可知，设置上面的代码，不管是否有 GPU 可用，都将工作。

如果你有一个 CUDA 支持的图形处理器，你应该看到这样的输出:

device(type='cuda', index=0)

6.3 训练

建立训练和测试循环，首先需要一个损失函数和一个优化器，选择与前面相同的函数 nn.L1Loss() 和 torch.optim.SGD()。需要在优化器中传递新模型的参数 (model.parameters()) ，以便优化器在训练期间对它们进行调整。学习率设置为 0.1 。

# Create loss function
loss_fn = nn.L1Loss()

# Create optimizer
optimizer = torch.optim.SGD(params=model_1.parameters(), # optimize newly created model's parameters
                            lr=0.01)

与前面的训练循环相比，我们在这个步骤中所做的唯一不同的事情是将数据放在目标设备上。这样，如果模型在 GPU 上，数据就在 GPU 上(反之亦然)。设置 epochs=1000。

torch.manual_seed(42)

# Set the number of epochs 
epochs = 1000 

# Put data on the available device
# Without this, error will happen (not all model/data on device)
X_train = X_train.to(device)
X_test = X_test.to(device)
y_train = y_train.to(device)
y_test = y_test.to(device)

for epoch in range(epochs):
    ### Training
    model_1.train() # train mode is on by default after construction

    # 1. Forward pass
    y_pred = model_1(X_train)

    # 2. Calculate loss
    loss = loss_fn(y_pred, y_train)

    # 3. Zero grad optimizer
    optimizer.zero_grad()

    # 4. Loss backward
    loss.backward()

    # 5. Step the optimizer
    optimizer.step()

    ### Testing
    model_1.eval() # put the model in evaluation mode for testing (inference)
    # 1. Forward pass
    with torch.inference_mode():
        test_pred = model_1(X_test)
    
        # 2. Calculate the loss
        test_loss = loss_fn(test_pred, y_test)

    if epoch % 100 == 0:
        print(f"Epoch: {epoch} | Train loss: {loss} | Test loss: {test_loss}")

Epoch: 0 | Train loss: 0.5551779866218567 | Test loss: 0.5739762187004089
Epoch: 100 | Train loss: 0.006215685047209263 | Test loss: 0.014086711220443249
Epoch: 200 | Train loss: 0.0012645043898373842 | Test loss: 0.013801807537674904
Epoch: 300 | Train loss: 0.0012645043898373842 | Test loss: 0.013801807537674904
Epoch: 400 | Train loss: 0.0012645043898373842 | Test loss: 0.013801807537674904
Epoch: 500 | Train loss: 0.0012645043898373842 | Test loss: 0.013801807537674904
Epoch: 600 | Train loss: 0.0012645043898373842 | Test loss: 0.013801807537674904
Epoch: 700 | Train loss: 0.0012645043898373842 | Test loss: 0.013801807537674904
Epoch: 800 | Train loss: 0.0012645043898373842 | Test loss: 0.013801807537674904
Epoch: 900 | Train loss: 0.0012645043898373842 | Test loss: 0.013801807537674904

注意: 由于机器学习的随机性，根据模型是在 CPU 还是 GPU 上进行训练，可能会得到略有不同的结果(不同的损失和预测值)。即使在任何设备上使用相同的随机种子，也是如此。如果差异很大，你可能希望查找错误，但是，如果差异很小(理想情况下是这样) ，则可以忽略它。

让我们检查模型已经学习的参数，并将它们与原始参数进行比较。

# Find our model's learned parameters
from pprint import pprint # pprint = pretty print, see: https://docs.python.org/3/library/pprint.html 
print("The model learned the following values for weights and bias:")
pprint(model_1.state_dict())
print("\nAnd the original values for weights and bias are:")
print(f"weights: {weight}, bias: {bias}")

The model learned the following values for weights and bias:
OrderedDict([('linear_layer.weight', tensor([[0.6968]], device='cuda:0')),
             ('linear_layer.bias', tensor([0.3025], device='cuda:0'))])

And the original values for weights and bias are:
weights: 0.7, bias: 0.3

6.4 做出预测

已经有了一个训练有素的模型，让我们打开它的评估模式，并作出一些预测。

# Turn model into evaluation mode
model_1.eval()

# Make predictions on the test data
with torch.inference_mode():
    y_preds = model_1(X_test)
y_preds

tensor([[0.8600],
        [0.8739],
        [0.8878],
        [0.9018],
        [0.9157],
        [0.9296],
        [0.9436],
        [0.9575],
        [0.9714],
        [0.9854]], device='cuda:0')

如果你使用 GPU 上的数据进行预测，你可能会注意到上面的输出结果是 device='cuda:0' 。这意味着数据是在 CUDA 设备0(第一个 GPU 你的系统可以访问由于零索引) ，如果你最终使用多个 GPU 在未来，这个数字可能会更高。

现在让我们绘制模型的预测图。

# plot_predictions(predictions=y_preds) # -> won't work... data not on CPU

# Put data on the CPU and plot it
plot_predictions(predictions=y_preds.cpu())

哇！看这些红点，它们和绿点几乎完美地排成一线。我想额外的 epochs 起了作用。

6.5 保存和加载模型

模型预测很好，将它保存到文件中，以便以后使用。

from pathlib import Path

# 1. Create models directory 
MODEL_PATH = Path("models")
MODEL_PATH.mkdir(parents=True, exist_ok=True)

# 2. Create model save path 
MODEL_NAME = "01_pytorch_workflow_model_1.pth"
MODEL_SAVE_PATH = MODEL_PATH / MODEL_NAME

# 3. Save the model state dict 
print(f"Saving model to: {MODEL_SAVE_PATH}")
torch.save(obj=model_1.state_dict(), # only saving the state_dict() only saves the models learned parameters
           f=MODEL_SAVE_PATH)

Saving model to: models/01_pytorch_workflow_model_1.pth

为了确保一切正常，我们把它装回去

创建 LinearRegressionModelV2() 类的新实例

使用 torch.nn.Module.load_state_dict() 加载模型状态

将模型的新实例发送到目标设备(以确保代码与设备无关)

# Instantiate a fresh instance of LinearRegressionModelV2
loaded_model_1 = LinearRegressionModelV2()

# Load model state dict 
loaded_model_1.load_state_dict(torch.load(MODEL_SAVE_PATH))

# Put model to target device (if your data is on GPU, model will have to be on GPU to make predictions)
loaded_model_1.to(device)

print(f"Loaded model:\n{loaded_model_1}")
print(f"Model on device:\n{next(loaded_model_1.parameters()).device}")

Loaded model:
LinearRegressionModelV2(
  (linear_layer): Linear(in_features=1, out_features=1, bias=True)
)
Model on device:
cuda:0

现在我们可以评估加载的模型，看看它的预测是否与之前的预测一致。

# Evaluate loaded model
loaded_model_1.eval()
with torch.inference_mode():
    loaded_model_1_preds = loaded_model_1(X_test)
y_preds == loaded_model_1_preds

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

好了，我们已经走了很长的路。你现在已经建立和训练了两个神经网络模型！

—THE END—