windows下最新yolov5转ncnn教程(支持u版yolov5(ultralytics版)v5.0)

参考

1. 详细记录u版YOLOv5目标检测ncnn实现, 这篇文章作者就是ncnn作者
2. YOLOv5转NCNN过程
3. TorchScript, ONNX, CoreML Export #251, 这是yolov5官方github上的issue
 

前提

1. 编译好ncnn win, 我用的是ncnn-20210525-full-source
2. git clone yolov5源码, 我是在20210630时间直接clone的主页
3. 根据TorchScript, ONNX, CoreML Export #251-Before You Start, 
pip install -r requirements.txt  # base requirements
pip install coremltools>=4.1 onnx>=1.9.0 scikit-learn==0.19.2  # export requirements
4.安装onnx-simplifier, 因为最新版yolov5导出onnx时可以直接调用onnxsim简化模型
 

背景

在最新yolov5代码下, 如果直接按照https://zhuanlan.zhihu.com/p/275989233教程导出的模型是有问题的, 比如检测框满屏都是, 出现这个问题的原因是最新的代码在导出onnx模型时, 把anchor box decode过程也包含进来的了.
经过测试发现, 去掉anchor box decode后再导出onnx模型, 按照https://zhuanlan.zhihu.com/p/275989233教程可以顺利转成ncnn模型.
 

一, 在导出onnx时排除box decode过程.

定位到文件models/yolo.py,
找到class Detect的foward函数, 这里包括了box decode的后处理过程, 修改后加入一个export_detect来决定是否执行后处理. 最后的return也要加入export_detect条件.

def forward(self, x):
	# x = x.copy()  # for profiling
	z = []  # inference output
	for i in range(self.nl):
		x[i] = self.m[i](x[i])  # conv
		bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
		x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()

		if not hasattr(self, 'export_detect'):
			self.export_detect = True # export_detect表示要不要导出box decode处理
		if (not self.training) and self.export_detect:  # inference
		# if not self.training:  # inference
			if self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic:
				self.grid[i] = self._make_grid(nx, ny).to(x[i].device)

			y = x[i].sigmoid()
			if self.inplace:
				y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy
				y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
			else:  # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
				xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy
				wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i].view(1, self.na, 1, 1, 2)  # wh
				y = torch.cat((xy, wh, y[..., 4:]), -1)
			z.append(y.view(bs, -1, self.no))

	return x if self.training or (not self.export_detect) else (torch.cat(z, 1), x)
	# return x if self.training else (torch.cat(z, 1), x)

定位到export.py, 在ONNX export这里加入export_detect开关

    # ONNX export ------------------------------------------------------------------------------------------------------
    if 'onnx' in include:
        # export_detect表示要不要导出box decode处理
        detect_module = model.model[-1]
        detect_module.export_detect = False
        print("==============\nexclude the detect module in onnx\n==============")
        prefix = colorstr('ONNX:')
        try:
            import onnx

            print(f'{prefix} starting export with onnx {onnx.__version__}...')
            f = weights.replace('.pt', '.onnx')  # filename
            torch.onnx.export(model, img, f, verbose=False, opset_version=opset_version,
                              training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
                              do_constant_folding=not train,
                              input_names=['images'],
                              output_names=['output'],
                              dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'},  # shape(1,3,640,640)
                                            'output': {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)
                                            } if dynamic else None)

现在就可以转换onnx了.

python export.py --weights yolov5s.pt --img 640 --batch 1 --simplify

上面加入--simplify可以直接简化模型, 这样的话后面就不用单独再用onnxsim简化了.
python -m onnxsim yolov5s.onnx yolov5s-sim.onnx
 

二, onnx模型转ncnn模型

上面已经得到了简化后的yolov5s.onnx, 下一步直接用onnx2ncnn工具得到yolov5s.param和yolov5s.bin

onnx2ncnn yolov5s.onnx yolov5s.param yolov5s.bin

此时就算提示Unsupported slice step !, 得到的yolov5s.bin也是可以正常用的(当然前提是前面onnx模型没错)

然后修改yolov5s.param,

(param里第一行表示版本信息, 第二行表示层layer数量数据blob数量,
第四行开始每一行前几个参数意思大概是:
层类型-层名-输入blob数量-输出blob数量-输入blob name0-输入blob name1...-输出blob name0-输出blob name1...-具体层参数...)

修改前, 从Split层到下面的Concat层 这10层是yolov5的focus层, ncnn目前不支持这样, 需要写一个自定义的YoloV5Focus层来代替这10层, 注意我这里层的输入blob名字是images, 输出blob名字是167. 不同网络的blob也不一样:

修改后, 注意YOLOV5Focus层的输入是image输出是167, 上面的185改为176, 因为用1层替换了10层:

保存yolov5s.param后, 此时的yolov5s.param和yolov5s.bin已经可以正常使用了.

但还是可以继续用ncnnoptimize工具修正param中的blob_count, 后面的65536表示转为fp16模型(不是0表示fp32,1表示fp16吗???)
(根据param文件的第二行两个数字表示layer_count和blob_count, layer_count 要对应修改，但 blob_count 只需确保大于等于实际数量即可)

ncnnoptimize yolov5s.param yolov5s.bin yolov5s-opt.param yolov5s-opt.bin 65536

后续的使用步骤继续参考详细记录u版YOLOv5目标检测ncnn实现, 官方也给了模型的使用例子examples/yolov5.cpp

done.