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