海思 YOLOv5 pytorch 转 onnx 转 Caffe 再转 wk 的转化详解

目录

        前沿

        YOLOv5模型的选取与修改

        YOLOv5 pytorch 转 onnx 转 Caffe

        YOLOv5 Caffe转wk文件

        总结

        参考


前沿

        作者在将YOLOv5 pytorch版本转成wk文件中碰到了很多的坑,但在研究中也收获了很多知识,针对转化过程,作者进行了详细的解读来给大家避坑,如果在转化过程中有什么问题也可以联系作者,我们可以一起进行分析和讨论。

YOLOv5模型的选取与修改

        现如今,官方最新版本的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 pytorch 转 onnx 转 Caffe

        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函数下面的内容。

例如:

海思 YOLOv5 pytorch 转 onnx 转 Caffe 再转 wk 的转化详解_第1张图片

         其中,convertToCaffe函数选用最后一个即可,主要原因是Focus网络层在训练的时候已经被改为了Conv层,所以,无需再用其他函数。

如果转化成功,会出现以下信息:

海思 YOLOv5 pytorch 转 onnx 转 Caffe 再转 wk 的转化详解_第2张图片

        注意:如果转化成功,便会出现类似以下的信息,如果转化失败,代码会报错。转化失败的原因可能是因为前面生成的simple_onnx不满足转化条件,存在Caffe中无法解析的网络层,或是在生成simple_onnx前修改了yolo.py、common.py等源码,或是onnx的版本号可能不符合需求,或是无法生成节点。因此在转化模型前,先用netron(https://netron.app/)(只需将onnx文件导入即可)来看一下神经网络结构。

YOLOv5 Caffe转wk文件

        安装并下载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张模型测试的图片。

海思 YOLOv5 pytorch 转 onnx 转 Caffe 再转 wk 的转化详解_第3张图片

         打开文件最下方的 xxx.cfg,导入caffe模型文件,并设置模型转化的对应参数。

海思 YOLOv5 pytorch 转 onnx 转 Caffe 再转 wk 的转化详解_第4张图片

参数说明: 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 pytorch 转 onnx 转 Caffe 再转 wk 的转化详解_第5张图片                            修改后:海思 YOLOv5 pytorch 转 onnx 转 Caffe 再转 wk 的转化详解_第6张图片

         因为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 海思

你可能感兴趣的:(c++,目标检测,华为,caffe,pytorch)