TensorRT模型推理加速實踐
一、TensorRT簡介
TensorRT是由C++、CUDA、python三種語言編寫成的庫,有助于在 NVIDIA GPU上進行高性能推理?;谀壳爸髁鞯纳疃葘W習框架得到的模型都可通過TensorRT實現推理加速。
圖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的構造輸入名字要保持一致,具體代碼如下:
圖2 onnx轉TensorRT代碼
這兩種方式生成的trt模型的效果都是一樣的,但第一種方式由于要對輸入模型和本地環境做一下額外的校驗,通過python API這種初始化設定參數的方式要快很多,兩種方式都需要添加最小輸入、常規輸入、最大輸入,以保證模型能夠動態地處理最大長度以內的句子。
3.TensorRT推理
通過TensorRT的模型轉換后,外部訓練好的模型都被tensorRT統一成tensorRT可識別的engine文件(并優化過)。在推理時,只要通過TensorRT的推理SDK就可以完成推理。具體的推理過程如下:
圖3 TensorRT推理流程
代碼實現如下:
圖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使得模型整體吞吐量提升了幾十倍,著實驚艷。
圖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 |
三、總結
不管是在模型精度損失、以及模型推理性能加速上,TensorRT確實表現亮眼,如想要更進一步提升性能,可以嘗試int8的量化方式。
當然,TensorRT也不是沒有缺陷。TensorRT對CUDA版本、GPU版本的要求還是比較苛刻的,包括針對不同的硬件環境,它的可移植性是相對較差,所以在實際的應用部署過程中,盡可能的需要保證研發環境和實際的部署環境一致。
參考文獻
1、https://zhuanlan.zhiku.com/p/446477459
2、TensorRT SDK | NVIDIA Developer
3、https://zhuanlan.zhihu.com/p/446477075
4、https://docs.nvidia.com/deeplearning/tensorrt/quick-start-guide
5、https://blog.csdn.net/weixin_48026885/article/details/124025362
6、Deploying Deep Neural Networks with NVIDIA TensorRT | NVIDIA Technical Blog
7、NVIDIA Announces TensorRT 8 Slashing BERT-Large Inference Down to 1 Millisecond | NVIDIA Technical Blog
本文轉載自??AI遇見云??,作者: 錢博文
