如何将nuScenes中图像数据转换为yolov5可以训练的数据

转载

网络安全战士 2025-01-07 11:15:57

文章标签 深度学习 ncnn python github 目标检测 文章分类 架构后端开发

背景

在最新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也不一样:

如何将nuScenes中图像数据转换为yolov5可以训练的数据_ncnn