【模型推理加速系列】07: 以BERT为例全面评测各种推理加速方案

简介

冬至夜月圆，霜露点清银。枫叶落尽霜，寒风凛冽水。山河冰封路，长夜独酌愁。火炉照暖家，窗外雪花飞。小伙伴们好，我是微信公众号小窗幽记机器学习的首席试药师：卖布洛芬的小男孩。上面这首诗是最近顶流ChatGPT模仿李白风格创作的，仅供欣赏。前文以CV领域中的resnet18模型为例综合评测各种推理加速方案，具体评测结果可以参考之前的小作文：模型推理加速系列 ｜ 06: 基于resnet18加速方案评测。今天这篇小作文尝试以NLP领域中的常用模型BERT为例(仅将输入文本进行encode)，综合评测包括Pytorch、ONNX、JIT、TensorRT和OpenVino在内这5种推理方案的性能。

更多、更新文章欢迎关注微信公众号：小窗幽记机器学习。后续会持续输出模型推理加速和工程部署相关系列，敬请期待～

模型导出

导出 ONNX

如何将Pytorch版BERT导出为ONNX格式可以参考之前的文章：模型推理加速系列｜04：BERT模型推理加速 TorchScript vs. ONNX。

导出 TorchScript

如何将Pytorch版BERT导出为TorchScript格式可以参考之前的文章：模型推理加速系列｜04：BERT模型推理加速 TorchScript vs. ONNX，更多关于TorchScript模型格式的介绍可以参考之前的文章：模型推理加速系列｜05：TorchScript模型格式简介及其使用。

导出 TensorRT

在TensorRT这部分依次评测变长和定长版(指文本输入长度上的限制)，所以在导出TensorRT模型格式的时候，导出支持变长和定长文本输入的两个版本。

导出输入文本定长版 TensorRT engine格式模型示例如下：

CUDA_VISIBLE_DEVICES=0 trtexec --onnx=/home/model_zoo/nlp/onnx/bert-base-chinese/bert_dynamic.onnx  --minShapes=input_ids:1x32,attention_mask:1x32,token_type_ids:1x32 --optShapes=input_ids:64x32,attention_mask:64x32,token_type_ids:64x32 --maxShapes=input_ids:128x32,attention_mask:128x32,token_type_ids:128x32 --saveEngine=/home/model_zoo/nlp/tensorrt/bert_static_32.engine

PS：

以上导出的模型在batch size维度上的动态的，但是只支持长度为32的输入文本。

导出输入文本变长版 TensorRT engine格式模型示例如下：

CUDA_VISIBLE_DEVICES=0 trtexec --onnx=/home/model_zoo/nlp/onnx/bert-base-chinese/bert_dynamic.onnx --minShapes=input_ids:1x8,attention_mask:1x8,token_type_ids:1x8 --optShapes=input_ids:64x32,attention_mask:64x32,token_type_ids:64x32 --maxShapes=input_ids:128x64,attention_mask:128x64,token_type_ids:128x64 --saveEngine=/home/model_zoo/nlp/tensorrt/bert_dynamic_max64.engine

导出 OpenVino

使用的镜像是openvino/ubuntu20_runtime:latest。先用之前Pytorch格式模型转为的ONNX格式模型，再将ONNX格式模型转为OpenVino所需要的格式。这里需要用到mo命令，安装方式如下：pip3 install openvino-dev[all]。将ONNX模型转为openvino的命令如下：

mo --input_model /home/model_zoo/nlp/onnx/bert-base-chinese/bert_dynamic.onnx --input "input_ids[-1 128],attention_mask[-1 128],token_type_ids[-1 128]" --output_dir /home/model_zoo/nlp/openvino/bert/ --data_type FP32

评测结果

本次实验硬件信息：

GPU：1张 Nvidia T4

CPU：10 Intel® Xeon® Platinum 8255C CPU @ 2.50GHz

以下实验结果是batch size的推理耗时，单位是ms

输入文本长度为 8

CPU：

batch size	Pytorch	ONNX	JIT	OpenVino(定长版)	OpenVino(变长版)
1	24.5	10.1	18.8	12.1	12.5
2	46.7	12.0	35.5	14.7	16.2
4	53.9	18.0	40.9	25.4	25.5
8	65.2	30.3	50.6	35.0	36.8
16	82.6	53.6	65.3	66.1	66.2
32	126.3	99.1	105.9	112.7	112.1
64	236.2	194.2	192.9	198.5	200.1
128	474.9	368.2	360.3	362.1	364.6

GPU：

