TensorRT模型推理加速實(shí)踐

發(fā)布于 2024-7-18 09:52

瀏覽

0收藏

一、TensorRT簡介

TensorRT是由C++、CUDA、python三種語言編寫成的庫，有助于在 NVIDIA GPU上進(jìn)行高性能推理。基于目前主流的深度學(xué)習(xí)框架得到的模型都可通過TensorRT實(shí)現(xiàn)推理加速。

TensorRT模型推理加速實(shí)踐-AI.x社區(qū)

圖1 TensorRT轉(zhuǎn)換過程

2021年7月，NVIDIA 發(fā)布了 TensorRT 8.0版本，新版本極大提升了Transformers結(jié)構(gòu)的推理新能。TensorRT性能優(yōu)化主要依賴下面兩種方式：

1、權(quán)重與激活精度校準(zhǔn)：在推理中使用FP16或者INT8精度計(jì)算，通過降低計(jì)算精度，提高計(jì)算效率，實(shí)現(xiàn)更低的顯存占用率和延遲以及更高的吞吐率。結(jié)合TensorRT提供的完全自動的校準(zhǔn)（Calibration）過程，減少精度降低帶來的性能損耗。

2、層與張量融合：通過Kernel縱向融合、Kernel橫向融合、消除concatenation層這三種方式，實(shí)現(xiàn)層與張量的融合，TensorRT可以獲得更小、更快、更高效地計(jì)算流圖，其擁有更少層網(wǎng)絡(luò)結(jié)構(gòu)以及更少Kernel Launch次數(shù)，以提高 GPU 的計(jì)算效率。

目前TensorRT已廣泛應(yīng)用于國內(nèi)外各大廠家，如微信也使用TensorRT來進(jìn)行搜索加速。

二、TensorRT實(shí)踐

1.環(huán)境部署

TensorRT是必須依賴GPU的環(huán)境的，要根據(jù)自己本地的CUDA版本來確認(rèn)所需要安裝的TensorRT版本（這里面有EA搶先版和GA通用版本，建議下載GA通用版本，穩(wěn)定一點(diǎn)），筆者本地配置為

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

可供參考，具體可參考官網(wǎng)[4]。

2.ONNX轉(zhuǎn)TRT

onnx模型轉(zhuǎn)成TRT模型的辦法有很多，這里介紹簡單介紹兩種常用方式：

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

通過下面的腳本可直接進(jìn)行轉(zhuǎn)換，這里面主要涉及自有的基于bert實(shí)現(xiàn)的文本分類模型（負(fù)面情感等級分析），原訓(xùn)練后得到模型已轉(zhuǎn)換為onnx模型，大家可以根據(jù)自己實(shí)際的模型進(jìn)行相應(yīng)地修改。

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這種形式轉(zhuǎn)換，首先以trt的Logger為參數(shù)，使用builder創(chuàng)建計(jì)算圖類型，然后使用onnx等網(wǎng)絡(luò)框架下的結(jié)構(gòu)填充計(jì)算圖，最后由計(jì)算圖創(chuàng)建cuda環(huán)境下的引擎，并以序列化的形式寫入到文件中，方便后續(xù)推理過程。特別的，這里面要注意輸入名要與onnx的構(gòu)造輸入名字要保持一致，具體代碼如下：

TensorRT模型推理加速實(shí)踐-AI.x社區(qū)

圖2 onnx轉(zhuǎn)TensorRT代碼

這兩種方式生成的trt模型的效果都是一樣的，但第一種方式由于要對輸入模型和本地環(huán)境做一下額外的校驗(yàn)，通過python API這種初始化設(shè)定參數(shù)的方式要快很多，兩種方式都需要添加最小輸入、常規(guī)輸入、最大輸入，以保證模型能夠動態(tài)地處理最大長度以內(nèi)的句子。

3.TensorRT推理

通過TensorRT的模型轉(zhuǎn)換后，外部訓(xùn)練好的模型都被tensorRT統(tǒng)一成tensorRT可識別的engine文件（并優(yōu)化過）。在推理時，只要通過TensorRT的推理SDK就可以完成推理。具體的推理過程如下：

TensorRT模型推理加速實(shí)踐-AI.x社區(qū)

圖3 TensorRT推理流程

代碼實(shí)現(xiàn)如下：

TensorRT模型推理加速實(shí)踐-AI.x社區(qū)

圖4 模型推理代碼

說明：

通過TensorRT運(yùn)行時，加載轉(zhuǎn)換好的engine。

推理前準(zhǔn)備：

1）在CPU中處理好輸入（數(shù)據(jù)讀入、標(biāo)準(zhǔn)化、Bert分詞等）；

2）利用TensorRT的推理SDK中common模塊進(jìn)行輸入和輸出GPU顯存分配。

執(zhí)行推理：

1）將CPU的輸入拷貝到GPU中；

2）在GPU中進(jìn)行推理，并將模型輸出放入GPU顯存中。

推理后處理：

1）、將輸出從GPU顯存中拷貝到CPU中；

2）在CPU中進(jìn)行其他后處理。

4.實(shí)驗(yàn)驗(yàn)證

下圖是TensorRT在公開模型推理加速的實(shí)際表現(xiàn)，采用的是先進(jìn)的GPU卡V100,較其他方式，在保證同等低延遲的情況下，結(jié)合tensorRT使得模型整體吞吐量提升了幾十倍，著實(shí)驚艷。

TensorRT模型推理加速實(shí)踐-AI.x社區(qū)

圖5 基于公開模型TensorRT的表現(xiàn)

筆者也對 TensorRT的效果進(jìn)行了本地的性能驗(yàn)證，為了保證本地環(huán)境與實(shí)際線上保持一致，這里采用了docker鏡像的方式進(jìn)行實(shí)際的封裝。最后我們比較了幾個模型的整體效果和性能，TensorRT的整體性能大約能提升6-7倍，與 BERT 原模型40QPS 的性能相比，TensorRT 整體性能達(dá)到了258QPS，推理速度提升極大。

表1 基于負(fù)面情感分析模型的性能比較

模型	準(zhǔn)確率（%）	性能（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

三、總結(jié)

不管是在模型精度損失、以及模型推理性能加速上，TensorRT確實(shí)表現(xiàn)亮眼，如想要更進(jìn)一步提升性能，可以嘗試int8的量化方式。

當(dāng)然，TensorRT也不是沒有缺陷。TensorRT對CUDA版本、GPU版本的要求還是比較苛刻的，包括針對不同的硬件環(huán)境，它的可移植性是相對較差，所以在實(shí)際的應(yīng)用部署過程中，盡可能的需要保證研發(fā)環(huán)境和實(shí)際的部署環(huán)境一致。