TorchScript 简介
Contents
TorchScript 简介#
TorchScript 是 PyTorch 模型的中间表示形式(Module
的子类),可以在 C++ 等高性能环境中运行。
PyTorch 模型创建的基础知识#
从定义简单的 Module
开始。Module
是 PyTorch 中的基本组成单元。它包含:
构造函数,它为调用模块做准备
一组
Parameters
和子Module
。它们由构造函数初始化,module
可以在调用期间使用forward
函数。这是调用模块时运行的代码。
简单示例:
import torch
class MyCell(torch.nn.Module):
def __init__(self):
super().__init__()
def forward(self, x, h):
new_h = torch.tanh(x + h)
return new_h
my_cell = MyCell()
x = torch.rand(3, 4)
h = torch.rand(3, 4)
print(my_cell(x, h))
tensor([[0.0702, 0.3656, 0.8002, 0.7619],
[0.9325, 0.5827, 0.3802, 0.6648],
[0.8181, 0.7089, 0.7629, 0.5327]])
class MyCell(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(4, 4)
def forward(self, x, h):
new_h = torch.tanh(self.linear(x) + h)
return 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.2894, -0.2291, 0.0220, 0.1802],
[ 0.3967, 0.1058, 0.4551, 0.3112],
[ 0.3052, -0.0664, 0.7393, 0.3936]], 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):
super().__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
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.2625, -0.3511, 0.9434, -0.4812],
[ 0.8007, 0.2158, 0.8831, 0.0060],
[ 0.6084, -0.1552, 0.9591, -0.3329]], grad_fn=<TanhBackward0>)
TorchScript 基础#
现在让我们以运行的示例为例,看看如何应用 TorchScript。简而言之,即使考虑到 PyTorch 的灵活性和动态特性,TorchScript 也提供了捕获模型定义的工具。让我们从研究所谓的跟踪开始。
class MyCell(torch.nn.Module):
def __init__(self):
super().__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.9178, -0.6887, 0.8918, -0.0055],
[ 0.6206, -0.4067, 0.9349, 0.2997],
[ 0.9303, -0.1039, 0.8448, 0.2920]], grad_fn=<TanhBackward0>),
tensor([[ 0.9178, -0.6887, 0.8918, -0.0055],
[ 0.6206, -0.4067, 0.9349, 0.2997],
[ 0.9303, -0.1039, 0.8448, 0.2920]], grad_fn=<TanhBackward0>))
type(traced_cell)
torch.jit._trace.TopLevelTracedModule
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]() # /tmp/ipykernel_631109/2138254843.py:7:0
%12 : Float(3, 4, strides=[4, 1], requires_grad=1, device=cpu) = aten::add(%20, %h, %11) # /tmp/ipykernel_631109/2138254843.py:7:0
%13 : Float(3, 4, strides=[4, 1], requires_grad=1, device=cpu) = aten::tanh(%12) # /tmp/ipykernel_631109/2138254843.py:7: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 模块相同的结果:
print(my_cell(x, h))
print(traced_cell(x, h))
(tensor([[ 0.9178, -0.6887, 0.8918, -0.0055],
[ 0.6206, -0.4067, 0.9349, 0.2997],
[ 0.9303, -0.1039, 0.8448, 0.2920]], grad_fn=<TanhBackward0>), tensor([[ 0.9178, -0.6887, 0.8918, -0.0055],
[ 0.6206, -0.4067, 0.9349, 0.2997],
[ 0.9303, -0.1039, 0.8448, 0.2920]], grad_fn=<TanhBackward0>))
(tensor([[ 0.9178, -0.6887, 0.8918, -0.0055],
[ 0.6206, -0.4067, 0.9349, 0.2997],
[ 0.9303, -0.1039, 0.8448, 0.2920]],
grad_fn=<DifferentiableGraphBackward>), tensor([[ 0.9178, -0.6887, 0.8918, -0.0055],
[ 0.6206, -0.4067, 0.9349, 0.2997],
[ 0.9303, -0.1039, 0.8448, 0.2920]],
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().__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)
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)
/tmp/ipykernel_631109/3110858787.py:3: 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!
if x.sum() > 0:
查看 .code
输出,我们可以看到 if-else 分支无处可寻!为什么?跟踪所做的正是我们所说的:运行代码,记录所发生的运算,并构造 ScriptModule 来完成这些工作。不幸的是,像控制流这样的东西被删除了。
我们如何在 TorchScript 中忠实地表示这个模块?Torch 提供了脚本编译器,它直接分析您的 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)
# New inputs
x, h = torch.rand(3, 4), torch.rand(3, 4)
traced_cell(x, h)
(tensor([[ 0.6796, 0.0787, 0.8243, 0.4358],
[ 0.4151, -0.1153, 0.8075, 0.5238],
[ 0.8052, 0.4922, 0.8099, 0.2496]], grad_fn=<TanhBackward0>),
tensor([[ 0.6796, 0.0787, 0.8243, 0.4358],
[ 0.4151, -0.1153, 0.8075, 0.5238],
[ 0.8052, 0.4922, 0.8099, 0.2496]], grad_fn=<TanhBackward0>))
混合脚本和跟踪#
有些情况需要使用跟踪而不是脚本(例如,模块有许多基于常量 Python 值的架构决策,而我们不希望这些常量 Python 值出现在 TorchScript 中)。在这种情况下,脚本可以与跟踪组合在一起:torch.jit.script
将内联被跟踪模块的代码,而跟踪将内联被脚本模块的代码。
class MyRNNLoop(torch.nn.Module):
def __init__(self):
super().__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().__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)
这样,当情况需要脚本和跟踪时,就可以使用它们并将它们一起使用。
保存和加载模型#
Torch 提供 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 执行。