【目标檢測】使用TensorRT加速YOLOv5

前言

今天是程式員節，當然是以程式員的方式來度過節日。

很早就聽說TensorRT可以加速模型推理，但一直沒時間去進行實踐，今天就來把這個陳年舊坑填補一下。

背景知識

在實踐之前有必要了解一下相關知識。

TensorRT簡介

TensorRT是可以在NVIDIA各種GPU硬體平台下運作的一個C++推理架構。我們利用Pytorch、TF或者其他架構訓練好的模型，可以轉化為TensorRT的格式，然後利用TensorRT推理引擎去運作我們這個模型，進而提升這個模型在英偉達GPU上運作的速度[1]。

TensorRT支援幾乎所有主流深度學習架構，将python架構轉換成C++的TensorRT，進而可以加速推理。

具體而言，TensorRT主要做了以下幾點來加快推理速度[1]：

• 算子融合(層與張量融合)：簡單來說就是通過融合一些計算op或者去掉一些多餘op來減少資料流通次數以及顯存的頻繁使用來提速
• 量化：量化即IN8量化或者FP16以及TF32等不同于正常FP32精度的使用，這些精度可以顯著提升模型執行速度并且不會保持原先模型的精度
• 核心自動調整：根據不同的顯示卡構架、SM數量、核心頻率等(例如1080TI和2080TI)，選擇不同的優化政策以及計算方式，尋找最合适目前構架的計算方式
• 動态張量顯存：我們都知道，顯存的開辟和釋放是比較耗時的，通過調整一些政策可以減少模型中這些操作的次數，進而可以減少模型運作的時間
• 多流執行：使用CUDA中的stream技術，最大化實作并行操作

當然，TensorRT主要缺點是與特定GPU綁定[1]，在不同型号上轉換出來的模型不能通用(這一點筆者暫未去從實踐證明)

TensorRT官方在其​​倉庫​​提供了三個開源工具，之後有需要可以使用。

三個工具大緻用途[1]：

• ONNX GraphSurgeon

  可以修改我們導出的ONNX模型，增加或者剪掉某些節點，修改名字或者次元等等

• Polygraphy

  各種小工具的集合，例如比較ONNX和trt模型的精度，觀察trt模型每層的輸出等等，主要用-來debug一些模型的資訊

• PyTorch-Quantization

  可以在Pytorch訓練或者推理的時候加入模拟量化操作，進而提升量化模型的精度和速度，并且支援量化訓練後的模型導出ONNX和TensorRT

ONNX簡介

Open Neural Network Exchange（ONNX，開放神經網絡交換）格式，是微軟和Facebook提出用來表示深度學習模型的開放格式，定義了一組和環境，平台均無關的标準格式，可使模型在不同架構之間進行轉移[2]。

假設模型在 TensorFlow 中，但想要在TensorRT中使用，或者模型在PyTorch中，想要在TFLite中使用，就可以利用 ONNX 作為中介，進行模型轉換。

典型的幾個線路[3]：

• Pytorch -> ONNX -> TensorRT
• Pytorch -> ONNX -> TVM
• TF -> onnx -> ncnn
• Pytorch -> ONNX -> tensorflow

ONNX結構是将每一個網絡的每一層或者每一個算子當作節點Node，再由這些Node去建構一個Graph，最後将Graph和這個onnx模型的其他資訊結合在一起，生成一個model。

可以通過線上網站​​https://netron.app/​​來檢視ONNX模型結構。

實踐上手

下面來進行實踐上手。

實驗環境

這是我的電腦環境：

• 作業系統：Windows10
• 顯示卡：RTX2060
• CUDA版本：11.6
• Pytorch版本：1.7.1
• Python版本：3.8

檢視CUDA版本

TensortRT非常依賴CUDA版本，在安裝之前，需要先檢視本機安裝好的CUDA版本，檢視方式有多種，第一種方式可以通過NVIDIA 控制台檢視；第二種方式可以通過在控制台輸入​

​nvcc -V​

​進行檢視。

安裝TensorRT

首先需要到​​Nvidia官網​​去下載下傳對應Cuda版本的TensorRT安裝包。

我這裡下載下傳的是紅框選中的這一個，這個版本支援CUDA11.0-11.7

下載下傳好後，将壓縮包進行解壓縮，然後到Anaconda環境中，以此進入到以下幾個檔案夾，安裝這幾個whl檔案

cd TensorRT-8.4.3.1\python
pip install tensorrt-8.4.3.1-cp38-none-win_amd64.whl

cd TensorRT-8.4.3.1\graphsurgeon
pip install graphsurgeon-0.4.6-py2.py3-none-any.whl

cd TensorRT-8.4.3.1\onnx_graphsurgeon
pip install onnx_graphsurgeon-0.3.12-py2.py3-none-any.whl

cd TensorRT-8.4.3.1\uff
pip install uff-0.6.9-py2.py3-none-any.whl      

然後需要移動安裝包裡的一些檔案：

• 将​

  ​TensorRT-8.4.3.1\include​

  ​​中頭檔案拷貝到​

  ​C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.6\include​

  ​
• 将​

  ​TensorRT-8.4.3.1\lib​

  ​​中所有lib檔案拷貝到​

  ​C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.6\lib\x64​

  ​
• 将​

  ​TensorRT-8.4.3.1\lib​

  ​​中所有dll檔案拷貝到​

  ​C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.6\bin​

  ​

注：這裡的v11.6根據自己的Cuda版本号即可

之後，需要手動将​

​C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.6\bin​

​路徑添加到使用者Path環境變量中

添加完之後，重新開機電腦，讓環境變量生效。

之後進行驗證：

