基于TensorRT量化部署RepVGG模型

文章转载于:GiantPandaCV
作者:willer

【GiantPandaCV导语】本文为大家介绍了一个TensorRT int8 量化部署 RepVGG 模型的教程，并开源了全部代码。主要是教你如何搭建tensorrt环境，对pytorch模型做onnx格式转换，onnx模型做tensorrt int8量化，及对量化后的模型做推理，实测在1070显卡做到了不到1ms一帧！

0x0. RepVGG简介

一个简单但功能强大的卷积神经网络架构，该架构具有类似于VGG的推理时间主体，该主体仅由3×3卷积和ReLU的堆栈组成，而训练时间模型具有多分支拓扑。训练时间和推理时间体系结构的这种解耦是通过结构性重新参数化技术实现的，因此模型是名为RepVGG。论文解读可以看：RepVGG|让你的ConVNet一卷到底，plain网络首次超过80%top1精度

开源代码地址如下：https://github.com/Wulingtian...\_TensorRT\_int8

0x1. 环境配置

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/1Nl5X... 密码: c4dn
cuda安装
如果系统有安装驱动，运行如下命令卸载
sudo apt-get purge nvidia*
禁用nouveau，运行如下命令
sudo vim /etc/modprobe.d/blacklist.conf
在末尾添加
blacklist nouveau
然后执行
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 激活环境
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...\_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目录下