TensorRT模型推理加速实践

zhcs333

发布于 2024-7-18 09:52

浏览

0收藏

一、TensorRT简介

TensorRT是由C++、CUDA、python三种语言编写成的库，有助于在 NVIDIA GPU上进行高性能推理。基于目前主流的深度学习框架得到的模型都可通过TensorRT实现推理加速。

TensorRT模型推理加速实践-AI.x社区

图1 TensorRT转换过程

2021年7月，NVIDIA 发布了 TensorRT 8.0版本，新版本极大提升了Transformers结构的推理新能。TensorRT性能优化主要依赖下面两种方式：

1、权重与激活精度校准：在推理中使用FP16或者INT8精度计算，通过降低计算精度，提高计算效率，实现更低的显存占用率和延迟以及更高的吞吐率。结合TensorRT提供的完全自动的校准（Calibration）过程，减少精度降低带来的性能损耗。

2、层与张量融合：通过Kernel纵向融合、Kernel横向融合、消除concatenation层这三种方式，实现层与张量的融合，TensorRT可以获得更小、更快、更高效地计算流图，其拥有更少层网络结构以及更少Kernel Launch次数，以提高 GPU 的计算效率。

目前TensorRT已广泛应用于国内外各大厂家，如微信也使用TensorRT来进行搜索加速。

二、TensorRT实践

1.环境部署

TensorRT是必须依赖GPU的环境的，要根据自己本地的CUDA版本来确认所需要安装的TensorRT版本（这里面有EA抢先版和GA通用版本，建议下载GA通用版本，稳定一点），笔者本地配置为

硬件（GPU）：T4
软件：tensorRT8.0.1.6 + cuda11.4 + pycuda2020.1 + torch1.12.1+cu113

可供参考，具体可参考官网[4]。

2.ONNX转TRT

onnx模型转成TRT模型的办法有很多，这里介绍简单介绍两种常用方式：

（1）使用TensorRT自带的trtexec命令

通过下面的脚本可直接进行转换，这里面主要涉及自有的基于bert实现的文本分类模型（负面情感等级分析），原训练后得到模型已转换为onnx模型，大家可以根据自己实际的模型进行相应地修改。

1../trtexec   
2.--notallow="pytorch_model_fp16.onnx"    
3.--saveEngine="bert.trt"   
4.--fp16   
5.--minShapes=input_ids:1x1,attention_mask:1x1,token_type_ids:1x1   
6.--optShapes=input_ids:1x256,attention_mask:1x256,token_type_ids:1x256  
7.--maxShapes=input_ids:1x256,attention_mask:1x256,token_type_ids:1x256 
8.--device=1

（2）通过python API这种形式转换，首先以trt的Logger为参数，使用builder创建计算图类型，然后使用onnx等网络框架下的结构填充计算图，最后由计算图创建cuda环境下的引擎，并以序列化的形式写入到文件中，方便后续推理过程。特别的，这里面要注意输入名要与onnx的构造输入名字要保持一致，具体代码如下：

TensorRT模型推理加速实践-AI.x社区

图2 onnx转TensorRT代码

这两种方式生成的trt模型的效果都是一样的，但第一种方式由于要对输入模型和本地环境做一下额外的校验，通过python API这种初始化设定参数的方式要快很多，两种方式都需要添加最小输入、常规输入、最大输入，以保证模型能够动态地处理最大长度以内的句子。

3.TensorRT推理

通过TensorRT的模型转换后，外部训练好的模型都被tensorRT统一成tensorRT可识别的engine文件（并优化过）。在推理时，只要通过TensorRT的推理SDK就可以完成推理。具体的推理过程如下：

TensorRT模型推理加速实践-AI.x社区

图3 TensorRT推理流程

代码实现如下：

TensorRT模型推理加速实践-AI.x社区

图4 模型推理代码

说明：

通过TensorRT运行时，加载转换好的engine。

推理前准备：

1）在CPU中处理好输入（数据读入、标准化、Bert分词等）；

2）利用TensorRT的推理SDK中common模块进行输入和输出GPU显存分配。

执行推理：

1）将CPU的输入拷贝到GPU中；

2）在GPU中进行推理，并将模型输出放入GPU显存中。

推理后处理：

1）、将输出从GPU显存中拷贝到CPU中；

2）在CPU中进行其他后处理。

4.实验验证

下图是TensorRT在公开模型推理加速的实际表现，采用的是先进的GPU卡V100,较其他方式，在保证同等低延迟的情况下，结合tensorRT使得模型整体吞吐量提升了几十倍，着实惊艳。

TensorRT模型推理加速实践-AI.x社区

图5 基于公开模型TensorRT的表现

笔者也对 TensorRT的效果进行了本地的性能验证，为了保证本地环境与实际线上保持一致，这里采用了docker镜像的方式进行实际的封装。最后我们比较了几个模型的整体效果和性能，TensorRT的整体性能大约能提升6-7倍，与 BERT 原模型40QPS 的性能相比，TensorRT 整体性能达到了258QPS，推理速度提升极大。

表1 基于负面情感分析模型的性能比较

模型	准确率（%）	性能（QPS）
bert原模型	89.05	40
onnx模型	89.05	63
onnx+fp16模型	89.05	185
onnx+fp16+opt模型	89.05	214
tensorRT+fp16模型	89.15	258