【模型推理加速系列】04：BERT加速方案对比 TorchScript vs. ONNX

简介

本文以 BERT-base 的为例，介绍2种常用的推理加速方案：ONNX 和 TorchScript，并实测对比这两种加速方案与原始Pytorch模型格式的inference性能。本文通过实测发现：

(1)GPU环境中，小batch size场景下，ONNX表现最好；

(2)GPU环境中，大batch size场景下，TorchScript 性能最好；

(3)CPU(AMD)环境中，ONNX性能最好；

PS：

以上结论仅基于笔者自己的测试环境，具体测试细节请参考推理速度评测部分。
完整实验代码请在微信公众号：小窗幽记机器学习 内回复【bert】获取下载链接。

基于ONNX

导出ONNX模型

关于BERT如何基于ONNX做推理加速可以参考之前的两篇小作文：如何用ONNX加速BERT特征抽取 和 如何用ONNX加速BERT特征抽取-part2。本文在此只再次强调Pytorch模型导出ONNX模型过程中的2个关键点：

(1) 创建符合数据规范的伪输入(dummy inputs)，并利用伪输入在模型中做一次前向inference。Pytorch模型转ONNX调用的接口是torch.onnx.export。该接口支持2种模式：tracing 和 scripting，默认模式是基于tracing(追踪)。本文这里导出BERT的ONNX格式也是使用tracing模式。在这个过程中会tracing(追踪)tensor的相关操作形成静态图。tracing方式不会捕获根据输入数据而改变的行为，比如对于if语句，只会记录执行的那一条分支，对于for循环的次数，则会导出与跟踪运行时完全相同的静态图。因此，如果待导出的模型推理阶段(forward函数)中存在判断和非固定循环次数时，Tracing模式并不适用。对于这类使用动态控制流导出模型，则需要使用Scripting模式导出，即使用torch.jit.script。

(2) 选择和设置输入和输出的name及其动态轴。需要特别注意的是设置input_names的顺序需要与实际相符合。比如BERT-base中文版这里的模型，如果变动input_names=['input_ids', 'attention_mask', 'token_type_ids']中的参数名顺序会有影响吗？此外，正常业务场景下，batch size 一般是变长的，所以会将输入参数中的第0维设置为动态。

示例代码

BERT 导出 ONNX具体示例代码如下：

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   convert_2_onnx.py
@Time    :   2022/11/05 10:10:16
@Author  :   Jason
@Desc    :   Pytorch模型转为ONNX格式
'''
import torch
from transformers import AutoModel
model_path = "/share/model_zoo/bert-base-chinese/"
MODEL_ONNX_PATH = "/share/result_models/ner/bert/onnx/raw_bert_dynamic.onnx"
OPERATOR_EXPORT_TYPE = torch._C._onnx.OperatorExportTypes.ONNX
model = AutoModel.from_pretrained(model_path)
model.eval()

def make_train_dummy_input(seq_len):
    org_input_ids = torch.LongTensor([[i for i in range(seq_len)]])
    org_token_type_ids = torch.LongTensor([[1 for i in range(seq_len)]])
    org_input_mask = torch.LongTensor([[0 for i in range(int(seq_len/2))] + [1 for i in range(seq_len - int(seq_len/2))]])
    return (org_input_ids, org_token_type_ids, org_input_mask)
    
org_dummy_input = make_train_dummy_input(int(args.seq_len))
output = torch.onnx.export(model,
                           org_dummy_input,
                           MODEL_ONNX_PATH,
                           verbose=True,
                           operator_export_type=OPERATOR_EXPORT_TYPE,
                           opset_version=12,
                           input_names=['input_ids', 'attention_mask', 'token_type_ids'],  # 需要注意顺序！不可随意改变, 否则结果与预期不符
                           output_names=['last_hidden_state', 'pooler_output'],  # 需要注意顺序, 否则在推理阶段可能用错output_names
                           do_constant_folding=True,
                           dynamic_axes={"input_ids": {0: "batch_size", 1: "length"},
                                         "token_type_ids": {0: "batch_size", 1: "length"},
                                         "attention_mask": {0: "batch_size", 1: "length"},
                                         "pooler_output": {0: "batch_size"},
                                         "last_hidden_state": {0: "batch_size"}}
                           )
print("Export of {} complete!".format(MODEL_ONNX_PATH))

PS：

上述代码设置输入的文本长度和batch size 都是变长，即无需对输入做padding和攒到固定batch size，再做inference。

基于TorchScript

JIT

在介绍TorchScript之前，先介绍下JIT。因为Pytorch中TorchScript的使用是在torch.jit这个模块下。JIT 是一种概念，全称是 Just In Time Compilation，中文译为即时编译，是一种程序优化的方法。

在深度学习中 JIT 的思想更随处可见，最明显的就是 Keras 框架的 model.compile。PyTorch 从面世以来一直以易用性著称(使用eager模式设计)，使用最贴合原生 Python 的开发方式。这得益于 PyTorch 的动态图结构。开发者可以在 PyTorch 的模型前向中加任何 Python 的流程控制语句， 甚至是断点单步跟进都不会有任何问题，但是如果是 TensorFlow，则需要使用 tf.cond 等 TensorFlow 自己开发的流程控制，甚是麻烦。

动态图模型通过牺牲一些高级特性来换取易用性，那到底 JIT 有哪些特性，在什么情况下不得不用到 JIT 呢？PyTorch JIT 模式就是一种将原本eager模式的表达转变并固定为一个计算图，便于进行优化和序列化。那TorchScript(PyTorch 的 JIT 实现) 到底带来了哪些好处？主要有以下3点：

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

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

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负责解析运行这个模型文件。如果使用过Triton的小伙伴应该还记得配置文件中backend中一个选项是libtorch，正是对应这里的TorchScript模型格式的部署。

TorchScript 保存模型有两种模式：trace 模式和 script 模式。本文将BERT以trace模式导出 TorchScript。

PS：目前以 script 模式导出BERT尚未成功，big sad～欢迎导出成功的小伙伴指教。

示例代码

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   convert_2_jit.py
@Time    :   2022/10/31 20:10:41
@Author  :   Jason
@Desc    :   Pytorch模型转为TorchScript格式
'''

