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

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

丶Endeavor

7619人浏览 · 2022-07-14 15:50:50

丶Endeavor · 2022-07-14 15:50:50 发布

前沿

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

例如：

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

如果转化成功，会出现以下信息：

注意：如果转化成功，便会出现类似以下的信息，如果转化失败，代码会报错。转化失败的原因可能是因为前面生成的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张模型测试的图片。

打开文件最下方的 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层为例子）。

修改前：修改后：