👇踩坑指南——onnx系列

Posted 蟹老板的汉堡店-蟹堡王

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了👇踩坑指南——onnx系列相关的知识,希望对你有一定的参考价值。

模型推理加速系列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 sizebaselinejitONNX
11870220036.2
8861
161160

因为这款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 sizebaselinejitONNX
15.642
85.54.42.6
166.04.12.7
325.44.23.6
646.56.46.7
12811.110.812.7
25622.321.825.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

以上是关于👇踩坑指南——onnx系列的主要内容,如果未能解决你的问题,请参考以下文章

深入浅出SpringCloud原理及实战「SpringCloud-Alibaba系列」微服务模式搭建系统基础架构实战指南及版本规划踩坑分析

MACOS arco design 初学者踩坑指南

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

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

ARouter踩坑指南

ARouter踩坑指南