TorchScript 简介
原文:https://pytorch.org/tutorials/beginner/Intro_to_TorchScript_tutorial.html
James Reed (jamesreed@fb.com),Michael Suo (suo@fb.com),修订 2
本教程是 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__)
出:
1.7.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.8837, 0.5372, 0.4951, 0.9124],
[0.6124, 0.7072, 0.6395, 0.9585],
[0.6178, 0.8701, 0.8071, 0.2415]]), tensor([[0.8837, 0.5372, 0.4951, 0.9124],
[0.6124, 0.7072, 0.6395, 0.9585],
[0.6178, 0.8701, 0.8071, 0.2415]]))
因此,我们已经:
- 创建了一个子类
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.5042, 0.8137, -0.1593, 0.4167],
[ 0.1716, 0.8078, -0.2267, 0.7011],
[ 0.5616, 0.8753, 0.1597, -0.3899]], grad_fn=<TanhBackward>), tensor([[ 0.5042, 0.8137, -0.1593, 0.4167],
[ 0.1716, 0.8078, -0.2267, 0.7011],
[ 0.5616, 0.8753, 0.1597, -0.3899]], grad_fn=<TanhBackward>))
我们已经重新定义了模块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.8636, 0.5572, 0.6262, 0.8546],
[0.7766, 0.5056, 0.5357, 0.8360],
[0.7293, 0.7581, 0.7117, 0.2432]], grad_fn=<TanhBackward>), tensor([[0.8636, 0.5572, 0.6262, 0.8546],
[0.7766, 0.5056, 0.5357, 0.8360],
[0.7293, 0.7581, 0.7117, 0.2432]], grad_fn=<TanhBackward>))
我们再次重新定义了MyCell
类,但是在这里我们定义了MyDecisionGate
。 该模块利用控制流。 控制流包括循环和if
语句之类的内容。
给定完整的程序表示形式,许多框架都采用计算符号派生的方法。 但是,在 PyTorch 中,我们使用梯度色带。 我们记录发生的操作,并在计算派生时向后回放。 这样,框架不必为语言中的所有构造显式定义派生类。
Autograd 的工作原理
TorchScript 的基础
现在,让我们以正在运行的示例为例,看看如何应用 TorchScript。
简而言之,即使 PyTorch 具有灵活和动态的特性,TorchScript 也提供了捕获模型定义的工具。 让我们开始研究所谓的跟踪。
跟踪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)
)
我们倒退了一点,并学习了MyCell
类的第二版。 和以前一样,我们实例化了它,但是这一次,我们调用了torch.jit.trace
,将其传递给Module
,并传递给了示例输入,网络可能会看到。
这到底是做什么的? 它调用了Module
,记录了运行Module
时发生的操作,并创建了torch.jit.ScriptModule
的实例(其中TracedModule
是实例)
TorchScript 将其定义记录在中间表示(或 IR)中,在深度学习中通常称为图。 我们可以检查带有.graph
属性的图:
print(traced_cell.graph)
出:
graph(%self.1 : __torch__.MyCell,
%input : Float(3:4, 4:1, requires_grad=0, device=cpu),
%h : Float(3:4, 4:1, requires_grad=0, device=cpu)):
%19 : __torch__.torch.nn.modules.linear.Linear = prim::GetAttr[name="linear"](%self.1)
%21 : Tensor = prim::CallMethod[name="forward"](%19, %input)
%12 : int = prim::Constant[value=1]() # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0
%13 : Float(3:4, 4:1, requires_grad=1, device=cpu) = aten::add(%21, %h, %12) # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0
%14 : Float(3:4, 4:1, requires_grad=1, device=cpu) = aten::tanh(%13) # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0
%15 : (Float(3:4, 4:1, requires_grad=1, device=cpu), Float(3:4, 4:1, requires_grad=1, device=cpu)) = prim::TupleConstruct(%14, %14)
return (%15)
但是,这是一个非常低级的表示形式,图中包含的大多数信息对最终用户没有用。 相反,我们可以使用.code
属性来给出代码的 Python 语法解释:
print(traced_cell.code)
出:
def forward(self,
input: Tensor,
h: Tensor) -> Tuple[Tensor, Tensor]:
_0 = torch.add((self.linear).forward(input, ), h, alpha=1)
_1 = torch.tanh(_0)
return (_1, _1)
那么为什么我们要进行所有这些操作? 有以下几个原因:
- TorchScript 代码可以在其自己的解释器中调用,该解释器基本上是受限制的 Python 解释器。 该解释器不获取全局解释器锁定,因此可以在同一实例上同时处理许多请求。
- 这种格式允许我们将整个模型保存到磁盘上,然后将其加载到另一个环境中,例如在以 Python 以外的语言编写的服务器中
- TorchScript 为我们提供了一种表示形式,其中我们可以对代码进行编译器优化以提供更有效的执行
- TorchScript 允许我们与许多后端/设备运行时进行交互,与单个运算符相比,它们要求更广泛的程序视图。
我们可以看到,调用traced_cell
会产生与 Python 模块相同的结果:
print(my_cell(x, h))
print(traced_cell(x, h))
出:
(tensor([[-0.3869, 0.0678, 0.5692, 0.6332],
[ 0.1230, 0.4653, 0.8051, 0.3346],
[-0.5288, 0.2767, 0.9063, 0.4727]], grad_fn=<TanhBackward>), tensor([[-0.3869, 0.0678, 0.5692, 0.6332],
[ 0.1230, 0.4653, 0.8051, 0.3346],
[-0.5288, 0.2767, 0.9063, 0.4727]], grad_fn=<TanhBackward>))
(tensor([[-0.3869, 0.0678, 0.5692, 0.6332],
[ 0.1230, 0.4653, 0.8051, 0.3346],
[-0.5288, 0.2767, 0.9063, 0.4727]], grad_fn=<TanhBackward>), tensor([[-0.3869, 0.0678, 0.5692, 0.6332],
[ 0.1230, 0.4653, 0.8051, 0.3346],
[-0.5288, 0.2767, 0.9063, 0.4727]], grad_fn=<TanhBackward>))
使用脚本转换模块
有一个原因是我们使用了模块的第二版,而不是使用带有大量控制流的子模块。 现在让我们检查一下:
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)
出:
def forward(self,
input: Tensor,
h: Tensor) -> Tuple[Tensor, Tensor]:
_0 = self.dg
_1 = (self.linear).forward(input, )
_2 = (_0).forward(_1, )
_3 = torch.tanh(torch.add(_1, h, alpha=1))
return (_3, _3)
查看.code
输出,可以发现找不到if-else
分支! 为什么? 跟踪完全按照我们所说的去做:运行代码,记录发生的操作,并构造一个执行此操作的ScriptModule
。 不幸的是,诸如控制流之类的东西被擦除了。
我们如何在 TorchScript 中忠实地表示此模块? 我们提供了脚本编译器,它可以直接分析您的 Python 源代码以将其转换为 TorchScript。 让我们使用脚本编译器转换MyDecisionGate
:
scripted_gate = torch.jit.script(MyDecisionGate())
my_cell = MyCell(scripted_gate)
traced_cell = torch.jit.script(my_cell)
print(traced_cell.code)
出:
def forward(self,
x: Tensor,
h: Tensor) -> Tuple[Tensor, Tensor]:
_0 = (self.dg).forward((self.linear).forward(x, ), )
new_h = torch.tanh(torch.add(_0, 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)
出:
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).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,
argument_1: Tensor) -> Tensor:
_0, h, = (self.loop).forward(argument_1, )
return torch.relu(h)
这样,当情况需要它们时,可以使用脚本和跟踪并将它们一起使用。
保存和加载模型
我们提供 API,以存档格式将 TorchScript 模块保存到磁盘或从磁盘加载 TorchScript 模块。 这种格式包括代码,参数,属性和调试信息,这意味着归档文件是模型的独立表示形式,可以在完全独立的过程中加载。 让我们保存并加载包装好的 RNN 模块:
traced.save('wrapped_rnn.zip')
loaded = torch.jit.load('wrapped_rnn.zip')
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,
argument_1: Tensor) -> Tensor:
_0, h, = (self.loop).forward(argument_1, )
return torch.relu(h)
如您所见,序列化保留了模块层次结构和我们一直在研究的代码。 也可以将模型加载到 C++ 中,以实现不依赖 Python 的执行。
进一步阅读
我们已经完成了教程! 有关更多涉及的演示,请查看 NeurIPS 演示来使用 TorchScript 转换机器翻译模型。
脚本的总运行时间:(0 分钟 0.269 秒)
下载 Python 源码:Intro_to_TorchScript_tutorial.py
下载 Jupyter 笔记本:Intro_to_TorchScript_tutorial.ipynb
由 Sphinx 画廊生成的画廊