将mmdetection中的快速实例分割算法yolact转为x3支持的onnx模型
安装好最新的mmdetection, mmdet官网有指导教程
mmdet提供了yolact算法,需要修改的相关目录如下:
Yolact配置文件: configs/yolact/yolact_r50_1x8_coco.py
配置文件中的关注数据处理模块,yolact在推理阶段直接将图片resize到目标尺寸,因此在工具链里面的数据预处理要做相应修改。
Yolact模型文件:/mmdet/models/detectors/yolact.py
模型文件主关注simple_test函数,这个函数是在推理时调用,所有需要改动的代码均在这个函数以及它的子函数里面。
Yolact head文件:/mmdet/models/dense_heads/yolact_head.py
Yolact调用的子函数,模型的后处理均在这个文件下完成,x3不支持的算子也可在该文件中发现,因此该文件需要重点修改。
yolact模型(不包含后处理)
输入:1x3x550x550
输出:1)bbox_pred; 2) scores; 3) mask_coefficients; 4) mask_proto
对yolact模型不熟的可去看原文
mmdet提供了转换模型到onnx模型的工具
工具目录:mmdetection/tools/deployment/pytorch2onnx.py
目前mmdet里面已经支持一些方法直接转为onnx模型(包含了所有后处理部分),但由于算子的不支持,这些模型均无法直接在x3上使用。
在改代码前建议阅读以下资料,对onnx有一定了解。
https://zhuanlan.zhihu.com/p/272767300
https://zhuanlan.zhihu.com/p/273566106
https://zhuanlan.zhihu.com/p/286298001
一些有用工具:
OK,下面直接上代码
1. Pytorch2onnx修改部分,参数可参考图片内容自己根据模型修改
由于x3仅支持10版本,须将onnx版本改为10
Yolact模型输出有四个元素(不考虑后处理),因此修改输出部分代码
到此,pytorch2onnx工具已经设订好了,接下来就是模型代码部分了。
2. Yolact模型修改
方法一:完全丢到mmdet里面处理部分,直接输出模型的四种成分,工具链中yolov5好像就是这种做法,所有后处理均在工具链中完成。
修改如下
直接修改simple_test函数,仅需添加上述三行代码。运行pytorch2onnx函数,可获得成功转换onnx模型。
工具链端实测改模型可行,除最后一个relu算子在cpu上执行,其它所有算子均在BPU上执行。
方法二:方法一需要我们将所有的后处理重新在板端实现,因此我想尽可能多的操作在模型转换过程中实现,从而降低后处理难度。主要是将锚框生成部分并入onnx模型,改动代码分为两部分,如下
1.Simple_test函数
将添加的代码放到get_bboxes函数后,因为我要从get_bboxes里面拿到anchor信息
2.get_bboxes函数(在yolact_head.py里面)
现在运行pytorch2onnx应该可以转化成功,且用01_check.sh检测通过,所有算子均支持,由于输出中包含了anchor信息,因此也可以简化后处理部分内容。
采用第二种方法在工具链里面得到的一些结果
01_check.sh
由于添加了一些后处理信息,reshape, concat, transpose等一些算子在cpu上运行,其它的均在gpu上运行,暂时没有检验相比第一种方式效率是否降低了。
03_build.sh
04_inference.sh
针对方法二转换的模型,我自己写了一些后处理,简单的实现目标检测任务,实例分割部分还没实(纯靠print大法实在难受)。然后显而易见,得到了一个非常糟糕的结果, 没有高效的debug工具太难找哪里的问题了。
直接上图
捉急的结果,都检测了些啥。。。。
第一次接触onnx,改得很粗糙, 但是work?
Anyway, Debug长路漫漫,先分享到这里了,弄完整实例分割后再与大家分享