batch size	Pytorch	ONNX	JIT	TensorRT(定长版)	TensorRT(变长版)
1	10.7	6.5	8.3	4.8	7.9
2	11.6	6.6	8.3	5.1	7.9
4	10.9	7.1	7.2	5.5	8.0
8	11.6	10.0	11.3	7.8	8.0
16	15.8	14.9	15.4	13.0	20.9
32	24.3	23.8	23.9	20.6	75.3
64	44.6	42.6	44.4	37.2	287.9
128	86.5	87.6	83.9	73.5	568.1

输入文本长度为 16

CPU：

batch size	Pytorch	ONNX	JIT	OpenVino(定长版)	OpenVino(变长版)
1	46.1	12.2	35.5	14.6	16.3
2	52.7	18.1	41.9	21.8	24.2
4	64.1	30.0	51.9	31.9	38.6
8	81.1	52.7	68.0	52.5	66.2
16	125.4	97.4	111.1	102.6	112.4
32	235.2	190.8	216.7	185.8	197.1
64	430.8	362.8	352.6	340.7	366.3
128	754.7	725.4	769.1	620.1	718.5

GPU：

batch size	Pytorch	ONNX	JIT	TensorRT(定长版)	TensorRT(变长版)
1	11.0	6.6	8.0	5.1	21.1
2	11.2	7.1	7.2	5.5	21.1
4	11.8	10.2	11.4	7.9	21.4
8	16.2	15.2	15.8	13.5	21.9
16	25.1	24.7	24.8	20.9	20.8
32	45.9	43.7	45.4	38.0	74.3
64	88.8	90.1	87.2	75.8	286.9
128	174.2	174.8	170.8	148.6	571.2

TensorRT 模型是定长输入，如果使用变长版TensorRT，但是输入的文本长度设置为16。

输入文本长度为 32

CPU：

batch size	Pytorch	ONNX	JIT	OpenVino(定长版)	OpenVino(变长版)
1	45.8	18.0	39.7	21.8	22.7
2	64.7	29.8	51.1	32.2	32.4
4	80.6	52.9	65.6	52.7	52.0
8	129.4	96.8	115.2	95.4	92.6
16	244.9	191.7	213.5	178.8	196.3
32	420.1	363.7	362.3	312.2	318.2
64	794.5	733.0	701.9	636.4	713.7
128	1648.0	1477.6	1501.5	1281.9	1441.1

GPU：

batch size	Pytorch	ONNX	JIT	TensorRT(定长版)	TensorRT(变长版-maxlen512)
1	10.6	7.5	7.2	5.5	74.6
2	12.0	10.5	11.7	7.9	74.9
4	16.3	15.6	16.0	13.5	75.4
8	24.8	24.7	24.8	21.8	76.4
16	44.9	44.1	45.6	37.3	74.9
32	86.5	89.4	86.6	74.5	74.4
64	173.3	174.6	170.7	146.1	288.0
128	326.9	346.8	322.3	285.6	1188.4

输入文本长度为 64

CPU：

batch size	Pytorch	ONNX	JIT	OpenVino(定长版)	OpenVino(变长版)
1	56.7	30.2	49.4	32.1	32.6
2	72.2	53.3	72.0	52.8	53.3
4	112.4	98.3	105.8	93.0	93.2
8	218.4	197.1	205.6	208.9	177.3
16	402.5	367.9	392.9	310.0	391.7
32	764.8	739.3	752.7	653.2	711.3
64	1531.4	1488.3	1442.0	1356.5	1580.3
128	3244.5	3139.7	2883.6	3063.2	3133.2

GPU：

batch size	Pytorch	ONNX	JIT	TensorRT(定长版)	TensorRT(变长版-maxlen512)	TensorRT(变长版-maxlen64)
1	11.8	10.5	11.5	8.0	287.7	287.4
2	16.3	15.4	16.1	13.7	287.8	287.5
4	25.3	24.6	25.3	22.0	288.6	291.6
8	44.9	43.6	46.2	39.4	293.1	291.1
16	88.9	89.8	86.7	74.7	287.3	288.9
32	176.2	174.9	170.5	146.2	287.9	288.9
64	332.7	348.1	322.7	287.3	288.6	288.6
128	675.4	683.1	655.1	569.2	1185.5	574.9

结论

根据上述实验结果可以得出以下结论：

• 对于CPU，短文本和小batch size 场景，ONNX推理速度最快，OpenVnio次之。而在长文本和大batch size的场景，OpenVino推理速度最快
• 对于OpenVino，在短文本和小batch size场景，定长输入和变长输入性能差别不大。随着文本长度增加和batch size增大，定长版OpenVino显著优于变长版OpenVino
• 对于GPU，定长版TensorRT显著优于变长版TensorRT。但是，定长版在实际使用中可能会涉及对输入文本的padding操作，需要在具体使用场景中评测该部分的耗时。

更多、更新文章欢迎关注微信公众号：小窗幽记机器学习。后续会持续输出模型推理加速和工程部署相关系列，敬请期待～