基于TensorRT量化部署RepVGG模型
【GiantPandaCV導語】本文為大家介紹了一個TensorRT int8 量化部署 RepVGG 模型的教程，并開源了全部代碼。主要是教你如何搭建tensorrt環(huán)境，對pytorch模型做onnx格式轉換，onnx模型做tensorrt int8量化，及對量化后的模型做推理，實測在1070顯卡做到了不到1ms一幀！
0x0. RepVGG簡介
一個簡單但功能強大的卷積神經網絡架構，該架構具有類似于VGG的推理時間主體，該主體僅由3×3卷積和ReLU的堆棧組成，而訓練時間模型具有多分支拓撲。訓練時間和推理時間體系結構的這種解耦是通過結構性重新參數化技術實現的，因此模型是名為RepVGG。論文解讀可以看：RepVGG|讓你的ConVNet一卷到底，plain網絡首次超過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安裝包（編譯好了，也可以自己從官網下載編譯）可以從鏈接: https://pan.baidu.com/s/1Nl5XTAsUOyTbY6VbigsMNQ 密碼: c4dn cuda安裝 如果系統有安裝驅動，運行如下命令卸載 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 是否接受協議: accept 然后選擇Install 最后回車 vim ~/.bashrc 添加如下內容： 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（根據自己的python版本安裝） pip install pycuda 安裝python版本的cuda 定位到用戶根目錄 tar -xzvf opencv-3.4.2.zip 以備推理調用
0x2. RepVGG模型訓練以及轉換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行，修改類別數 vim train.py 修改IMAGENET_TRAINSET_SIZE參數 指定訓練圖片的數量 根據自己的訓練數據及配置設置data（數據集路徑），arch（我選擇的是最小的模型RepVGG-A0），epochs，lr，batch-size，model_path（設置ImageNet預訓練模型路徑，就是上面convert_model.py轉換得到的模型）等參數 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模型轉換為 int8 tensorrt引擎
cd RepVGG_TensorRT_int8/repvgg_tensorrt_int8_tools vim convert_trt_quant.py 修改如下參數 BATCH_SIZE 模型量化一次輸入多少張圖片 BATCH 模型量化次數 height width 輸入圖片寬和高 CALIB_IMG_DIR 量化圖片路徑(把訓練的圖片放到一個文件夾下，然后把這個文件夾設置為此參數，注意BATCH_SIZE*BATCH要小于或等于訓練圖片數量） 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參數為自己的用戶根目錄 vim repvgg_infer.cc修改如下參數 output_name repvgg模型有1個輸出 我們可以通過netron查看模型輸出名 pip install netron 安裝netron vim netron_repvgg.py 把如下內容粘貼 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 訓練的模型有多少類 參數配置完畢
mkdir?build
cd?build
cmake?..
make
./RepVGGsEngine 輸出平均推理時間，實測平均推理時間小于1ms一幀，不得不說，RepVGG真的很香！至此，部署完成！由于我訓練的是貓狗識別下面放一張貓狗同框的圖片結尾。
順便放一下我的數據集鏈接 鏈接: https://pan.baidu.com/s/1Mh6GxTLoXRTCRQh-TPUc3Q 密碼: 3dt3
歡迎關注GiantPandaCV, 在這里你將看到獨家的深度學習分享，堅持原創(chuàng)，每天分享我們學習到的新鮮知識。( ? ?ω?? )?
有對文章相關的問題，或者想要加入交流群，歡迎添加BBuf微信：
為了方便讀者獲取資料以及我們公眾號的作者發(fā)布一些Github工程的更新，我們成立了一個QQ群，二維碼如下，感興趣可以加入。
評論
圖片
表情