import torch
import pdb
from transformers import BertModel, BertTokenizer, BertConfig, AutoTokenizer
device = "cuda"
is_use_gpu = False

model_path = "/share/model_zoo/bert-base-chinese/"
test_text = ["我是卖铁观音的小男孩,毕业于华中科技大学"]
tokenizer = BertTokenizer.from_pretrained(model_path)
inputs = tokenizer(test_text, max_length=32, return_tensors="pt")
config = BertConfig.from_pretrained(model_path, cache_dir=None, torchscript=True, use_cache=False)
model = BertModel.from_pretrained(model_path, config=config, cache_dir=None, )
if is_use_gpu:
    model = model.to(device)
    inputs = inputs.to(device)
model.eval()
print("inputs=", inputs)
temp = (inputs['input_ids'], inputs['attention_mask'], inputs['token_type_ids'], )
traced_model = torch.jit.trace(model, temp)
if is_use_gpu:
    torch.jit.save(traced_model, "traced_bert_gpu.pt")
else:
    torch.jit.save(traced_model, "traced_bert_cpu.pt")

推理速度评测

以下分享不同加速方案、在不同batch size下 在CPU和GPU上的推理耗时。

CPU

评测所使用CPU 型号：

128 AMD EPYC 7713 64-Core Processor

AMD EPYC 7713 型号CPU 不同batch size大小下的平均耗时(单位ms)，即单个batch size的推理耗时(如果想要得到单个样本的推理耗时只需要再除以batch size即可)：

batch size	baseline	jit	ONNX
1	1870	2200	36.2
8	–	–	861
16	–	–	1160

因为这款AMD型号的CPU在baseline 和 jit 上的推理性能实在是太差，放弃在baseline 和jit上的进一步评测。但是，可以得出一个结论，ONNX格式在该型号CPU上性能仍然较好，从而侧面看出ONNX模型格式在特殊场景下仍然具有一定优越性。

GPU

评测所使用GPU型号：

NVIDIA GeForce RTX 3090 GPU

NVIDIA GeForce RTX 3090 型号GPU 不同batch size大小下的平均耗时(单位ms)，即单个batch size的推理耗时(如果想要得到单个样本的推理耗时只需要再除以batch size即可)：

batch size	baseline	jit	ONNX
1	5.6	4	2
8	5.5	4.4	2.6
16	6.0	4.1	2.7
32	5.4	4.2	3.6
64	6.5	6.4	6.7
128	11.1	10.8	12.7
256	22.3	21.8	25.2

基于上述实验，可以得出以下结论：

• batch size 较小时，ONNX性能显著优于其他方案
• batch size 较大时， TorchScript格式(jit) ，显著优于其他方案，特别是ONNX方案。从而可以看出，ONNX不太适合于大 batch size场景。

需要特别强调的是，上述实验结果基于笔者特定软硬件环境评测所得。不同环境下的评测结果结论可能不同，欢迎小伙伴们基于自己的机器进行实测后分享实验结果。

附录

硬件信息：

CPU: AMD EPYC 7713(2个物理CPU，每个物理CPU中core的个数是64，逻辑CPU的个数是128)
GPU: NVIDIA GeForce RTX 3090(单卡)

核心包的版本：

onnxruntime-gpu Version: 1.13.1
transformers Version: 4.21.3
torch Version: 1.13.0a0+08820cb