TensorRT之yolov4模型量化

0 背景

在《DeepStream5.0系列之yolov4使用》中我介绍过 deepstream 调用 yolov4 的方法,但是只支持 FP32 和 FP16 精度,如果要想调用 INT8 精度的 engine 模型,需要额外生成一个量化表。

本文是在 Jetson Xavier AGX 上部署测试,另外 NX 也支持 int8,其它型号的 jetson 就不要试了

1 环境配置

项目基于 https://github.com/jkjung-avt/tensorrt_demos,因此先配置相关环境

git clone https://github.com/jkjung-avt/tensorrt_demos.git
cd tensorrt_demos/ssd
./install_pycuda.sh
sudo pip3 install onnx==1.4.1
cd ../plugins
make

2 数据集准备

 NVIDIA 官方文档说生成量化表有 500 张图片即可,为了提高精度,这里使用 1000 张图片

创建文件夹,并从自己模型的数据集中拷贝 1000 张图片过来

mkdir calib_images
for jpg in $(ls -1 /your/dataset/location/*.jpg | sort -R | head -1000); do     cp ${jpg} calib_images/;   done

3 模型量化

这里以我们自己训练的 yolov4.weights 和 yolov4.cfg 为例,类别是 1 类,输入大小是 608*608,将模型和配置文件拷贝到 tensorrt_demos/yolo 目录下

创建模型软链接

cat yolov4.cfg | sed -e '2s/batch=64/batch=1/' > yolov4-608.cfg
ln -sf yolov4.weights yolov4-608.weights

生成 onnx 模型

python yolo_to_onnx.py -m yolov4-608 -c 1

创建 int8 软链接

ln -s yolov4-608.cfg yolov4-int8-608.cfg
ln -s yolov4-608.onnx yolov4-int8-608.onnx

开始量化

python onnx_to_tensorrt.py -v --int8 -m yolov4-int8-608 -c 1

这个过程比较慢,在 Xavier AGX 上用了 15 分钟左右。

运行完毕后,在目录中生成 yolov4-int8-608.trt 文件以及 calib_yolov4-int8-608.bin 文件,即为最终需要的 engine 以及对应的量化表。

4 模型测试

返回上一级目录,运行 trt_yolo.py 脚本可以对生成的 engine 进行测试,需要自己准备一张测试图片 test.jpg

cd ..
python trt_yolo.py --image ./yolo/test.jpg -m yolov4-int8-608

会在窗口显示检测结果以及帧率,如图所示

为了进行比较,这里生成一个 FP32 精度的模型进行测试,方法和上边类似,指令如下

python onnx_to_tensorrt.py -v -m yolov4-608 -c 1

测试图片如下

可以看到,int8 精度的帧率加速了大约 32%,效果还是很明显的

已标记关键词 清除标记
©️2020 CSDN 皮肤主题: 游动-白 设计师:白松林 返回首页