备注
点击此处下载完整示例代码
TorchScript 入门介绍
创建于：2025 年 4 月 1 日 | 最后更新：2025 年 4 月 1 日 | 最后验证：2024 年 11 月 5 日
作者：James Reed (jamesreed@fbcom)，Michael Suo (suo@fbcom)，rev2
警告
TorchScript 不再处于活跃开发状态。
本教程是关于 TorchScript 的介绍，它是 PyTorch 模型（ nn.Module 的子类）的中间表示形式，可以在 C++等高性能环境中运行。
在本教程中，我们将涵盖：
使用 PyTorch 进行模型编写的入门，包括：
模块
定义
forward函数将模块组合成模块层次结构
将 PyTorch 模块转换为 TorchScript 的特定方法，我们的高性能部署运行时
追踪现有模块
使用脚本直接编译模块
如何组合两种方法
保存和加载 TorchScript 模块
希望您完成本教程后，将继续学习后续教程，该教程将带您通过一个从 C++调用 TorchScript 模型的示例。
import torch # This is all you need to use both PyTorch and TorchScript!
print(torch.__version__)
torch.manual_seed(191009) # set the seed for reproducibility
PyTorch 模型编写基础 ¶
让我们从定义一个简单的 Module 开始。 Module 是 PyTorch 中的基本组成单元。它包含：
一个构造函数，用于准备模块的调用
一组
Parameters和子Modules。这些由构造函数初始化，并在模块调用期间可以使用。这是一个函数。这是当模块被调用时运行的代码。
让我们看看一个小例子：
class MyCell(torch.nn.Module):
def __init__(self):
super(MyCell, self).__init__()
def forward(self, x, h):
new_h = torch.tanh(x + h)
return new_h, new_h
my_cell = MyCell()
x = torch.rand(3, 4)
h = torch.rand(3, 4)
print(my_cell(x, h))
因此我们：
创建了一个继承自
torch.nn.Module的类。定义了一个构造函数。构造函数并没有做什么，只是调用了
super的构造函数。定义了一个
forward函数，该函数接受两个输入并返回两个输出。forward函数的实际内容并不重要，它有点像是一个假的 RNN 单元——也就是说，它是一个循环应用的功能。
我们实例化了模块，并创建了 x 和 h ，它们只是 3x4 的随机值矩阵。然后我们用 my_cell(x, h) 调用了这个单元。这反过来又调用了我们的 forward 函数。
让我们做一些更有趣的事情：
class MyCell(torch.nn.Module):
def __init__(self):
super(MyCell, self).__init__()
self.linear = torch.nn.Linear(4, 4)
def forward(self, x, h):
new_h = torch.tanh(self.linear(x) + h)
return new_h, new_h
my_cell = MyCell()
print(my_cell)
print(my_cell(x, h))
我们重新定义了我们的模块 MyCell ，但这次我们增加了一个 self.linear 属性，并在前向函数中调用了 self.linear 。
这里到底发生了什么？ torch.nn.Linear 是 PyTorch 标准库中的一个 Module 。就像 MyCell 一样，它可以使用调用语法来调用。我们正在构建一个 Module 的层次结构。
在 Module 上的 print 将给出 Module 的子类层次结构的可视化表示。在我们的例子中，我们可以看到我们的 Linear 子类及其参数。
以这种方式组合 Module ，我们可以简洁且易于阅读地编写具有可重用组件的模型。
你可能已经注意到输出中的 grad_fn 。这是 PyTorch 自动微分方法 autograd 的一个细节。简而言之，这个系统允许我们通过可能复杂的程序来计算导数。这种设计使得在模型编写方面具有巨大的灵活性。
现在我们来考察这种灵活性：
class MyDecisionGate(torch.nn.Module):
def forward(self, x):
if x.sum() > 0:
return x
else:
return -x
class MyCell(torch.nn.Module):
def __init__(self):
super(MyCell, self).__init__()
self.dg = MyDecisionGate()
self.linear = torch.nn.Linear(4, 4)
def forward(self, x, h):
new_h = torch.tanh(self.dg(self.linear(x)) + h)
return new_h, new_h
my_cell = MyCell()
print(my_cell)
print(my_cell(x, h))
我们再次重新定义了我们的 MyCell 类，但这次我们定义了 MyDecisionGate 。这个模块利用了控制流。控制流包括循环和 if -语句。
许多框架采用在给定完整程序表示的情况下计算符号导数的方法。然而，在 PyTorch 中，我们使用梯度带。我们记录操作的发生，并在计算导数时反向回放它们。这样，框架不必显式地为语言中的所有构造定义导数。
自动微分是如何工作的
TorchScript 基础
现在让我们来看一下我们的运行示例，看看我们如何应用 TorchScript。
简而言之，TorchScript 提供了工具来捕获您模型的定义，即使在 PyTorch 灵活和动态的本质面前。让我们首先来了解一下我们所说的跟踪（tracing）。
跟踪 Modules ¶
class MyCell(torch.nn.Module):
def __init__(self):
super(MyCell, self).__init__()
self.linear = torch.nn.Linear(4, 4)
def forward(self, x, h):
new_h = torch.tanh(self.linear(x) + h)
return new_h, new_h
my_cell = MyCell()
x, h = torch.rand(3, 4), torch.rand(3, 4)
traced_cell = torch.jit.trace(my_cell, (x, h))
print(traced_cell)
traced_cell(x, h)
我们回滚了一点，并采用了我们 MyCell 类的第二个版本。和之前一样，我们实例化了它，但这次我们调用了 torch.jit.trace ，传入 Module ，并传入了网络可能看到的示例输入。
这究竟做了什么？它调用了 Module ，记录了当 Module 运行时发生的操作，并创建了一个 torch.jit.ScriptModule 的实例（其中 TracedModule 是一个实例）
TorchScript 将其定义记录在中间表示（或 IR），在深度学习中通常被称为图。我们可以通过 .graph 属性来检查这个图：
print(traced_cell.graph)
然而，这是一个非常低级的表示，图中包含的大部分信息对最终用户来说并不有用。相反，我们可以使用 .code 属性来给出代码的 Python 语法解释：
print(traced_cell.code)
那么我们为什么要做所有这些？有几个原因：
TorchScript 代码可以在自己的解释器中调用，这基本上是一个受限的 Python 解释器。这个解释器不会获取全局解释器锁，因此可以在同一实例上同时处理许多请求。
这种格式允许我们将整个模型保存到磁盘，并在另一个环境中加载，例如在用非 Python 语言编写的服务器中。
TorchScript 为我们提供了一种表示形式，可以在代码上执行编译器优化，以提供更高效的执行
TorchScript 允许我们与许多需要比单个操作符更广泛程序视图的后端/设备运行时进行接口
我们可以看到，调用 traced_cell 产生的结果与 Python 模块相同：
print(my_cell(x, h))
print(traced_cell(x, h))
使用脚本转换模块
我们之所以使用模块的第二个版本，而不是带有控制流子模块的那个版本，是有原因的。现在让我们来分析一下：
class MyDecisionGate(torch.nn.Module):
def forward(self, x):
if x.sum() > 0:
return x
else:
return -x
class MyCell(torch.nn.Module):
def __init__(self, dg):
super(MyCell, self).__init__()
self.dg = dg
self.linear = torch.nn.Linear(4, 4)
def forward(self, x, h):
new_h = torch.tanh(self.dg(self.linear(x)) + h)
return new_h, new_h
my_cell = MyCell(MyDecisionGate())
traced_cell = torch.jit.trace(my_cell, (x, h))
print(traced_cell.dg.code)
print(traced_cell.code)
查看输出 .code ，我们可以看到 if-else 分支根本不存在！为什么？追踪器正是我们所说的那样：运行代码，记录发生的操作，并构建一个 ScriptModule 来执行这些操作。不幸的是，像控制流这样的东西被抹去了。
我们如何才能在 TorchScript 中忠实表示这个模块呢？我们提供了一个脚本编译器，它可以直接分析你的 Python 源代码并将其转换为 TorchScript。让我们使用脚本编译器转换 MyDecisionGate ：
scripted_gate = torch.jit.script(MyDecisionGate())
my_cell = MyCell(scripted_gate)
scripted_cell = torch.jit.script(my_cell)
print(scripted_gate.code)
print(scripted_cell.code)
哈喽！现在我们已经忠实地捕捉到了程序在 TorchScript 中的行为。现在让我们尝试运行这个程序：
# New inputs
x, h = torch.rand(3, 4), torch.rand(3, 4)
print(scripted_cell(x, h))
混合脚本和追踪
有些情况下需要使用追踪而不是脚本（例如，一个模块有许多基于常量 Python 值的架构决策，我们不希望这些值出现在 TorchScript 中）。在这种情况下，脚本可以与追踪组合： torch.jit.script 将内联已追踪模块的代码，而追踪将内联已脚本模块的代码。
第一种情况的一个例子：
class MyRNNLoop(torch.nn.Module):
def __init__(self):
super(MyRNNLoop, self).__init__()
self.cell = torch.jit.trace(MyCell(scripted_gate), (x, h))
def forward(self, xs):
h, y = torch.zeros(3, 4), torch.zeros(3, 4)
for i in range(xs.size(0)):
y, h = self.cell(xs[i], h)
return y, h
rnn_loop = torch.jit.script(MyRNNLoop())
print(rnn_loop.code)
第二种情况的一个例子：
class WrapRNN(torch.nn.Module):
def __init__(self):
super(WrapRNN, self).__init__()
self.loop = torch.jit.script(MyRNNLoop())
def forward(self, xs):
y, h = self.loop(xs)
return torch.relu(y)
traced = torch.jit.trace(WrapRNN(), (torch.rand(10, 3, 4)))
print(traced.code)
这样，当需要使用脚本和跟踪时，可以分别使用它们，并将它们一起使用。
保存和加载模型
我们提供 API 来将 TorchScript 模块保存到/从磁盘的存档格式中。此格式包括代码、参数、属性和调试信息，这意味着存档是模型的独立表示，可以在完全不同的进程中加载。让我们保存和加载我们的包装 RNN 模块：
traced.save('wrapped_rnn.pt')
loaded = torch.jit.load('wrapped_rnn.pt')
print(loaded)
print(loaded.code)
如您所见，序列化保留了模块层次结构和我们所检查的代码。该模型还可以加载到 C++中，以实现无 Python 的执行。
进一步阅读
我们已经完成了教程！想看更深入的演示，请查看使用 TorchScript 转换机器翻译模型的 NeurIPS 演示：https://colab.research.google.com/drive/1HiICg6jRkBnr5hvK2-VnMi88Vi9pUzEJ
脚本总运行时间：（0 分钟 0.000 秒）
