YOLOv5模型部署TensorRT之 FP32、FP16、INT8推理

点击上方蓝字关注我们

微信公众号：OpenCV学堂 关注获取更多计算机视觉与深度学习知识

引言

YOLOv5最新版本的6.x已经支持直接导出engine文件并部署到TensorRT上了。

FP32推理TensorRT演示

可能很多人不知道YOLOv5新版本6.x中已经支持一键导出Tensor支持engine文件，而且只需要一条命令行就可以完成：演示如下：

python export.py --weights yolov5s.pt --include onnx engine --device 0

其中onnx表示导出onnx格式的模型文件，支持部署到：

- OpenCV DNN- OpenVINO- TensorRT- ONNXRUNTIME

但是在TensorRT上推理想要速度快，必须转换为它自己的engine格式文件，参数engine就是这个作用。上面的命令行执行完成之后，就会得到onnx格式模型文件与engine格式模型文件。--device 0参数表示GPU 0，因为我只有一张卡！上述导出的FP32的engine文件。

使用tensorRT推理

YOLOv5 6.x中很简单，一条命令行搞定了，直接执行：

python detect.py --weights yolov5s.engine --view-img --source data/images/zidane.jpg

FP16推理TensorRT演示

在上面的导出命令行中修改为如下

python export.py --weights yolov5s.onnx --include engine --half --device 0

其中就是把输入的权重文件改成onnx格式，然后再添加一个新的参 --half 表示导出半精度的engine文件。就这样直接执行该命令行就可以导出生成了，图示如下：

对比可以发现相比FP32大小的engine文件，FP16的engine文件比FP32的engine大小减少一半左右，整个文件只有17MB大小左右。

推理执行的命令跟FP32的相同，直接运行，显示结果如下：

对比发现FP32跟FP16版本相比，速度提升了但是精度几乎不受影响！

INT8量化与推理TensorRT演示

TensorRT的INT量化支持要稍微复杂那么一点点，最简单的就是训练后量化。只要完成Calibrator这个接口支持，我用的TensorRT版本是8.4.0.x的，它支持以下几种Calibrator：

不同的量化策略，得到的结果可能稍有差异，另外高版本上的INT8量化之后到低版本的TensorRT机器上可能无法运行，我就遇到过！所以建议不同平台要统一TensorRT版本之后，再量化部署会比较好。上面的Calibrator都必须完成四个方法，分别是：

#使用calibrator验证时候每次张数，跟显存有关系，最少1张get_batch_size  #获取每个批次的图像数据，组装成CUDA内存数据get_batch  #如果以前运行过保存过，可以直接读取量化，低碳给国家省电read_calibration_cache  #保存calibration文件，量化时候会用到write_calibration_cache 

这块对函数集成不懂建议参考TensorRT自带的例子：

TensorRT-8.4.0.6\samples\python\int8_caffe_mnist

几乎是可以直接用的！Copy过来改改就好了！

搞定了Calibrator之后，需要一个验证数据集，对YOLOv5来说，其默认coco128数据集就是一个很好的验证数据，在data文件夹下有一个coco128.yaml文件，最后一行就是就是数据集的下载URL，直接通过URL下载就好啦。

完成自定义YOLOv5的Calibrator之后，就可以直接读取onnx模型文件，跟之前的官方转换脚本非常相似了，直接在上面改改，最重要的配置与生成量化的代码如下：

# build trt enginebuilder.max_batch_size = 1config.max_workspace_size = 1 << 30config.set_flag(trt.BuilderFlag.INT8)config.int8_calibrator = calibratorprint('Int8 mode enabled')plan = builder.build_serialized_network(network, config)

主要就是设置config中的flag为INT8，然后直接运行，得到plan对象，反向序列化为engine文件，保存即可。最终得到的INT8量化engine文件的大小在9MB左右。

数据太少，只有128张， INT8量化之后的YOLOv5s模型推理结果并不尽如人意。但是我也懒得再去下载COCO数据集， COCO训练集一半数据作为验证完成的量化效果是非常好。

这里，我基于YOLOv5s模型自定义数据集训练飞鸟跟无人机，对得到模型，直接用训练集270张数据做完INT8量化之后的推理效果如下：

量化效果非常好，精度只有一点下降，但是速度比FP32的提升了1.5倍左右（3050Ti）。

已知问题与解决

量化过程遇到这个错误

[09/22/2022-23:01:13] [TRT] [I]   Calibrated batch 127 in 0.30856 seconds.[09/22/2022-23:01:16] [TRT] [E] 2: [quantization.cpp::nvinfer1::DynamicRange::DynamicRange::70] Error Code 2: Internal Error (Assertion min_ <= max_ failed. )[09/22/2022-23:01:16] [TRT] [E] 2: [builder.cpp::nvinfer1::builder::Builder::buildSerializedNetwork::619] Error Code 2: Internal Error (Assertion engine != nullptr failed. )Failed to create the engineTraceback (most recent call last):

解决方法，把Calibrator中getBtach方法里面的代码：

img = np.ascontiguousarray(img, dtype=np.float32)

to

img = np.ascontiguousarray(img, dtype=np.float16)

这样就可以避免量化失败。

具体解释可以查看这个帖子！

https://github.com/NVIDIA/TensorRT/issues/1634

扫码获取YOLOv5 TensorRT INT8量化脚本与视频教程

扫码查看OpenCV+OpenVIO+Pytorch系统化学习路线图

 推荐阅读 

CV全栈开发者说 - 从传统算法到深度学习怎么修炼

2022入坑深度学习，我选择Pytorch框架！

Pytorch轻松实现经典视觉任务

教程推荐 | Pytorch框架CV开发-从入门到实战

OpenCV4 C++学习 必备基础语法知识三

OpenCV4 C++学习 必备基础语法知识二

OpenCV4.5.4 人脸检测+五点landmark新功能测试

OpenCV4.5.4人脸识别详解与代码演示

OpenCV二值图象分析之Blob分析找圆

OpenCV4.5.x DNN + YOLOv5 C++推理

OpenCV4.5.4 直接支持YOLOv5 6.1版本模型推理

OpenVINO2021.4+YOLOX目标检测模型部署测试

比YOLOv5还厉害的YOLOX来了，官方支持OpenVINO推理