作者在将YOLOv5 pytorch版本转成wk文件中碰到了很多的坑,但在研究中也收获了很多知识,针对转化过程,作者进行了详细的解读来给大家避坑,如果在转化过程中有什么问题也可以联系作者,我们可以一起进行分析和讨论。
现如今,官方最新版本的YOLOv5是6.0以上的版本,而6.0版本相比于6.0以下版本(不包含6.0)发生了不少的变化,其中最为突出的改动是将force层修改为成步长为2,卷积为 2 x 2的Conv层 (GitHub - ultralytics/yolov5: YOLOv5 in PyTorch > ONNX > CoreML > TFLite)。
YOLOv5 5.0和6.0以YOLOv5s.yaml为例
# yolov5s.yaml YOLOv5 5.0版本
# parameters
nc: 10 # number of classes
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.50 # layer channel multiple
# anchors
anchors:
- [10,13, 16,30, 33,23] # P3/8
- [30,61, 62,45, 59,119] # P4/16
- [116,90, 156,198, 373,326] # P5/32
# YOLOv5 backbone
backbone:
# [from, number, module, args]
[[-1, 1, Focus, [64, 3]], # 这块部分 0-P1/2
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
[-1, 3, C3, [128]],
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
[-1, 9, C3, [256]],
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
[-1, 9, C3, [512]],
[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
[-1, 1, SPP, [1024, [5, 9, 13]]],
[-1, 3, C3, [1024, False]], # 9
]
# YOLOv5 head
head:
[[-1, 1, Conv, [512, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 6], 1, Concat, [1]], # cat backbone P4
[-1, 3, C3, [512, False]], # 13
[-1, 1, Conv, [256, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 4], 1, Concat, [1]], # cat backbone P3
[-1, 3, C3, [256, False]], # 17 (P3/8-small)
[-1, 1, Conv, [256, 3, 2]],
[[-1, 14], 1, Concat, [1]], # cat head P4
[-1, 3, C3, [512, False]], # 20 (P4/16-medium)
[-1, 1, Conv, [512, 3, 2]],
[[-1, 10], 1, Concat, [1]], # cat head P5
[-1, 3, C3, [1024, False]], # 23 (P5/32-large)
[[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
]
# YOLOv5s.yaml YOLOv5 6.0版本
# Parameters
nc: 10 # number of classes
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.50 # layer channel multiple
anchors:
- [10,13, 16,30, 33,23] # P3/8
- [30,61, 62,45, 59,119] # P4/16
- [116,90, 156,198, 373,326] # P5/32
# YOLOv5 v6.0 backbone
backbone:
# [from, number, module, args]
[[-1, 1, Conv, [64, 6, 2, 2]], # 这块部分 0-P1/2
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
[-1, 3, C3, [128]],
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
[-1, 6, C3, [256]],
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
[-1, 9, C3, [512]],
[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
[-1, 3, C3, [1024]],
[-1, 1, SPPF, [1024, 5]], # 9
]
# YOLOv5 v6.0 head
head:
[[-1, 1, Conv, [512, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 6], 1, Concat, [1]], # cat backbone P4
[-1, 3, C3, [512, False]], # 13
[-1, 1, Conv, [256, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 4], 1, Concat, [1]], # cat backbone P3
[-1, 3, C3, [256, False]], # 17 (P3/8-small)
[-1, 1, Conv, [256, 3, 2]],
[[-1, 14], 1, Concat, [1]], # cat head P4
[-1, 3, C3, [512, False]], # 20 (P4/16-medium)
[-1, 1, Conv, [512, 3, 2]],
[[-1, 10], 1, Concat, [1]], # cat head P5
[-1, 3, C3, [1024, False]], # 23 (P5/32-large)
[[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
]
而wk(nnie)并不支持的Focus层,因此,如果用YOLOv5 6.0以下版本,需要将Focus层进行替换,替换成YOLOv5 6.0版本的Conv层,写法保持一致。
经查阅其他博客发现,上采样nn.Upsample在海思的运用会影响检测速度,因而文档作者将nn.Upsample用上采样的另一种表现形式nn.ConvTranspose2d来替换,替换后的结果如下:
# YOLOv5 6.0 yolov5s.yaml
# Parameters
nc: 10 # number of classes
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.50 # layer channel multiple
anchors:
- [10,13, 16,30, 33,23] # P3/8
- [30,61, 62,45, 59,119] # P4/16
- [116,90, 156,198, 373,326] # P5/32
# YOLOv5 v6.0 backbone
backbone:
# [from, number, module, args]
[[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
[-1, 3, C3, [128]],
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
[-1, 6, C3, [256]],
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
[-1, 9, C3, [512]],
[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
[-1, 3, C3, [1024]],
[-1, 1, SPPF, [1024, 5]], # 9
]
# YOLOv5 v6.0 head
head:
[[-1, 1, Conv, [512, 1, 1]],
#[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[-1, -1, nn.ConvTranspose2d,[256, 256, 2, 2]],
[[-1, 6], 1, Concat, [1]], # cat backbone P4
[-1, 3, C3, [512, False]], # 13
[-1, 1, Conv, [256, 1, 1]],
# [-1, 1, nn.Upsample, [None, 2, 'nearest']],
[-1, -1, nn.ConvTranspose2d,[128, 128, 2, 2]],
[[-1, 4], 1, Concat, [1]], # cat backbone P3
[-1, 3, C3, [256, False]], # 17 (P3/8-small)
[-1, 1, Conv, [256, 3, 2]],
[[-1, 14], 1, Concat, [1]], # cat head P4
[-1, 3, C3, [512, False]], # 20 (P4/16-medium)
[-1, 1, Conv, [512, 3, 2]],
[[-1, 10], 1, Concat, [1]], # cat head P5
[-1, 3, C3, [1024, False]], # 23 (P5/32-large)
[[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
]
模型修改后需要重新进行训练,之后的模型转化均是基于训练后的 .pt 或 .pth 文件进行的。
YOLOv5官方代码文件中自带 pt (pth) 转 onnx 代码文件export.py,6.0版本的export.py文件位置放置在主目录下。6.0版本以下的export.py文件一般放置在models文件夹下。
模型转化所需要的python库为 onnx 和 onnx-simplifier,版本号可根据YOLOv5官方代码文件中的 requirements.txt 来确定。
运行export.py文件如下:
python export.py --data mydata.yaml --weights yolov5s.pt --batch 1 --img 640 640 --train --simplify --include onnx --opset 10
其中,--img,--data和 --weights 参数是基于读者自己输入模型尺寸的大小,训练集的yaml文件地址以及生成的 .pt (.pth) 文件地址而定,剩余参数可依据情况而定(6.0版本的export.py文件进行的模型转化尽量都加上,要不容易产生错误,具体原因还没有查清)。
将生成的onnx文件进行进一步的简化,运行如下:
python -m onnxsim yolov5s.onnx yolov5s-sim.onnx
其中,onnxsim后的两个参数分别是上一步生成的onnx的地址以及生成简化后onnx的地址。
注:如果发生报错,请查看网络层是否发生了更改,参数是否正确等问题。
作者将onnx模型转化为Caffe模型是在linux的Ubuntu 18.04版本下进行的,这里需要搭建一下Caffe环境,采用的转化代码为Wulingtian老师的代码。
(GitHub - Wulingtian/yolov5_caffe: yolov5 onnx caffe)
将上一步生成的 yolov5s-sim.onnx 文件放入代码文件中,并打开convertCaffe.py文件,根据模型网络层的情况与simple_onnx文件地址修改main函数下面的内容。
例如:
其中,convertToCaffe函数选用最后一个即可,主要原因是Focus网络层在训练的时候已经被改为了Conv层,所以,无需再用其他函数。
如果转化成功,会出现以下信息:
注意:如果转化成功,便会出现类似以下的信息,如果转化失败,代码会报错。转化失败的原因可能是因为前面生成的simple_onnx不满足转化条件,存在Caffe中无法解析的网络层,或是在生成simple_onnx前修改了yolo.py、common.py等源码,或是onnx的版本号可能不符合需求,或是无法生成节点。因此在转化模型前,先用netron(https://netron.app/)(只需将onnx文件导入即可)来看一下神经网络结构。
安装并下载mingw64和RuyiStudio,RuyiStudio的版本为2.0.38。
打开RuyiStudio.exe,并创建NNIE Studio Project,然后赋予项目名称,设置相应的配置信息。在NNIE Studio Project中创建models和images文件夹(名称可以自己定),models中放入上一步转成Caffe的 xxx.caffemodel和 xxx.prototxt文件,images中放入20-30张模型测试的图片。
打开文件最下方的 xxx.cfg,导入caffe模型文件,并设置模型转化的对应参数。
参数说明: prototxt ------------------------------- 导入caffe模型的prototxt文件
caffemodel --------------------------- 导入caffe模型的caffemodel文件
net_type ------------------------------ 模型类别(图像都是CNN)
is_simulation ------------------------ 是否进行仿真(默认即可)
marked_prototxt -------------------- 预生成wk文件的prototxt文件 (需要进行修改)
output_wk_name -------------------- 导出的wk文件的命名
complie_model ---------------------- 导出模型的精度
log_level ------------------------------ 是否开启日志文件
align_bytes ---------------------------- 模型组对齐方式(默认)
batch_num ----------------------------- 处理图片的批次
sparse_rate ---------------------------- FC参数稀疏化操作(默认)
image_type ---------------------------- 输入模型的图像类型(与后处 理有关)
RGB_order ---------------------------- 输入模型图像的通道顺序
image_list ----------------------------- 参考图片
norm_type ----------------------------- 模型输入数据初始化
mean_file ------------------------------ 均值文件(一般不用加)
注意:所有参数均可在《HiSVP开发指南》102-106页进行查看,里面有更加详尽的参数说明。
打开mark_prototxt文件夹下的xxx.prototxt,删除xxx.prototxt最后几行中的Reshape和Transpose层(即model_hand中3个output区域的Reshape和Transpose层),然后保存修改后的xxx.prototxt,利用marked_prototxt中的Mark来查看修改后的xxx.prototxt(以其中一组Reshape和Transpose层为例子)。
因为YOLOv5模型采用onnx输出后的神经网络为5维(可以用netron开查看YOLOv5模型的网络结果,netron网页查看地址:https://netron.app/),而RuyiStudio只支持4维输出结构,输出形式与YOLOv5后处理方式有关。
如果wk生成成功,则会显示wk生成成功的提示,例如:
注意:如果转化失败,请查看xxx.prototxt中的网络层结构是否符合RuyiStudio所支持的网络层结构。
总的来说,模型转化并没有困难的知识点,只需要注意一些小的细节。后处理部分需要根据模型的格式进行一定的修改,这里不对后处理部分进行说明。希望大家都能顺利的转成wk文件,尽量避免一些不必要的坑。
海思开发:yolo v5的 focus层 移植到海思上的方法_tang-shopping的博客-CSDN博客_caffe focus层
海思开发:yolo v5s :pytorch->onnx->caffe->nnie_tang-shopping的博客-CSDN博客_yolov5 海思