注意
点击 这里 下载完整示例代码
TorchScript简介¶
创建日期: 2019年8月9日 | 最后更新日期: 2024年12月2日 | 最后验证日期: 2024年11月5日
作者： James Reed (jamesreed@fb.com), Michael Suo (suo@fb.com), 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
2.5.0+cu124
<torch._C.Generator object at 0x7fce85f71b30>
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 属性，并在前向函数中调用了 self.linear。
这里到底发生了什么？torch.nn.Linear 是 PyTorch 标准库中的一个 Module，就像MyCell一样，可以使用调用语法来调用它。我们正在构建一个 Module 层级。
print on a Module 将给出 Module 的子类层次结构的可视化表示。在我们的示例中，我们可以看到我们的 Linear 子类及其参数。
通过以这种方式组合 Modules，我们可以简洁且易读地编写可重用组件的模型。
您可能已经在输出中注意到了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 灵活和动态的特性下也是如此。 让我们开始了解一下我们称之为追踪的技术。
跟踪 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>))
我们稍微回溯了一步，使用了我们第2个版本的 MyCell
类。和之前一样，我们实例化了它，但这次，我们调用了 torch.jit.trace，
传入了 Module，并传入了网络可能见到的 示例输入。
它到底做了什么？它调用了Module，记录了当Module运行时发生的所有操作，并创建了一个torch.jit.ScriptModule的实例（TracedModule是该实例）
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)
然而，这只是一个非常低级的表示，图中包含的大部分信息对最终用户没有用处。相反，我们可以使用.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 解释器。这个解释器不会获取全局解释器锁，因此可以在同一实例上同时处理多个请求。
这种格式允许我们将整个模型保存到磁盘，并在另一个环境中加载它，例如用其他语言编写的服务器中。
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)
/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)
查看 .code 的输出，我们可以发现 if-else 分支竟然无处寻觅！为什么？追踪正是做了我们说它会做的事情：
运行代码，记录发生的操作 并构建一个 ScriptModule 来执行这些操作。不幸的是，像控制流这样的东西会被擦除。
我们如何忠实地在TorchScript中表示这个模块？我们提供了一个
脚本编译器，它可以对你的Python源代码进行直接分析，并将其转换为TorchScript。让我们使用脚本编译器将 MyDecisionGate
转换为TorchScript：
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)
这样，在需要的情况下可以单独使用脚本记录和追踪，也可以将它们结合起来使用。
保存和加载模型¶
我们提供了API来以存档格式将TorchScript模块保存到/从磁盘加载。这种格式包括代码、参数、属性和调试信息，意味着存档可以独立于其他过程被加载。让我们保存并加载我们的包装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的执行。
进一步阅读¶
我们完成了教程！要查看更详细的演示，请访问使用TorchScript转换机器翻译模型的NeurIPS演示： https://colab.research.google.com/drive/1HiICg6jRkBnr5hvK2-VnMi88Vi9pUzEJ
脚本总运行时间： ( 0 分钟 0.215 秒)