【模型推理加速系列】05: 推理加速格式TorchScript简介及其应用
简介
本文紧接前文：模型推理加速系列｜04：BERT模型推理加速 TorchScript vs. ONNX 实验结果：在动态文本长度且大batch size的场景下，TorchScript 格式的inference性能具有一定优越性，进一步介绍TorchScript的基本概念及其在 C++ 中的使用。下一篇会以resnet18模型为例，对比Pytorch模型格式、ONNX、TorchScript等模型格式在GPU上的inference性能，敬请期待。更多、更新文章欢迎关注微信公众号：小窗幽记机器学习
JIT
在介绍TorchScript之前先简要说明一下什么是JIT。JIT 其实是一种概念，全称是 Just In Time Compilation，中文译为即时编译，是一种程序优化方法。常见的JIT使用场景是正则表达式。例如，在 Python 中使用正则表达式：
pat = re.compile(pattern)
result = pat.match(string)
或：
result = re.match(pattern, string)
两种写法从结果上来说是等价的。但第一种写法，会先对正则表达式进行compile再使用。如果多次使用某一个正则表达式，建议先对其进行compile，再使用compile之后得到的对象来做正则匹配。而这个compile的过程就是JIT(即时编译)。
深度学习中 JIT 的思想随处可见，比如 keras 框架的 model.compile，TensorFlow 中的 Graph 也是一种 JIT。PyTorch中的JIT是内生于PyTorch框架的一种DSL (Domain-specific language, DSL) 和Compiler栈的集合，目的是为了让PyTorch的使用者拥有方便且高性能模型推理的方法。Pytorch生态支持2种开发模式，一种是大家常用的 Eager 模式(主要用于训练、debug等环境)。这种模式使得PyTorch从面世以来以易用性收获一大波忠实粉丝。这种易用性得益于 PyTorch eager模式下的动态图结构。Pytorch的eager模式让使用者能过以贴合原生 Python方式进行开发。使用者可以在 PyTorch 的模型前向推理过程加任何Python的流程控制语句，甚至是断点单步跟进。但早期的TensorFlow(目前TensorFlow也有Eager模式)，则需要使用 tf.cond 等 TensorFlow 自己开发的流程控制，比较麻烦。Pytorch的另一种开发模式是 Script 模式，主要用于生产环境，包含2个主要模块 Pytorch JIT 和 TorchScript。其中 Pytorch JIT 是PyTorch程序优化的编译器。TorchScript将Python语言定义的PyTorch模型转化成一个序列化文件，即脚本文件。那么为什么 PyTorch 还要引入 JIT 呢？
动态图模型通过牺牲一些高级特性来换取易用性，那JIT有哪些特性，又会在什么情况下使用JIT？PyTorch JIT 模式就是一种将原本eager模式的表达转变并固定为一个计算图，便于进行优化和序列化。
Pytorch 中引入 JIT 主要有以下几个好处：
-
模型部署的可移植性
PyTorch 的 1.0 版本发布的最核心的两个新特性就是 JIT 和 C++ API，这两个特性一起发布不是没有道理的，JIT 是 Python 和 C++ 的桥梁。开发者可以使用 Python 训练模型，然后通过 JIT 将模型转为语言无关的模块， 从而让 C++ 可以非常方便得调用，从此使用 Python 训练模型，使用 C++ 将模型部署到生产环境对 PyTorch 来说成为了一件很容易的事。而因为使用了 C++， 开发者几乎可以把 PyTorch 模型部署到任意平台和设备上：树莓派、iOS、Android 等。 -
性能提升
JIT 既然是为生产部署而生，那免不了在性能上面做了极大优化。如果推断的场景对性能要求高，则可以考虑将模型(torch.nn.Module)转换为 TorchScript Module， 再进行推断。 -
模型可视化
TensorFlow 或 Keras 对模型可视化(TensorBoard等工具)非常友好，因为本身就是静态图的编程模型，在模型定义好后整个模型的结构和前向推理逻辑就已经清楚了；但 PyTorch 本身是不支持的，所以 PyTorch 模型在可视化上一直表现得不好，但 JIT 改善了这一情况。现在可以使用 JIT 的 trace 功能来得到 PyTorch 模型针对某一输入的前向推理逻辑，通过前向推理逻辑可以得到模型大致的结构，但如果在forward方法中有很多条件控制语句，则需要使用 PyTorch JIT 中的Scripting 的方式，这两种方式将在下文详细介绍。
TorchScript
TorchScript是PyTorch模型的一种中间表示(即常见的intermediate representation，IR)，是nn.Module的子类，基于这种模型格式可以在高性能的环境如C++中运行。TorchScript这种PyTorch模型表示格式可以由 TorchScript 编译器理解、编译和序列化而成。从根本上说，TorchScript 本身就是一种编程语言。它是使用 PyTorch API 的 Python 的子集。简单来说，TorchScript 软件栈可以将 Python 代码转换成 C++ 代码。TorchScript 软件栈包括两部分：TorchScript(Python)和 LibTorch(C++)。TorchScript负责将 Python 代码转成一个模型文件，LibTorch负责解析运行这个模型文件。后续会从代码层面介绍如何保存为TorchScript模型格式及其在Python和C++中如何加载使用。
保存TorchScript模型
将Pytorch模型保存为 TorchScript 格式有两种模式：trace 模式和 script 模式。
trace 模式
trace 模式顾名思义就是跟踪模型的执行，然后记录执行过程中的路径。在使用 trace 模式时，需要构造一个符合要求的输入，然后使用 TorchScript tracer 运行一遍，记录整个运行过程。在 trace 模式中运行时，每执行一个算子，就会往当前的 graph 加入一个 node。所有代码执行完毕，每一步的操作就会以一个计算图里的某个节点的形式被保存下来。值得一提的是，PyTorch 导出 ONNX 也是使用了这部分代码，所以理论上能够导出 ONNX 的模型也能够使用 trace 模式导出 TorchScript 格式模型。
trace 模式有以下2点限制：
- 不能有 if-else 等控制流
- 只支持 Tensor 操作。不支持非Tensor 操作，如List、Tuple、Map 等容器操作
通过上述对实现方式的解释，很容易理解为什么有这种限制：
- 跟踪出的 graph 是静态的。 如果有控制流，那么记录下来的只是当时生成模型时走的那条路径
- 追踪代码是跟 Tensor 算子绑定在一起的，如果是非 Tensor 的操作则无法被记录。
通过 trace 模式的特点可以看出，trace 模式通常用于深度模型的导出。常见的深度模型通常没有 if-else 控制流且通常没有非 Tensor 操作。
简单示例：
import torch
torch.manual_seed(42)
class TestTraceCell(torch.nn.Module):
def __init__(self):
super(TestTraceCell, 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
test_cell = TestTraceCell()
x, h = torch.rand(3, 4), torch.rand(3, 4)
traced_cell = torch.jit.trace(test_cell, (x, h))
print(traced_cell)
print("trace module=",traced_cell(x, h))
运行结果如下：
TestTraceCell(
original_name=TestTraceCell
(linear): Linear(original_name=Linear)
)
trace module= (tensor([[0.9567, 0.6879, 0.2618, 0.7927],
[0.8227, 0.7464, 0.4165, 0.5366],
[0.8193, 0.1679, 0.8132, 0.9052]], grad_fn=), tensor([[0.9567, 0.6879, 0.2618, 0.7927],
[0.8227, 0.7464, 0.4165, 0.5366],
[0.8193, 0.1679, 0.8132, 0.9052]], grad_fn=))
torch.jit.trace_module调用Module的实例并记录Module运行时发生的操作，并创建了torch.jit.TracedModule的实例（TracedModule是其实例）
TorchScript将其定义记录在中间表示(或IR)中，在深度学习中通常称为graph。可以用.graph属性检查graph的详细结果：
print(traced_cell.graph)
打印结果如下：
graph(%self.1 : __torch__.TestTraceCell,
%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]() # test_torchscript.py:9:0
%12 : Float(3, 4, strides=[4, 1], requires_grad=1, device=cpu) = aten::add(%20, %h, %11) # test_torchscript.py:9:0
%13 : Float(3, 4, strides=[4, 1], requires_grad=1, device=cpu) = aten::tanh(%12) # test_torchscript.py:9: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)
这是一个非常底层的表示形式，graph中包含的大多数信息对最终用户没有用。为此，可以使用.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)
可以看出打印结果和真实代码中定义的forward一样。
那这样做有什么好处呢？ 有以下几个原因：
-
TorchScript代码可以在其自己的解释器中调用，该解释器基本上是受限制的Python解释器。该解释器不被全局解释器锁定，因此可以在同一实例上同时处理许多请求。
-
这种格式使开发人员可以将整个模型保存到磁盘上，并将其加载到另一个环境中。比如以Python以外的语言编写的服务。
-
TorchScript提供了一种表示形式，并可以对其代码进行编译器优化以提供更有效的运行。
-
TorchScript提供多种后端或设备runtime的接口
可以看到，调用traced_cell产生的结果与Python模块相同：
print(test_cell(x, h))
print(traced_cell(x, h))
打印结果如下：
(tensor([[0.9567, 0.6879, 0.2618, 0.7927],
[0.8227, 0.7464, 0.4165, 0.5366],
[0.8193, 0.1679, 0.8132, 0.9052]], grad_fn=), tensor([[0.9567, 0.6879, 0.2618, 0.7927],
[0.8227, 0.7464, 0.4165, 0.5366],
[0.8193, 0.1679, 0.8132, 0.9052]], grad_fn=))
(tensor([[0.9567, 0.6879, 0.2618, 0.7927],
[0.8227, 0.7464, 0.4165, 0.5366],
[0.8193, 0.1679, 0.8132, 0.9052]], grad_fn=), tensor([[0.9567, 0.6879, 0.2618, 0.7927],
[0.8227, 0.7464, 0.4165, 0.5366],
[0.8193, 0.1679, 0.8132, 0.9052]], grad_fn=))
script 模式
在介绍 script 模式之前，先来看以下示例：
import torch
torch.manual_seed(42)
class TestDecisionGate(torch.nn.Module):
def forward(self, x):
if x.sum() > 0:
return x
else:
return -x
class TestCell(torch.nn.Module):
def __init__(self, dg):
super(TestCell, 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
test_cell = TestCell(TestDecisionGate())
x, h = torch.rand(3, 4), torch.rand(3, 4)
traced_cell = torch.jit.trace(test_cell, (x, h))
print(traced_cell.code)
运行结果如下：
test_torchscript.py:14: 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:
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分支！ 为什么？trace完全按照之前所说的去做：运行代码，记录发生的操作，并构造一个可以做到这一点的ScriptModule。在这个过程中诸如控制流之类的东西被抹去了。
那么如何在TorchScript中真实地表示此模块？Pytorch提供了一个脚本编译器(script compiler)， 它可以直接分析Python源代码以将其转换为TorchScript。
script 模式和trace 模式的使用方式很接近， 但是实现原理却大相径庭。TorchScript 实现了一个完整的编译器以支持 script 模式。保存模型阶段对应着编译器的前端(语法分析、语义分析、类型检查、中间代码生成)。即，对code进行词法分析、语法分析、句法分析，形成AST树(抽象语法树)，最后再将AST树线性化。在保存模型时，TorchScript 编译器解析 Python 代码，并构建代码的AST。script 模式在的限制比较小，不仅支持 if-else 等控制流，还支持非 Tensor 操作，如 List、Tuple、Map 等容器操作。script相当于一个嵌入在Python/Pytorch的DSL，其语法只是pytorch语法的子集，这意味着存在一些op和语法script不支持，这样在编译的时候就会遇到问题。此外，script的编译优化方式更像是CPU上的传统编译优化，重点对于图进行硬件无关优化，并对IF、loop这样的statement进行优化。
使用脚本编译器转换TestDecisionGate代码如下所示：
scripted_gate = torch.jit.script(TestDecisionGate())
test_cell = TestCell(scripted_gate)
scripted_cell = torch.jit.script(test_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(test_cell(x, h))
print(scripted_cell(x, h))
打印结果如下：
(tensor([[ 0.1018, -0.2095, 0.9220, 0.8287],
[ 0.4238, -0.0612, 0.7071, 0.8443],
[-0.2990, 0.4254, 0.7227, 0.8770]], grad_fn=), tensor([[ 0.1018, -0.2095, 0.9220, 0.8287],
[ 0.4238, -0.0612, 0.7071, 0.8443],
[-0.2990, 0.4254, 0.7227, 0.8770]], grad_fn=))
(tensor([[ 0.1018, -0.2095, 0.9220, 0.8287],
[ 0.4238, -0.0612, 0.7071, 0.8443],
[-0.2990, 0.4254, 0.7227, 0.8770]], grad_fn=), tensor([[ 0.1018, -0.2095, 0.9220, 0.8287],
[ 0.4238, -0.0612, 0.7071, 0.8443],
[-0.2990, 0.4254, 0.7227, 0.8770]], grad_fn=))
混合Scripting和Tracing
在某些情况下，需要使用Tracing而不是Scripting，例如，module中有许多决策，这些决策是基于Python中固定值做出的，而我们又不希望在TorchScript中出现。在这种情况下，可以通过tracing来组成scripting：torch.jit.script将内联被traced module 的代码，而tracing将内联被scripted module 的代码。
示例1：
import torch
torch.manual_seed(42)
class TestTraceCell(torch.nn.Module):
def __init__(self):
super(TestTraceCell, 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
class TestDecisionGate(torch.nn.Module):
def forward(self, x):
if x.sum() > 0:
return x
else:
return -x
class TestCell(torch.nn.Module):
def __init__(self, dg):
super(TestCell, 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
x, h = torch.rand(3, 4), torch.rand(3, 4)
scripted_gate = torch.jit.script(TestDecisionGate())
class TestRNNLoop(torch.nn.Module):
def __init__(self):
super(TestRNNLoop, self).__init__()
self.cell = torch.jit.trace(TestCell(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(TestRNNLoop())
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)
示例2：
class WrapRNN(torch.nn.Module):
def __init__(self):
super(WrapRNN, self).__init__()
self.loop = torch.jit.script(TestRNNLoop())
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)
模型保存和加载
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=TestRNNLoop
(cell): RecursiveScriptModule(
original_name=TestCell
(dg): RecursiveScriptModule(original_name=TestDecisionGate)
(linear): RecursiveScriptModule(original_name=Linear)
)
)
)
def forward(self,
xs: Tensor) -> Tensor:
loop = self.loop
_0, y, = (loop).forward(xs, )
return torch.relu(y)
LibTorch使用示例
要在 C++ 中加载序列化的 PyTorch 模型TorchScript格式则必须依赖于 PyTorch C++ API，也称为 LibTorch。运行模型阶段对应着编译器后端(代码优化、目标代码生成、目标代码优化)。LibTorch 还实现了一个可以运行该编译器所生成代码的解释器。在运行代码时，在 LibTorch 中，AST 被加载， 在进行一系列代码优化后生成目标代码(并非机器码)，然后由解释器运行。LibTorch 发行版包含共享库，头文件和 CMake 构建配置文件的集合。以下将使用 CMake 和 LibTorch 构建一个 C++示例。
编写示例程序
从官网下载LibTorch。
C++ 示例代码如下(仅仅是打印一个随机数)：
// example-app.cpp
#include
#include
int main() {
torch::manual_seed(42);
// torch::Tensor tensor = torch::rand({2, 3}); //cpu版
// std::cout << tensor << std::endl; //cpu版
torch::Device device(torch::kCUDA); //GPU版
torch::Tensor tensor = torch::rand({2, 3}).to(device);
std::cout << tensor << std::endl;
}
编写CMakeLists.txt
本示例的CMakeLists.txt如下：
cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(example-app)
find_package(Torch REQUIRED)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
add_executable(example-app example-app.cpp)
target_link_libraries(example-app "${TORCH_LIBRARIES}")
set_property(TARGET example-app PROPERTY CXX_STANDARD 14)
注意：
本示例使用C++14，而使用C++14需要的gcc版本不能太低！
编译和运行
mkdir build
cd build
cmake -DCMAKE_PREFIX_PATH=/absolute/path/to/libtorch ..
cmake --build . --config Release
本示例中目录结构如下所示：
example-app/
├── CMakeLists.txt
├── build
├── example-app.cpp
GPU 上运行结果：
0.8823 0.9150 0.3829
0.9593 0.3904 0.6009
[ CUDAFloatType{2,3} ]
CPU 上运行结果：
0.8823 0.9150 0.3829
0.9593 0.3904 0.6009
[ CPUFloatType{2,3} ]
可以看出，CPU的运行结果与GPU运行结果相同，符合预期。
小结
本文主要介绍推理加速格式TorchScript的一些基本概念和简单使用方法。后续会基于resnet18模型进一步评测，敬请期待。