模型推理加速系列bert加速方案对比torchscriptvs.onnx(代码片段)

JasonLiu1919 JasonLiu1919     2022-12-12     776

关键词:

文章目录

简介

本文以 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

模型推理加速系列07:以bert为例全面评测各种推理加速方案(代码片段)

...李白风格创作的,仅供欣赏。前文以CV领域中的resnet18模型为例综合评测各种推理加速方 查看详情

模型推理加速系列07:以bert为例全面评测各种推理加速方案(代码片段)

...李白风格创作的,仅供欣赏。前文以CV领域中的resnet18模型为例综合评测各种推理加速方 查看详情

模型推理加速系列06:基于resnet18加速方案评测(代码片段)

...称重师:卖麻辣烫的小男孩。今天这篇文章以resnet18模型为例,对比Pytorch、ONNX、TorchScript、TensorRT模型格式在不同硬件(包括CPU和GPU)上的inference性能。由于此前TorchScript模型在AMDCPU上的评测结果是负向效果(远慢于Pytorch)ÿ... 查看详情

模型推理加速系列06:基于resnet18加速方案评测(代码片段)

...称重师:卖麻辣烫的小男孩。今天这篇文章以resnet18模型为例,对比Pytorch、ONNX、TorchScript、TensorRT模型格式在不同硬件(包括CPU和GPU)上的inference性能。由于此前T 查看详情

模型推理一文看懂winograd卷积加速算法(代码片段)

...篇卷积加速相关的文章,感兴趣的同学可以查阅《【模型推理】一文看懂Img2Col卷积加速算法》、《【模型推理】一文看懂GoogleTPU脉动阵列加速卷积计算原理》、《【模型推理】谈谈为什么 查看详情

模型推理谈谈为什么量化能加速推理(代码片段)

... o_O O_o ~_~ o_O 本文主要讨论一下为什么量化能加速模型推理。 前面已经写过几篇关于模型量化相关的文章:《【模型推理】谈谈几种量化策略:MinMax、KLD、ADMM、EQ》、《【模型推理】谈谈模型量化组织方式》、... 查看详情

bert/transformer模型加速

...ontext信息编码或长时依赖建模能力,能够使图像/视觉模型具备更强的特征表达能力、与领域适应性。因此在计算机视觉领域(CV)& 查看详情

如何加速 Tensorflow 2 keras 模型进行推理?

】如何加速Tensorflow2keras模型进行推理?【英文标题】:HowtospeedupTensorflow2kerasmodelforinference?【发布时间】:2020-03-0420:50:55【问题描述】:所以现在有一个重大更新,从TensorFlow1.X迁移到2.X。在TF1.X中,我习惯了一个管道,它帮助我... 查看详情

模型推理谈谈为什么卷积加速更喜欢nhwclayout(代码片段)

 本文主要讨论一下为什么卷积加速更加喜欢NHWC的数据排布。 我目前接触过的数据排布类型(主要针对卷积)有NCHW(pytorch、caffe),NHWC(Tensorflow,也是TVMGPU和寒武纪MLUCore上更喜欢的dataLayout),CHW(TensorRT里不考虑动态batch... 查看详情

模型推理一文看懂googletpu脉动阵列加速卷积计算原理

 本教程详细解释了GoogleTPU脉动阵列加速卷积计算原理。 TPU中计算卷积的方式和GPU不同,主要是依靠一种称为“脉动阵列”的硬件电路结构来实现的。脉动阵列的主体部分是一个二维的滑动阵列,其中每一个节点都是... 查看详情

cpu推理|使用英特尔sapphirerapids加速pytorchtransformers(代码片段)

...我们将重点关注推理。使用基于PyTorch的HuggingFacetransformers模型,我们首先在IceLake服务器上分别测量它们在长、短两种文本序列上的性能。然后,我们在SapphireRapids服务器和最新版本的HuggingFaceOptimumIntel上执行相同的测试,并比较... 查看详情

tensorrt模型加速1-输入输出部署流程

...TensorRT的输入、输出以及部署流程,了解TensorRT在部署模型中起到的作用。然后介绍TensorRT模型导入流程,针对不同的深度学习框架,使用不同的方法导入模型。一、TensorRT简介TensorRT是NVIDIA公司发布的一个高性能的深... 查看详情

tensorrt模型加速|网络结构优化|低精度推理

...张量融合”实现,低精度推理通过使用FP16或INT8进行模型推理实现。一、低精度推理背景:神经网络在训练时采用高精度保存参数,一般采用 32 位浮点数(Floating-Point, FP32),因此模型最后的 weights 也... 查看详情

tensorrt模型加速|网络结构优化|低精度推理

...张量融合”实现,低精度推理通过使用FP16或INT8进行模型推理实现。一、低精度推理背景:神经网络在训练时采用高精度保存参数,一般采用 32 位浮点数(Floating-Point, FP32),因此模型最后的 weights 也... 查看详情

rk3399proyolov5人工智能加速推理解决方案(代码片段)

...;本文介绍了YOLOv5s算法在国产瑞芯微电子RK3399ProD上的部署推理.介绍整个的流程,并基于RK3399Pro简单的介绍下RKNN的Python接口使用,并记录一些踩过的坑。本文仅做交流使用,文中有什么理解的不到位的或者谬误,... 查看详情

量化感知训练实践:实现精度无损的模型压缩和推理加速(代码片段)

简介:本文以近期流行的YOLOX[8]目标检测模型为例,介绍量化感知训练的原理流程,讨论如何实现精度无损的实践经验,并展示了量化后的模型能够做到精度不低于原始浮点模型,模型压缩4X、推理加速最高2.3X... 查看详情

tensorrt模型加速2-优化方式

...张量融合”实现,低精度推理通过使用FP16或INT8进行模型推理实现。一、低精度推理背景:神经网络在训练时采用高精度保存参数,一般采用 32 位浮点数(Floating-Point, FP32),因此模型最后的 weights 也... 查看详情

两行代码自动压缩vit模型!模型体积减小3.9倍,推理加速7.1倍

Transformer模型及其变体,因其更优的注意力机制能力和长时依赖等特性,已成为自然语言处理(NLP)、语音识别(ASR)、计算机视觉(CV)等领域的主流序列建模结构。根据PaperWithCode网站ObjectDetectiononCOCOtest-dev专栏数据,检测... 查看详情