python
import tensorrt
print(tensorrt.__version__)      

在​

​import​

​時發生報錯，

FileNotFoundError: Could not find: nvinfer.dll. Is it on your PATH?

此時隻需要将缺少的檔案找到，然後添加到上面的​

​bin​

​目錄下即可，我這裡是在安裝的torch中lib檔案下找到的部分檔案，缺什麼移什麼即可。

如無報錯，再次驗證，可以輸出tensorrt版本：

下面運作安裝包裡面的一個sample.py檔案，以確定tensorrt能夠正常工作。

進入到下圖所示的路徑，運作​​

​sample.py​

​，如果正常輸出，則代表tensorrt安裝成功。

如果提示沒裝pycuda，還需要再安裝一下

pip install pycuda      

YOLOv5使用TensorRT加速

這部分内容看到不少部落格都是用Cmake編譯生成yolov5的VS工程，非常繁瑣麻煩，主要是這些博文寫作時間較早。

而在YOLOv5 6.0版本更新後，官方新增了一個​​

​export.py​

​檔案，支援大部分架構模型的導出，包括TensorRT。

下面我所使用的是​​YOLOv5官方​​最新的6.2版本。

下面直接導出官方提供的yolov5s模型試試，終端輸入：

python export.py --weights yolov5s.pt --data data/coco128.yaml --include engine --device 0 --half      

注：這裡的​

​--half​

​表示半精度模型，使用半精度可以加快推理速度，但會損失一定精度，直接導出可以不加

初次導出，遇到如下報錯

ONNX: export failure 0.4s: Exporting the operator silu to ONNX opset version 12 is not supported.

這個報錯需要修改pytorch的激活函數，找到該函數位置：​

​D:\anaconda\envs\pytorch\Lib\sitepackages\torch\nn\modules\activation.py​

​(此處結合自己的anaconda實際安裝位置來更改)

修改代碼如下：

class SiLU(Module):

    __constants__ = ['inplace']
    inplace: bool
 
    def __init__(self, inplace: bool = False):
        super(SiLU, self).__init__()
        self.inplace = inplace
 
    def forward(self, input: Tensor) -> Tensor:
        # ------------------------------------- #
        # 把F.silu替換掉，修改後如下
        return input * torch.sigmoid(input)
 
        #原來的代碼
        return F.silu(input, inplace=self.inplace)      

再次導出，不再報錯。

值得注意的是，YOLOv5并不會直接導出TensorRT模型，而是會先導出ONNX模型，然後将ONNX模型轉換成TensorRT模型，是以導出完成後，會在模型位置處生成​

​yolov5s.onnx​

​​和​

​yolov5s.engine​

​​，​

​yolov5s.engine​

​就是可以用來直接推理的TensorRT模型。

經過實測，不添加半精度導出yolov5s模型花費時間99.5s，添加半精度之後，導出yolov5s模型花費時間404.2s。

實驗結果

圖檔檢測

首先是來檢測一下圖檔的推理速度，首先修改detect.py，統計程式花費時間。

if __name__ == "__main__":
    begin_time = time.time()
    opt = parse_opt()
    main(opt)
    end_time = time.time()
    print("程式花費時間{}秒".format(end_time-begin_time))      

然後依次使用不同模型進行推理

python detect.py --weights yolov5s.pt
python detect.py --weights yolov5s.engine
python val.py --weights yolov5s.pt
python val.py --weights yolov5s.engine      

這裡資料源選擇的是coco128中128張圖檔，整體實驗結果如下表所示：

模型	推理花費時間(s)	AP50
原始模型	8.39	71.4%
TensorRT(全精度)	5.45	71.1%
TensorRT(半精度)	4.83	70.8%

從資料可以發現，使用TensorRT加速之後，模型推理速度提升了約35%，但是模型精度并沒有下降太多。coco資料集128張圖檔，模型訓練和檢測都是用同一份資料，這可能會對AP産生一定影響，于是再換用Visdrone資料集在進行實驗。

下面對Visdrone資料集進行實驗，使用yolov5m模型，訓練100個epoch.

使用​

​VisDrone2019-DET-test-dev​

​中的1610張圖檔進行驗證和檢測：

驗證指令：

python val.py --data data/VisDrone.yaml --weights runs/train/exp3/weights/best.engine --batch-size 4 --task test      

使用​

​VisDrone2019-DET-test-dev​

​中的

檢測指令：

python detect.py --weights runs/train/exp3/half/best.engine --source D:/Desktop/Work/Dataset/VisDrone/VisDrone2019-DET-test-dev/images --data data/VisDrone.yaml      

實驗結果如下表所示：

模型	驗證花費時間(s)	P	R	AP50	推理花費時間(s)
yolov5m	101.28	43.1%	34.9%	32.0%	157.51
onnx	156.16	42.8%	34.9%	32.0%	158.97
TensorRT(全精度)	124.37	42.8%	34.9%	32.0%	144.93
TensorRT(半精度)	127.85	42.8%	34.8%	31.9%	139.97

由表可見，使用TensorRT加速之後，推理速度提升約了10%，同時精度隻掉了0.3%，AP50基本上變化不大。

下面選一張圖檔來直覺對比一下，左側圖為原始模型推理圖，右側圖為TensorRT(半精度) 推理圖，兩者大緻上差異不大，各有各的漏檢對象。

視訊檢測

視訊檢測用了王者榮耀資料集做一個實驗，比較了正常檢測，tensorrt和onnx推理速度和幀率。推理所花費的時間分别為：

• 正常推理 程式花費時間20.88s
• onnx推理 程式花費時間25.68s
• tensorrt推理 程式花費時間16.03s

視訊幀率如下視訊所示：