基于TensorRT量化部署RepVGG模型
【GiantPandaCV導語】本文為大家介紹了一個TensorRT int8 量化部署 RepVGG 模型的教程，并開源了全部代碼。主要是教你如何搭建tensorrt環(huán)境，對pytorch模型做onnx格式轉(zhuǎn)換，onnx模型做tensorrt int8量化，及對量化后的模型做推理，實測在1070顯卡做到了不到1ms一幀！
0x0. RepVGG簡介
一個簡單但功能強大的卷積神經(jīng)網(wǎng)絡架構(gòu)，該架構(gòu)具有類似于VGG的推理時間主體，該主體僅由3×3卷積和ReLU的堆棧組成，而訓練時間模型具有多分支拓撲。訓練時間和推理時間體系結(jié)構(gòu)的這種解耦是通過結(jié)構(gòu)性重新參數(shù)化技術實現(xiàn)的，因此模型是名為RepVGG。論文解讀可以看：RepVGG|讓你的ConVNet一卷到底，plain網(wǎng)絡首次超過80%top1精度
開源代碼地址如下：https://github.com/Wulingtian/RepVGG_TensorRT_int8
0x1. 環(huán)境配置
ubuntu：18.04 cuda：11.0 cudnn：8.0 tensorrt：7.2.16 OpenCV：3.4.2 cuda，cudnn，tensorrt和OpenCV安裝包（編譯好了，也可以自己從官網(wǎng)下載編譯）可以從鏈接: https://pan.baidu.com/s/1Nl5XTAsUOyTbY6VbigsMNQ 密碼: c4dn cuda安裝 如果系統(tǒng)有安裝驅(qū)動，運行如下命令卸載 sudo apt-get purge nvidia* 禁用nouveau，運行如下命令 sudo vim /etc/modprobe.d/blacklist.conf 在末尾添加 blacklist nouveau 然后執(zhí)行 sudo update-initramfs -u chmod +x cuda_11.0.2_450.51.05_linux.run sudo ./cuda_11.0.2_450.51.05_linux.run 是否接受協(xié)議: accept 然后選擇Install 最后回車 vim ~/.bashrc 添加如下內(nèi)容： export PATH=/usr/local/cuda-11.0/bin:$PATH export LD_LIBRARY_PATH=/usr/local/cuda-11.0/lib64:$LD_LIBRARY_PATH source .bashrc 激活環(huán)境 cudnn 安裝 tar -xzvf cudnn-11.0-linux-x64-v8.0.4.30.tgz cd cuda/include sudo cp *.h /usr/local/cuda-11.0/include cd cuda/lib64 sudo cp libcudnn* /usr/local/cuda-11.0/lib64 tensorrt及OpenCV安裝 定位到用戶根目錄 tar -xzvf TensorRT-7.2.1.6.Ubuntu-18.04.x86_64-gnu.cuda-11.0.cudnn8.0.tar.gz cd TensorRT-7.2.1.6/python，該目錄有4個python版本的tensorrt安裝包 sudo pip3 install tensorrt-7.2.1.6-cp37-none-linux_x86_64.whl（根據(jù)自己的python版本安裝） pip install pycuda 安裝python版本的cuda 定位到用戶根目錄 tar -xzvf opencv-3.4.2.zip 以備推理調(diào)用
0x2. RepVGG模型訓練以及轉(zhuǎn)換onnx
定位到用戶根目錄 git clone https://github.com/Wulingtian/RepVGG_TensorRT_int8.git cd RepVGG_TensorRT_int8/models vim convert_model.py 設置 num_classes，例如：我訓練的是貓狗識別，則設置為2 python convert_model.py 生成可加載的ImageNet預訓練模型路徑 cd RepVGG_TensorRT_int8 vim repvgg.py 定位到154行，修改類別數(shù) vim train.py 修改IMAGENET_TRAINSET_SIZE參數(shù) 指定訓練圖片的數(shù)量 根據(jù)自己的訓練數(shù)據(jù)及配置設置data（數(shù)據(jù)集路徑），arch（我選擇的是最小的模型RepVGG-A0），epochs，lr，batch-size，model_path（設置ImageNet預訓練模型路徑，就是上面convert_model.py轉(zhuǎn)換得到的模型）等參數(shù) python train.py，開始訓練，模型保存在當前目錄，名為model_best.pth.tar python convert.py model_best.pth.tar RepVGG-A0-deploy.pth -a RepVGG-A0（指定模型類型，我訓練的是RepVGG-A0） vim export_onnx.py 設置arch，weights_file（convert.py生成的模型），output_file（輸出模型名稱），img_size（圖片輸入大?。?，batch_size（推理的batch） python export_onnx.py得到onnx模型
0x3. onnx模型轉(zhuǎn)換為 int8 tensorrt引擎
cd RepVGG_TensorRT_int8/repvgg_tensorrt_int8_tools vim convert_trt_quant.py 修改如下參數(shù) BATCH_SIZE 模型量化一次輸入多少張圖片 BATCH 模型量化次數(shù) height width 輸入圖片寬和高 CALIB_IMG_DIR 量化圖片路徑(把訓練的圖片放到一個文件夾下，然后把這個文件夾設置為此參數(shù)，注意BATCH_SIZE*BATCH要小于或等于訓練圖片數(shù)量） onnx_model_path onnx模型路徑（上面運行export_onnx.py得到的onnx模型） python convert_trt_quant.py 量化后的模型存到models_save目錄下
0x4. tensorrt模型推理
cd RepVGG_TensorRT_int8/repvgg_tensorrt_int8 cd yolov5_tensorrt_int8 vim CMakeLists.txt 修改USER_DIR參數(shù)為自己的用戶根目錄 vim repvgg_infer.cc修改如下參數(shù) output_name repvgg模型有1個輸出 我們可以通過netron查看模型輸出名 pip install netron 安裝netron vim netron_repvgg.py 把如下內(nèi)容粘貼 import netron netron.start('此處填充簡化后的onnx模型路徑', port=3344) python netron_repvgg.py 即可查看 模型輸出名 trt_model_path 量化的tensorrt推理引擎（models_save目錄下trt后綴的文件） test_img 測試圖片路徑 INPUT_W INPUT_H 輸入圖片寬高 NUM_CLASS 訓練的模型有多少類 參數(shù)配置完畢
mkdir?build
cd?build
cmake?..
make
./RepVGGsEngine 輸出平均推理時間，實測平均推理時間小于1ms一幀，不得不說，RepVGG真的很香！至此，部署完成！由于我訓練的是貓狗識別下面放一張貓狗同框的圖片結(jié)尾。
順便放一下我的數(shù)據(jù)集鏈接 鏈接: https://pan.baidu.com/s/1Mh6GxTLoXRTCRQh-TPUc3Q 密碼: 3dt3
歡迎關注GiantPandaCV, 在這里你將看到獨家的深度學習分享，堅持原創(chuàng)，每天分享我們學習到的新鮮知識。( ? ?ω?? )?
有對文章相關的問題，或者想要加入交流群，歡迎添加BBuf微信：
為了方便讀者獲取資料以及我們公眾號的作者發(fā)布一些Github工程的更新，我們成立了一個QQ群，二維碼如下，感興趣可以加入。