备注
点击 这里 下载完整示例代码
TorchScript 入门¶
Created On: Aug 09, 2019 | Last Updated: Dec 02, 2024 | Last Verified: Nov 05, 2024
作者: James Reed (jamesreed@fb.com), Michael Suo (suo@fb.com), 修订版2
警告
TorchScript 不再处于活跃开发中。
本教程是 TorchScript 的入门教程，TorchScript 是 PyTorch 模型（nn.Module 的子类）的中间表示，可以在高性能环境如 C++ 中运行。
在本教程中我们将覆盖以下内容：
PyTorch 模型编写的基础知识，包括：
模块
定义
forward函数将模块组成模块的层级结构
将 PyTorch 模块转换为 TorchScript 的特定方法，即我们的高性能部署运行时
追踪现有模块
利用脚本语法直接编译模块
如何合并这两种方法
保存和加载 TorchScript 模块
我们希望在完成本教程之后，您可以通过`后续教程 <https://pytorch.org/tutorials/advanced/cpp_export.html>`_进一步了解如何从 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
2.7.0+cu126
<torch._C.Generator object at 0x7f49be586a70>
PyTorch 模型编写基础¶
让我们从定义一个简单的 Module 开始。 Module 是 PyTorch 中的基本组合单元。 它包含：
一个构造函数，用于准备模块的调用
一组
Parameters和子Modules。它们由构造函数初始化，并可在模块调用期间使用。一个
forward函数。这是模块调用时运行的代码。
让我们看一下一个简单的例子：
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))
(tensor([[0.8219, 0.8990, 0.6670, 0.8277],
[0.5176, 0.4017, 0.8545, 0.7336],
[0.6013, 0.6992, 0.2618, 0.6668]]), tensor([[0.8219, 0.8990, 0.6670, 0.8277],
[0.5176, 0.4017, 0.8545, 0.7336],
[0.6013, 0.6992, 0.2618, 0.6668]]))
因此，我们已经完成了以下操作：
创建了一个继承自
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(
(linear): Linear(in_features=4, out_features=4, bias=True)
)
(tensor([[ 0.8573, 0.6190, 0.5774, 0.7869],
[ 0.3326, 0.0530, 0.0702, 0.8114],
[ 0.7818, -0.0506, 0.4039, 0.7967]], grad_fn=<TanhBackward0>), tensor([[ 0.8573, 0.6190, 0.5774, 0.7869],
[ 0.3326, 0.0530, 0.0702, 0.8114],
[ 0.7818, -0.0506, 0.4039, 0.7967]], grad_fn=<TanhBackward0>))
我们重新定义了模块 MyCell，但这次我们添加了一个 self.linear 属性，并在 forward 函数中调用了 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(
(dg): MyDecisionGate()
(linear): Linear(in_features=4, out_features=4, bias=True)
)
(tensor([[ 0.8346, 0.5931, 0.2097, 0.8232],
[ 0.2340, -0.1254, 0.2679, 0.8064],
[ 0.6231, 0.1494, -0.3110, 0.7865]], grad_fn=<TanhBackward0>), tensor([[ 0.8346, 0.5931, 0.2097, 0.8232],
[ 0.2340, -0.1254, 0.2679, 0.8064],
[ 0.6231, 0.1494, -0.3110, 0.7865]], grad_fn=<TanhBackward0>))
我们再次重新定义了 MyCell 类，但这次我们定义了 MyDecisionGate。 该模块利用了 控制流。 控制流包括循环和 if 语句。
许多框架使用的是根据完整程序表示计算符号导数的方法。然而，在 PyTorch 中，我们使用的是一种梯度带的方法。我们记录发生的操作，并在计算导数时反向重放它们。通过这种方式，框架无需为语言中的所有构造显式定义导数。
autograd 的工作原理¶
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(
original_name=MyCell
(linear): Linear(original_name=Linear)
)
(tensor([[-0.2541, 0.2460, 0.2297, 0.1014],
[-0.2329, -0.2911, 0.5641, 0.5015],
[ 0.1688, 0.2252, 0.7251, 0.2530]], grad_fn=<TanhBackward0>), tensor([[-0.2541, 0.2460, 0.2297, 0.1014],
[-0.2329, -0.2911, 0.5641, 0.5015],
[ 0.1688, 0.2252, 0.7251, 0.2530]], grad_fn=<TanhBackward0>))
我们回到 MyCell 类的第二版。与之前一样，我们实例化了它，但这次我们调用了 torch.jit.trace，传入了 Module 和网络可能看到的 示例输入。
这具体完成了什么？它调用了 Module，记录了运行 Module 时发生的操作，并创建了一个 torch.jit.ScriptModule 的实例（其中 TracedModule 是一个实例）。
TorchScript 将其定义记录在一个中间表示（IR）中，在深度学习中通常称为 图（graph）。我们可以通过 .graph 属性查看这个图：
print(traced_cell.graph)
graph(%self.1 : __torch__.MyCell,
%x : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu),
%h : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu)):
%linear : __torch__.torch.nn.modules.linear.Linear = prim::GetAttr[name="linear"](%self.1)
%20 : Tensor = prim::CallMethod[name="forward"](%linear, %x)
%11 : int = prim::Constant[value=1]() # /data1/lin/pytorch-tutorials/beginner_source/Intro_to_TorchScript_tutorial.py:191:0
%12 : Float(3, 4, strides=[4, 1], requires_grad=1, device=cpu) = aten::add(%20, %h, %11) # /data1/lin/pytorch-tutorials/beginner_source/Intro_to_TorchScript_tutorial.py:191:0
%13 : Float(3, 4, strides=[4, 1], requires_grad=1, device=cpu) = aten::tanh(%12) # /data1/lin/pytorch-tutorials/beginner_source/Intro_to_TorchScript_tutorial.py:191:0
%14 : (Float(3, 4, strides=[4, 1], requires_grad=1, device=cpu), Float(3, 4, strides=[4, 1], requires_grad=1, device=cpu)) = prim::TupleConstruct(%13, %13)
return (%14)
然而，这是一个非常底层的表示，图中包含的大多数信息对最终用户毫无用处。相反，我们可以使用 .code 属性来提供 Python 语法的代码解释：
print(traced_cell.code)
def forward(self,
x: Tensor,
h: Tensor) -> Tuple[Tensor, Tensor]:
linear = self.linear
_0 = torch.tanh(torch.add((linear).forward(x, ), h))
return (_0, _0)
那么我们为什么要做这些呢？有以下几个原因：
TorchScript 代码可以在其自己的解释器中调用，这是一个基本的受限 Python 解释器。该解释器不会获取全局解释器锁，因此可以在同一实例上同时处理多个请求。
这种格式允许我们将整个模型保存到磁盘，并加载到其他环境中，例如用非 Python 编写的服务器中。
TorchScript 为我们提供了一种表示，可以对代码进行编译器优化，以实现更高效的执行。
TorchScript 允许我们与需要比单个操作更广泛程序视图的许多后端/设备运行时进行接口。
我们可以看到调用 traced_cell 产生与 Python 模块相同的结果：
(tensor([[-0.2541, 0.2460, 0.2297, 0.1014],
[-0.2329, -0.2911, 0.5641, 0.5015],
[ 0.1688, 0.2252, 0.7251, 0.2530]], grad_fn=<TanhBackward0>), tensor([[-0.2541, 0.2460, 0.2297, 0.1014],
[-0.2329, -0.2911, 0.5641, 0.5015],
[ 0.1688, 0.2252, 0.7251, 0.2530]], grad_fn=<TanhBackward0>))
(tensor([[-0.2541, 0.2460, 0.2297, 0.1014],
[-0.2329, -0.2911, 0.5641, 0.5015],
[ 0.1688, 0.2252, 0.7251, 0.2530]], grad_fn=<TanhBackward0>), tensor([[-0.2541, 0.2460, 0.2297, 0.1014],
[-0.2329, -0.2911, 0.5641, 0.5015],
[ 0.1688, 0.2252, 0.7251, 0.2530]], grad_fn=<TanhBackward0>))
使用脚本将模块转换¶
我们之所以使用模块的第二个版本，而不是包含大量控制流的子模块，是有原因的。让我们现在研究这一点：
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)
/data1/lin/pytorch-tutorials/beginner_source/Intro_to_TorchScript_tutorial.py:263: TracerWarning:
Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
def forward(self,
argument_1: Tensor) -> NoneType:
return None
def forward(self,
x: Tensor,
h: Tensor) -> Tuple[Tensor, Tensor]:
dg = self.dg
linear = self.linear
_0 = (linear).forward(x, )
_1 = (dg).forward(_0, )
_2 = torch.tanh(torch.add(_0, h))
return (_2, _2)
从 .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)
def forward(self,
x: Tensor) -> Tensor:
if bool(torch.gt(torch.sum(x), 0)):
_0 = x
else:
_0 = torch.neg(x)
return _0
def forward(self,
x: Tensor,
h: Tensor) -> Tuple[Tensor, Tensor]:
dg = self.dg
linear = self.linear
_0 = torch.add((dg).forward((linear).forward(x, ), ), h)
new_h = torch.tanh(_0)
return (new_h, new_h)
成功了！ 我们现已将程序的行为忠实地捕捉在 TorchScript 中。现在让我们尝试运行程序：
# New inputs
x, h = torch.rand(3, 4), torch.rand(3, 4)
print(scripted_cell(x, h))
(tensor([[ 0.5679, 0.5762, 0.2506, -0.0734],
[ 0.5228, 0.7122, 0.6985, -0.0656],
[ 0.6187, 0.4487, 0.7456, -0.0238]], grad_fn=<TanhBackward0>), tensor([[ 0.5679, 0.5762, 0.2506, -0.0734],
[ 0.5228, 0.7122, 0.6985, -0.0656],
[ 0.6187, 0.4487, 0.7456, -0.0238]], grad_fn=<TanhBackward0>))
混合使用脚本与追踪¶
在某些情况下，应使用追踪而非脚本（例如，当模块有许多基于常量 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)
def forward(self,
xs: Tensor) -> Tuple[Tensor, Tensor]:
h = torch.zeros([3, 4])
y = torch.zeros([3, 4])
y0 = y
h0 = h
for i in range(torch.size(xs, 0)):
cell = self.cell
_0 = (cell).forward(torch.select(xs, 0, i), h0, )
y1, h1, = _0
y0, h0 = y1, h1
return (y0, h0)
接下来是第二种情况的例子：
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)
def forward(self,
xs: Tensor) -> Tensor:
loop = self.loop
_0, y, = (loop).forward(xs, )
return torch.relu(y)
通过这种方式，当需要时可以单独或结合使用脚本和追踪。
保存和加载模型¶
我们提供了用于以存档格式保存和加载 TorchScript 模块的 API。此格式包括代码、参数、属性和调试信息，这意味着存档是模型的独立表示，可以在完全独立的进程中加载。让我们保存和加载我们封装的 RNN 模块：
traced.save('wrapped_rnn.pt')
loaded = torch.jit.load('wrapped_rnn.pt')
print(loaded)
print(loaded.code)
RecursiveScriptModule(
original_name=WrapRNN
(loop): RecursiveScriptModule(
original_name=MyRNNLoop
(cell): RecursiveScriptModule(
original_name=MyCell
(dg): RecursiveScriptModule(original_name=MyDecisionGate)
(linear): RecursiveScriptModule(original_name=Linear)
)
)
)
def forward(self,
xs: Tensor) -> Tensor:
loop = self.loop
_0, y, = (loop).forward(xs, )
return torch.relu(y)
如您所见，序列化保留了模块层次结构以及我们一直在审视的代码。该模型还可以加载，例如，到 C++ 中 用于无需 Python 的执行。
进一步阅读¶
我们完成了本教程！如需更深入的展示，请查看 NeurIPS 演示，了解使用 TorchScript 转换机器翻译模型：https://colab.research.google.com/drive/1HiICg6jRkBnr5hvK2-VnMi88Vi9pUzEJ
脚本总运行时间: （0 分钟 0.159 秒）
