前言
最近一直忙于模型移植板端,用了不少厂家的sdk,发现挺多厂家的sdk都处于起步阶段,缺少一些技术支持,比如不支持五维向量,不支持一些onnx算子,不支持过深的模型结构,我最爱的Yolov5,v6,v7等高精度目标检测模型都无法移植上去,那就只能把眼光放回几年前尝试移植yolov3tiny和yolov4tiny,结果发现这玩意它要么darknet转onnx,要么是非官方实现的pytorch转onnx,转了之后还要求版本,算子,才能移植。
比如yolov3tiny网上下载的预训练权重本身就有问题(虽然之后修复但因为未知原因仍然没法移植过去),而yolov4tiny的问题在于不支持split算子,所以本文我来分享一下修改yolov4tiny的结构从而导出能移植的onnx,免得后来人又像我一样浪费几天时间折腾。
注:修改onnx有两种方式,导出前改和导出后改,在这里不推荐导出后改,难改也容易出问题,另外网上的onnx可视化修改工具我也试过了,十个里八个修改错误,基本没用。
准备
onnx
onnxsim
pytorch
这些前置要求请根据自己的环境配置。
我这里onnx1.12.0,torch1.11.0+cpu,onnxsim0.4.10。
1.下载yolov4的pytorch版源码和对应pth文件。
git clone https://github.com/bubbliiiing/yolov4-tiny-pytorch
进github页面里附带的百度网盘地址:
下载yolov4_tiny_weights_coco.pth即可。这里给不懂pytorch的说明一下,pth需配合模型网络代码才能成功加载,你如果只load pth只有一串字典。。
2.下载解压后打开目录下predict.py文件:
把mode里的值修改为export_onnx,就可以运行python predict.py产生onnx了
3.打开同一目录下yolo.py文件,查看或修改如下:
如果你没有显卡或者说没有安装cuda,把cuda设为False,这里默认是True。
这里是控制torch导出onnx的行为,其中最重要的是opset_version,请根据你所需要的算子版本来,这里默认是12,我需要的也是12,就不改了。此外默认是开启onnxsim简化的,这个简化很重要,能避免torch导出一些不必要的算子。
4.修改nets文件夹里的CSPdarknet53_tiny.py:
经过查找,这里面残差块的torch.split函数导致torch导出onnx时产生了split算子:
c = self.out_channels
# 对特征层的通道进行分割,取第二部分作为主干部分。
x = torch.split(x, c//2, dim = 1)[1]
# 对主干部分进行3x3卷积
x = self.conv2(x)
在netron里表现如下:
就是沿着第二个维度分成两半然后取第二部分。这其实和torch.chunk函数表现是相同的。
所以我们修改代码如下:
c = self.out_channels
# 对特征层的通道进行分割,取第二部分作为主干部分。
x = torch.chunk(x, 2, dim=1)[1]
# 对主干部分进行3x3卷积
x = self.conv2(x)
5.最后导出,命令行里运行python predict.py,并在model_data文件夹下用netron查看模型。
这里先展示没有开启onnxsim,导出的模型图,你会发现整了多余的shape和gather算子:
如果按照默认设置开启了onnxsim,就可以导出正确的模型结构图:
发现split已经完美变成了Slice,达到了我们的目的。
总结
我估计可能有看完的读者会问,既然如此为啥不直接用onnx的api来删除然后插入一个Slice算子呢?之前也讲过了,不仅很复杂,而且你插入之后onnx它不认!与其继续折腾不如在导出前就修改网络结构,只要明白torch和onnx算子的对应关系还是很容易做到的。另外这里预告一下,下一篇文章将介绍如何修改darknet转onnx的结构使得yolov3tiny.onnx能够正确运行,那个东西在github上被提了一堆issue反映shape inference error,但原始作者没有回答。