torchscript相关知识介绍（一）
TorchScript 是 PyTorch 模型（nn.Module的子类）的中间表示，可以在高性能环境（例如 C++）中运行。
注：onnx也是一种IR（中间表示）
- 将 PyTorch 模块转换为 TorchScript（我们的高性能部署运行时）的特定方法
- 跟踪现有模块
- 使用scripting直接编译模块
- 如何组合两种方法
- 保存和加载 TorchScript 模块
2、现在，让我们以正在运行的示例为例，看看如何应用 TorchScript。
简而言之，即使 PyTorch 具有灵活和动态的特性，TorchScript 也提供了捕获模型定义的工具。 让我们开始研究所谓的跟踪（tracing）。
（1）Tracing 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
mycell = Mycell()
x, h = torch.rand(3, 4), torch.rand(3, 4)
traced_model = torch.jit.trace(mycell, (x, h))
print(traced_model)
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)
那么为什么我们要进行所有这些操作？ 有以下几个原因：
（1）TorchScript 代码可以在其自己的解释器中调用（该解释器基本上是一个受限制的 Python 解释器）， 该解释器不获取全局解释器锁定，因此可以在同一实例上同时处理许多请求。
（2）这种格式允许我们将整个模型保存到磁盘上，然后将其加载到另一个环境中，例如在以 Python 以外的语言编写的服务器中
（3）TorchScript 为我们提供了一种表示形式（中间表示），其中我们可以对代码进行编译器优化以提供更有效的执行
（4）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=), 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=))
(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=), 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=))
二、Using Scripting to Convert Modules（使用 Scripting方式转换模块）
有一个原因是我们使用了模块的第二版，而不是使用带有大量控制流的子模块。 现在让我们检查一下：
import torch
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):
#MyCell类继承nn.Module类的__init__函数
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())
x, h = torch.rand(3, 4), torch.rand(3, 4)
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).forward((self.linear).forward(input, ), )
_1 = torch.tanh(torch.add(_0, h, alpha=1))
return (_1, _1)上面可以看到：查看.code输出，可以发现找不到if-else分支！ 为什么？ 跟踪完全按照我们所说的去做：运行代码，记录发生的操作，并构造一个执行此操作的ScriptModule。 不幸的是，诸如控制流之类的东西被擦除了。
我们如何在 TorchScript 中真实地表示此模块？我们提供了script 编译器，它可以直接分析您的 Python 源代码以将其转换为 TorchScript。
让我们使用script编译器转换MyDecisionGate：
import torch
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):
#MyCell类继承nn.Module类的__init__函数
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
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:
_0 = bool(torch.gt(torch.sum(x, dtype=None), 0))
if _0:
_1 = x
else:
_1 = torch.neg(x)
return _1
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 中程序的行为。 现在，让我们尝试运行该程序：
import torch
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):
#MyCell类继承nn.Module类的__init__函数
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
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)
# New inputs
x, h = torch.rand(3, 4), torch.rand(3, 4)
print(scripted_cell(x, h))输出：
(tensor([[ 0.7521, 0.2065, 0.4505, 0.5970],
[ 0.2100, -0.2474, 0.7773, 0.8026],
[ 0.4375, 0.5534, 0.5153, 0.5532]], grad_fn=), tensor([[ 0.7521, 0.2065, 0.4505, 0.5970],
[ 0.2100, -0.2474, 0.7773, 0.8026],
[ 0.4375, 0.5534, 0.5153, 0.5532]], grad_fn=)) 三、混合Scripting and Tracing两种机制
在某些情况下，需要使用跟踪（tracing）而不是脚本（scripting）（例如，一个模块具有许多基于不变的 Python 值做出的架构决策，而我们不希望它们出现在 TorchScript 中。
在这种情况下，脚本（scripting）可以与跟踪（tracing）组成：torch.jit.script 将内联已被跟踪（traced）的模块的代码，而跟踪（tracing）将内联已被脚本（scripted）化的模块的代码。
第一种情况的示例：
import torch
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):
#MyCell类继承nn.Module类的__init__函数
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
scripted_gate = torch.jit.script(MyDecisionGate())
x, h = torch.rand(3, 4), torch.rand(3, 4)
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)还有第二种情况的示例：
import torch
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):
#MyCell类继承nn.Module类的__init__函数
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
scripted_gate = torch.jit.script(MyDecisionGate())
x, h = torch.rand(3, 4), torch.rand(3, 4)
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
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, y, = (self.loop).forward(argument_1, )
return torch.relu(y)
这样，当情况需要它们时，可以一起混合使用脚本（scripting）机制和跟踪（tracing）机制。
scripting方式适用于程序拥有的大量控制流的情况，一些if-else的这种分支情况！！！
四、保存和加载模型
我们提供 API，以存档格式将 TorchScript 模块保存到磁盘或从磁盘加载 TorchScript 模块。 这种格式包括代码，参数，属性和调试信息，这意味着归档文件是模型的独立表示形式，可以在完全独立的过程中加载。 让我们保存并加载包装好的 RNN 模块：
import torch
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):
#MyCell类继承nn.Module类的__init__函数
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
scripted_gate = torch.jit.script(MyDecisionGate())
x, h = torch.rand(3, 4), torch.rand(3, 4)
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
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)))
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, y, = (self.loop).forward(argument_1, )
return torch.relu(y)如您所见，序列化保留了模块层次结构和我们一直在研究的代码。 也可以将模型加载到 C++ 中，以实现不依赖 Python 的执行。
参考：
https://pytorch.apachecn.org/#/docs/1.7/38
Introduction to TorchScript — PyTorch Tutorials 1.12.1+cu102 documentation