模型推理加速系列05: 推理加速格式TorchScript简介及其应用

Posted 2022-12-11 JasonLiu1919

简介

本文紧接前文：模型推理加速系列｜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 主要有以下几个好处：

1. 模型部署的可移植性
   PyTorch 的 1.0 版本发布的最核心的两个新特性就是 JIT 和 C++ API，这两个特性一起发布不是没有道理的，JIT 是 Python 和 C++ 的桥梁。开发者可以使用 Python 训练模型，然后通过 JIT 将模型转为语言无关的模块， 从而让 C++ 可以非常方便得调用，从此使用 Python 训练模型，使用 C++ 将模型部署到生产环境对 PyTorch 来说成为了一件很容易的事。而因为使用了 C++， 开发者几乎可以把 PyTorch 模型部署到任意平台和设备上：树莓派、ios、android 等。

2. 性能提升
   JIT 既然是为生产部署而生，那免不了在性能上面做了极大优化。如果推断的场景对性能要求高，则可以考虑将模型(torch.nn.Module)转换为 TorchScript Module， 再进行推断。

3. 模型可视化
   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=<TanhBackward0>), 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=<TanhBackward0>))

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=<TanhBackward0>), 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=<TanhBackward0>))
(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=<TanhBackward0>), 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=<TanhBackward0>))

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=<TanhBackward0>), 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=<TanhBackward0>))
(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=<TanhBackward0>), 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=<TanhBackward0>))

混合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 <torch/torch.h>
#include <iostream>

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
[ CUDAFloatType2,3 ]

CPU 上运行结果：

  0.8823  0.9150  0.3829
 0.9593  0.3904  0.6009
[ CPUFloatType2,3 ]

可以看出，CPU的运行结果与GPU运行结果相同，符合预期。

小结

本文主要介绍推理加速格式TorchScript的一些基本概念和简单使用方法。后续会基于resnet18模型进一步评测，敬请期待。