0. 前言
昨天美团开源了YOLOv6,又是一个YOLO系列的新作。此时距离YOLOX开源差不多刚好一年的时间。之前捏过很多YOLO系列的推理例子,比如YOLOv3、YOLOv4、YOLOv5、YOLOX、YOLOR和YOLOP等等。虽然最近已经没有在做detection方向了,但作为YOLO系列的老粉了,出来凑个热闹应该总是可以的。所以,这次也来凑凑YOLO系列的热闹,给出几个不同的推理引擎的例子,包括ONNXRuntime、MNN、NCNN和TNN,以及简单记录下模型的转换过程。总的来说,YOLOv6 的 C++推理,都是些重复性的工作,没什么太大的难度,刚好趁着周末,顺手捏一捏。这篇文章不会记录地很详细,只讲几个要点。
1. ONNX 和 TNN 模型转换
经过尝试,直接转换出来的ONNX和TNN模型文件在推理时,结果一切正常,不需要修改 YOLOv6 的 Detect 源码,使用官方提供的 deploy/ONNX/export_onnx.py 直接转换即可。但是 NCNN 和 MNN 都需要修改 Detect 的源码进行特殊处理才可正常推理。所以 ONNX 和 TNN 放在这一节讲,MNN 和 NCNN 的模型转换放在下一小节讲。
首先下载源码:
然后稍微修改下 export_onnx.py,源码是没有添加 onnxsim的,作为基操,我们把它添加上,修改后如下:
预训练好的pt模型文件可以从官方提供的链接下载,放在 YOLOv6 的根目录下,直接转换就行:
这个过程比较顺利,暂时没发现什么坑。接下来转换为 TNN 模型文件,命令如下:
TNN模型的转换需要用到 tnn-convert,如何使用tnn-convert就不展开了,感兴趣的同学可以看我之前写的一篇文章,传送门:
tnn-convert搭建简记-YOLOP转TNN5:https://zhuanlan.zhihu.com/p/431418709
2. MNN 和 NCNN 模型转换
NCNN 和 MNN 都需要修改 Detect 的源码进行特殊处理才可正常推理。MNN其实也可以直接转,但是转出来的模型文件虽然能推理,但是在decode完之后,结果很奇怪,所以我最后决定将MNN的模型文件转换采用和NCNN同样的处理方式。其实很多写的关于YOLO系列的部署文章都有提到过,该系列在部署时主要的一个问题就是如何处理 Detect Head 中关于 decode 部分的逻辑。这部分代码在 YOLOv6中长这样:
这部分虽然有些框架可以直接支持导出,但会产生大量的胶水op,所以一个可选的做法是,只导出decode之前的Raw部分的内容,在C++侧做decode。另外,我们可以看到,decode部分存在一个5维度的操作:
这部分在NCNN中应该是不支持的(按照我对ncnn::Mat的理解,它有c,h,w三个维度,并假设b=1,所以可以处理<=4维的张量),也不能直接导出。所以,这个5维的处理,我们也要做相应的修改。至于MNN,其实可以直接转换这部分decode的逻辑,但是我在推理时,发现出来的结果不太对,于是决定采用NCNN同样的处理方式,就是只导出decode前的部分,把decode放在c++侧处理,后来验证了这样做可以得到正常的推理结果。
那么,这段 Detect Head 的逻辑到底要怎么改呢?长话短说,我直接放一个我修改后的代码吧:
models/effidehead.py修改之后
这里,我们把原来的5维操作修改成4维操作,因为导出时,实际的no(num_outputs=85)、na(num_anchors=1)的值都是可以事先计算出来的,通过看 YOLOv6 中的 configs 文件夹中的配置文件,我们也可以确定实际上 na一直为1,而 ny,nx是每个特征图的大小,对于固定的输入shape,各个特征图的ny,nx也是固定的。为了变成4维操作,我们可以把原来在最后的两个维度(ny,nx)直接拉平,按照行主序的线性内存来理解,这样是可行的。于是有:
最后,在返回结果值之前,我们再做一个合并处理,这样就不需要在c++解码的时候单独对每个特征图都做一遍,合并后,只要对一个大的输出做decode就可以了。
这样修改之后,export_onnx.py的代码不用变,还是用原来的命令行直接导出就可以了(我在文件名增加了-for-ncnn作为后缀方便区分)。
还是先导出为 ONNX:
模型文件为:
用netron打开来看,发现decode那部分已经没有了:
带decode的onnx,用netron打开是长这样的:
可以发现两者有很明显的区别,不带decode的图,在输出部分要简单很多。
接下来,就是按常规流程将ONNX转换成NCNN和MNN模型,命令行如下:
一切正常。转换为MNN模型的命令为:
也是一切正常。
3. ONNX 和 TNN 模型 C++ 推理
ONNX和TNN的模型都是带decode的,因此在后处理时简单些,不用生成anchor了。模型推理直接输出的维度是(1,n,85),这个n表示总共输出的anchors个数,85的含义是:
由于输出的坐标就已经是归一化后的cx,cy和cw,ch,所以后处理就很简单了,直接转换成x1,y1,x2,y2格式就行。逻辑大概如下:
详细的推理代码就不展开了,会在文章最后放出。
4. NCNN 和 MNN 模型 C++ 推理
NCNN和MNN的模型文件没有导出decode部分,因此后处理复杂一点。后处理主要包括2部分,一是生成anchors,二是根据生成的anchors和输出的原始信息解码坐标。YOLOv6 的anchors生成逻辑其实和YOLOX基本是一致的,就是每个feature map上每个锚点生成一个anchor框,做过detection算法的同学应该很多理解这句话的意思,我也就不啰嗦了。直接放代码吧。
generate_anchors函数主要逻辑
坐标解码的主要逻辑
详细的推理代码就不展开了,会在文章最后放出。
5. YOLOv6 C++推理的使用例子
首先是放出YOLOv6 的4个推理引擎的C++ 源码,想必大家最关心就是能不能白嫖了。所有的代码都集成进了lite.ai.toolkit(https://github.com/DefTruth/lite.ai.toolkit) 工具箱中,零成本无压力白嫖。
- YOLOv6 ONNXRuntime C++ 源码:https://github.com/DefTruth/lite.ai.toolkit/blob/main/lite/ort/cv/yolov6.cpp
- YOLOv6 MNN C++ 源码:https://github.com/DefTruth/lite.ai.toolkit/blob/main/lite/mnn/cv/mnn_yolov6.cpp
- YOLOv6 NCNN C++ 源码:https://github.com/DefTruth/lite.ai.toolkit/blob/main/lite/ncnn/cv/ncnn_yolov6.cpp
- YOLOv6 TNN C++ 源码:https://github.com/DefTruth/lite.ai.toolkit/blob/main/lite/tnn/cv/tnn_yolov6.cpp
对源码感兴趣的同学,可自行选择关心的推理引擎版本进行阅读。接下来,在简单贴几个使用lite.ai.toolkit 工具箱一键调用的例子。
ONNXRuntime版本
MNN版本
NCNN版本
TNN版本
输出的结果如下:
延庆川北小区45孙老师 东屯 收卖废品垃圾破烂炒股 废品孙 一起人工智能把
没有篇 c++ yolo用不了 模型转换有用哦 虽然yolo7 也出来了 还么i找到哦
6. 总结
本文主要讲解了美团6月24开源的YOLOv6在不同推理引擎的C++工程化过程,包含了ORT、MNN、NCNN和TNN下的推理处理,以及不同推理框架下模型转换需要注意的问题。并且,所有的代码都集成进了lite.ai.toolkit 工具箱中,零成本无压力白嫖。
- YOLOv6 ONNXRuntime C++ 源码
- YOLOv6 MNN C++ 源码
- YOLOv6 NCNN C++ 源码
- YOLOv6 TNN C++ 源码