注意
单击此处下载完整的示例代码
TorchScript 简介¶
创建时间： Aug 09， 2019 |上次更新时间：2024 年 12 月 2 日 |上次验证： Nov 05， 2024
作者：詹姆斯·里德 （jamesreed@fb.com）、Michael Suo （suo@fb.com）、rev2
警告
TorchScript 不再处于积极开发阶段。
本教程是对 TorchScript 的介绍，TorchScript 是一个中级
PyTorch 模型（子类）的表示形式，该
然后，可以在 C++ 等高性能环境中运行。nn.Module
在本教程中，我们将介绍：
PyTorch 中模型创作的基础知识，包括：
模块
定义函数
forward将模块组合成模块层次结构
将 PyTorch 模块转换为 TorchScript 的具体方法，我们的 高性能部署运行时
跟踪现有模块
使用脚本直接编译模块
如何组合这两种方法
保存和加载 TorchScript 模块
我们希望在完成本教程后，您将继续完成后续教程，该教程将引导您完成实际调用 TorchScript 的示例 模型来自 C++。
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.5.0+cu124
<torch._C.Generator object at 0x7fce85f71b30>
PyTorch 模型创作的基础知识¶
让我们从定义一个简单的 .A 是
PyTorch 中的基本组成单位。它包含：ModuleModule
一个构造函数，用于准备模块以进行调用
一组 和 sub-。这些是初始化的 由构造函数使用，并且可以在调用期间由模块使用。
ParametersModules一个函数。这是在模块 被调用。
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定义了一个函数，该函数接受两个输入并返回 两个输出。函数的实际内容不是 真的很重要，但它有点像一个假的 RNN cell–那个 is——它是一个应用于循环的函数。
forwardforward
我们实例化了模块，并制作了 和 ，它们只是 3x4
随机值的矩阵。然后我们用 .这反过来又调用了我们的函数。xhmy_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>))
我们重新定义了我们的 module ，但这次我们添加了一个 attribute，并在 forward 中调用
功能。MyCellself.linearself.linear
这里到底发生了什么？ 是
PyTorch 标准库。就像 一样，它可以被调用
使用 CALL 语法。我们正在构建 s 的层次结构。torch.nn.LinearModuleMyCellModule
printon a 将提供 的子类层次结构的可视化表示。在我们的示例中，我们可以看到我们的子类及其参数。ModuleModuleLinear
通过以这种方式组合 s，我们可以简洁易读
使用可重用组件编写模型。Module
您可能已经注意到了输出。这是
PyTorch 的自动微分方法，称为 autograd。
简而言之，这个系统允许我们通过以下方式计算导数
可能复杂的程序。该设计允许大量的
模型创作的灵活性。grad_fn
现在让我们检查一下所述灵活性：
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>))
我们再次重新定义了我们的类，但这里我们定义了 。该模块利用 control flow。控制流
由循环和 -statement 等内容组成。MyCellMyDecisionGateif
许多框架采用计算符号导数的方法 给定完整的程序表示。但是，在 PyTorch 中，我们使用 渐变带。我们在操作发生时记录它们，并重播它们 倒退计算导数。这样，框架就不会 必须为 语言。
autograd 的工作原理¶
TorchScript 基础知识¶
现在让我们以运行示例为例，看看如何应用 TorchScript。
简而言之，TorchScript 提供了一些工具来捕获 模型，即使考虑到 PyTorch 的灵活和动态特性。 让我们首先研究一下我们所说的跟踪。
描图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>))
我们稍微倒带了一下，采用了我们类的第二个版本。和以前一样，我们已经实例化了它，但这次，我们调用了 ，传入了 ，传入了示例
网络可能会看到的 Inputs。MyCelltorch.jit.traceModule
这到底做了什么？它调用了 ，记录了
操作，并创建一个
的实例（其中是一个
实例）ModuleModuletorch.jit.ScriptModuleTracedModule
TorchScript 将其定义记录在中间表示中
（或 IR），在深度学习中通常称为图形。我们可以
检查具有以下属性的图形：.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]() # /var/lib/workspace/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) # /var/lib/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:191:0
%13 : Float(3, 4, strides=[4, 1], requires_grad=1, device=cpu) = aten::tanh(%12) # /var/lib/workspace/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)
但是，这是一个非常低级的表示形式，并且大多数
图表中包含的信息对最终用户没有用处。相反
我们可以使用 property 来给出 Python 语法解释
的代码中：.code
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 允许我们与许多后端/设备运行时交互 需要比单个操作员更广阔的程序视图。
我们可以看到，调用产生的结果与
Python 模块：traced_cell
(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>))
使用脚本转换模块¶
我们使用模块的第二个版本是有原因的，而不是 control-flow-loaded子模块。现在让我们来检查一下：
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)
/var/lib/workspace/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)
查看输出，我们可以看到分支
无处可寻！为什么？跟踪完全按照我们所说的那样：
运行代码，记录发生的操作，并构建一个完全执行此操作的 A。不幸的是，像 control 这样的事情
流被擦除。.codeif-elseScriptModule
我们如何在 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 的 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>))
混合编写脚本和跟踪¶
某些情况需要使用跟踪而不是脚本（例如
module 有许多基于常量的架构决策
我们不希望出现在 TorchScript 中的 Python 值）。在这个
case 时，脚本可以通过 tracing 来编写：will
内联 traced 模块的代码，tracing 将内联代码
对于脚本化模块。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)
这样，当情况需要时，可以使用脚本和跟踪 他们每个人都一起使用。
保存和加载模型¶
我们提供 API 用于将 TorchScript 模块保存到磁盘或从磁盘加载 archive 格式。此格式包括代码、参数、属性和 debug 信息，这意味着存档是独立的 模型的表示形式，可以加载到完全独立的 过程。让我们保存并加载包装好的 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.215 秒）