注意
here下载完整的示例代码
TorchScript 简介¶
James Reed (jamesreed@fb.com), Michael Suo (suo@fb.com), rev2
本教程介绍 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__)
输出：
1.3.1
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.9292, 0.9080, 0.4408, 0.9567],
[0.8918, 0.9411, 0.8546, 0.8362],
[0.8460, 0.7734, 0.9107, 0.7694]]), tensor([[0.9292, 0.9080, 0.4408, 0.9567],
[0.8918, 0.9411, 0.8546, 0.8362],
[0.8460, 0.7734, 0.9107, 0.7694]]))
因此，我们：
创建了一个类，该类将
torch.nn.Module.已定义构造函数。 构造函数不执行太多操作，只是调用
super。定义了
forward函数，该函数需要两个输入并返回两个输出。forward函数的实际内容并不重要，但它有点像一个假的 RNN 单元格，即应用于循环的函数。
我们实例化了模块，并制作了x和y它们只是随机值的 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.4868, -0.2578, 0.4380, 0.6714],
[ 0.2950, -0.2403, 0.7920, 0.3769],
[ 0.1951, -0.5589, 0.8387, -0.0487]], grad_fn=<TanhBackward>), tensor([[ 0.4868, -0.2578, 0.4380, 0.6714],
[ 0.2950, -0.2403, 0.7920, 0.3769],
[ 0.1951, -0.5589, 0.8387, -0.0487]], grad_fn=<TanhBackward>))
我们已经重新定义了模块MyCell但这次我们添加了一个self.linear属性，并在正向函数中调用self.linear。
这里到底发生了什么？ torch.nn.Linear是来自 PyTorch 标准库的Module。 与MyCell一样，可以使用调用语法调用它。 我们正在构建Module的层次结构。
在Module上print将提供Module子类层次结构的可视化表示形式。 在我们的示例中，我们可以看到Linear子类及其参数。
通过这种方式编写Module，我们可以用可重用的组件以可读的方式创作模型。
您可能已经注意到输出上的grad_fn 这是 PyTorch 的自动分化方法（称为自动梯度）的一个细节。
简而言之，该系统允许我们通过潜在的复杂程序计算导数。 该设计允许在模型创作方面具有很大的灵活性。
现在，让我们来检查一下所说的灵活性：
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.8270, 0.9801, 0.6031, 0.5930],
[ 0.6872, 0.9649, 0.8835, 0.4429],
[ 0.7594, 0.9094, 0.9068, -0.2839]], grad_fn=<TanhBackward>), tensor([[ 0.8270, 0.9801, 0.6031, 0.5930],
[ 0.6872, 0.9649, 0.8835, 0.4429],
[ 0.7594, 0.9094, 0.9068, -0.2839]], grad_fn=<TanhBackward>))
我们再次重新定义了 MyCell 类，但在这里我们定义了MyDecisionGate。 此模块利用控制流。 控制流由循环和if-语句等部分组成。
许多框架采用计算符号导数的方法，给出完整的程序表示。 但是，在 PyTorch 中，我们使用渐变胶带。 我们记录发生时的操作，并在计算导数中向后重放它们。 这样，框架不必显式定义语言中所有构造的导数。
自动分级的工作原理|
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)
输出：
TracedModule[MyCell](
original_name=MyCell
(linear): TracedModule[Linear](original_name=Linear)
)
我们退了一点，取了MyCell类的第二个版本。 和以前一样，我们已经实例化了它，但这次，我们调用torch.jit.trace在Module中传递，并在网络可能看到的示例输入中传递。
这到底做了什么？ 它调用了Module，记录了Module运行时发生的操作，并创建了一个torch.jit.ScriptModule实例（TracedModule是实例）
TorchScript在中间表示（或IR）中记录其定义，在深度学习中通常称为图形。 我们可以使用.graph属性检查图形：
print(traced_cell.graph)
输出：
但是，这是一个非常低级别的表示形式，并且图形中包含的大多数信息对最终用户没有用处。 相反，我们可以使用.code属性为代码提供 Python 语法解释：
print(traced_cell.code)
输出：
import __torch__
import __torch__.torch.nn.modules.linear
def forward(self,
input: Tensor,
h: Tensor) -> Tuple[Tensor, Tensor]:
_0 = self.linear
weight = _0.weight
bias = _0.bias
_1 = torch.addmm(bias, input, torch.t(weight), beta=1, alpha=1)
_2 = torch.tanh(torch.add(_1, h, alpha=1))
return (_2, _2)
那我们为什么要这么做呢？ 有几个原因：
TorchScript 代码可以在其自己的解释器中调用，这基本上是一个受限的 Python 解释器。 此解释器不获取全局解释器锁，因此可以在同一实例上同时处理许多请求。
此格式允许我们将整个模型保存到磁盘并将其加载到其他环境中，例如以 Python 以外的语言编写的服务器
TorchScript 为我们提供了一个表示形式，我们可以在代码上执行编译器优化，以提供更高效的执行
TorchScript 允许我们与许多后端/设备运行时进行接口，这些运行时需要比单个运算符更宽泛的程序视图。
我们可以看到，调用traced_cell产生的结果与 Python 模块相同：
print(my_cell(x, h))
print(traced_cell(x, h))
输出：
(tensor([[0.5565, 0.4329, 0.4924, 0.8865],
[0.3106, 0.6977, 0.8196, 0.7602],
[0.7294, 0.6298, 0.5052, 0.8642]], grad_fn=<TanhBackward>), tensor([[0.5565, 0.4329, 0.4924, 0.8865],
[0.3106, 0.6977, 0.8196, 0.7602],
[0.7294, 0.6298, 0.5052, 0.8642]], grad_fn=<TanhBackward>))
(tensor([[0.5565, 0.4329, 0.4924, 0.8865],
[0.3106, 0.6977, 0.8196, 0.7602],
[0.7294, 0.6298, 0.5052, 0.8642]],
grad_fn=<DifferentiableGraphBackward>), tensor([[0.5565, 0.4329, 0.4924, 0.8865],
[0.3106, 0.6977, 0.8196, 0.7602],
[0.7294, 0.6298, 0.5052, 0.8642]],
grad_fn=<DifferentiableGraphBackward>))
使用脚本转换模块|
我们使用模块的第二版是有原因的，而不是带有控制流子模块的子模块。 现在让我们来研究一下：
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.code)
输出：
import __torch__.___torch_mangle_14
import __torch__
import __torch__.torch.nn.modules.linear.___torch_mangle_15
def forward(self,
input: Tensor,
h: Tensor) -> Tuple[Tensor, Tensor]:
_0 = self.linear
weight = _0.weight
bias = _0.bias
x = torch.addmm(bias, input, torch.t(weight), beta=1, alpha=1)
_1 = torch.tanh(torch.add(x, h, alpha=1))
return (_1, _1)
查看.code输出，我们可以看到if-else分支无处可寻！ 为什么？ 跟踪完全按照我们所说的方式执行：运行代码，记录发生的操作，并构造一个脚本模块，该脚本模块正是这样做的。 不幸的是，像控制流之类的东西被抹去。
我们如何在 TorchScript 中忠实地表示此模块？ 我们提供一个脚本编译器，它直接分析你的Python源代码，将其转换为火炬脚本。 让我们使用脚本编译器转换MyDecisionGate
scripted_gate = torch.jit.script(MyDecisionGate())
my_cell = MyCell(scripted_gate)
traced_cell = torch.jit.script(my_cell)
print(traced_cell.code)
输出：
import __torch__.___torch_mangle_17
import __torch__.___torch_mangle_16
import __torch__.torch.nn.modules.linear.___torch_mangle_18
def forward(self,
x: Tensor,
h: Tensor) -> Tuple[Tensor, Tensor]:
_0 = self.linear
_1 = _0.weight
_2 = _0.bias
if torch.eq(torch.dim(x), 2):
_3 = torch.__isnot__(_2, None)
else:
_3 = False
if _3:
bias = ops.prim.unchecked_unwrap_optional(_2)
ret = torch.addmm(bias, x, torch.t(_1), beta=1, alpha=1)
else:
output = torch.matmul(x, torch.t(_1))
if torch.__isnot__(_2, None):
bias0 = ops.prim.unchecked_unwrap_optional(_2)
output0 = torch.add_(output, bias0, alpha=1)
else:
output0 = output
ret = output0
_4 = torch.gt(torch.sum(ret, dtype=None), 0)
if bool(_4):
_5 = ret
else:
_5 = torch.neg(ret)
new_h = torch.tanh(torch.add(_5, h, alpha=1))
return (new_h, new_h)
万岁！ 现在，我们已经忠实地在 TorchScript 中捕获了程序的行为。 现在，让我们尝试运行该程序：
# New inputs
x, h = torch.rand(3, 4), torch.rand(3, 4)
traced_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)
输出：
import __torch__
import __torch__.___torch_mangle_19
import __torch__.___torch_mangle_16
import __torch__.torch.nn.modules.linear.___torch_mangle_20
def forward(self,
xs: Tensor) -> Tuple[Tensor, Tensor]:
h = torch.zeros([3, 4], dtype=None, layout=None, device=None, pin_memory=None)
y = torch.zeros([3, 4], dtype=None, layout=None, device=None, pin_memory=None)
y0 = y
h0 = h
for i in range(torch.size(xs, 0)):
_0 = self.cell
_1 = torch.select(xs, 0, i)
_2 = _0.linear
weight = _2.weight
bias = _2.bias
_3 = torch.addmm(bias, _1, torch.t(weight), beta=1, alpha=1)
_4 = torch.gt(torch.sum(_3, dtype=None), 0)
if bool(_4):
_5 = _3
else:
_5 = torch.neg(_3)
_6 = torch.tanh(torch.add(_5, h0, alpha=1))
y0, h0 = _6, _6
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)
输出：
import __torch__
import __torch__.___torch_mangle_23
import __torch__.___torch_mangle_21
import __torch__.___torch_mangle_16
import __torch__.torch.nn.modules.linear.___torch_mangle_22
def forward(self,
argument_1: Tensor) -> Tensor:
_0 = self.loop
h = torch.zeros([3, 4], dtype=None, layout=None, device=None, pin_memory=None)
h0 = h
for i in range(torch.size(argument_1, 0)):
_1 = _0.cell
_2 = torch.select(argument_1, 0, i)
_3 = _1.linear
weight = _3.weight
bias = _3.bias
_4 = torch.addmm(bias, _2, torch.t(weight), beta=1, alpha=1)
_5 = torch.gt(torch.sum(_4, dtype=None), 0)
if bool(_5):
_6 = _4
else:
_6 = torch.neg(_4)
h0 = torch.tanh(torch.add(_6, h0, alpha=1))
return torch.relu(h0)
这样，当情况需要每个脚本和跟踪时，可以使用脚本和跟踪，并一起使用。
保存和加载模型|
我们提供 API，以存档格式将 TorchScript 模块保存到磁盘或从磁盘加载。 此格式包括代码、参数、属性和调试信息，这意味着存档是模型的独立表示形式，可以在完全独立进程中加载。 让我们保存并加载包装的 RNN 模块：
traced.save('wrapped_rnn.zip')
loaded = torch.jit.load('wrapped_rnn.zip')
print(loaded)
print(loaded.code)
输出：
ScriptModule(
original_name=WrapRNN
(loop): ScriptModule(
original_name=MyRNNLoop
(cell): ScriptModule(
original_name=MyCell
(dg): ScriptModule(original_name=MyDecisionGate)
(linear): ScriptModule(original_name=Linear)
)
)
)
import __torch__
import __torch__.___torch_mangle_23
import __torch__.___torch_mangle_21
import __torch__.___torch_mangle_16
import __torch__.torch.nn.modules.linear.___torch_mangle_22
def forward(self,
argument_1: Tensor) -> Tensor:
_0 = self.loop
h = torch.zeros([3, 4], dtype=None, layout=None, device=None, pin_memory=None)
h0 = h
for i in range(torch.size(argument_1, 0)):
_1 = _0.cell
_2 = torch.select(argument_1, 0, i)
_3 = _1.linear
weight = _3.weight
bias = _3.bias
_4 = torch.addmm(bias, _2, torch.t(weight), beta=1, alpha=1)
_5 = torch.gt(torch.sum(_4, dtype=None), 0)
if bool(_5):
_6 = _4
else:
_6 = torch.neg(_4)
h0 = torch.tanh(torch.add(_6, h0, alpha=1))
return torch.relu(h0)
如您所见，序列化将保留模块层次结构和我们在其中一直在检查的代码。 例如，也可以将模型加载到C++中，以便无 python 执行。
进一步阅读»
我们已经完成了我们的教程！ 有关更多演示，请查看 NeurIPS 演示，该演示用于使用 TorchScript 转换机器翻译模型：https://colab.research.google.com/drive/1HiICg6jRkBnr5hvK2-VnMi88Vi9pUzEJ
脚本总运行时间： （ 0 分 0.257 秒）