注意:首先确保本地环境支持YOLOv5能够正常运行。
链接直达
python export.py --weights yolov5s.pt --img 672 --batch 1 --opset=11
对于 YOLOv5 模型,我们在模型结构上的修改点主要在于几个输出节点处。由于目前的浮点转换工具链暂时不支持 5 维的 Reshape,所以在 prototxt中进行了删除,并将其移至后处理中执行。同时还添加了一个 transpose 算子,使该节点将以 NHWC 进行输出。这是因为在地平线芯片中, BPU 硬件本身以 NHWC 的layout 运行,这样修改后可以让 BPU 直接输出结果,而不在量化模型中引入额外的transpose。在转换onnx模型时需要先对代码进行修改。具体流程如下:
①、修改models目录下yolo.py
# x = x.copy() # for profiling
z = [] # inference output
self.training |= self.export
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()
x[i] = x[i].permute(0, 2, 3, 1).contiguous() #此项为新修改内容
if not self.training: # inference
if self.grid[i].shape[2:4] != x[i].shape[2:4]:
self.grid[i] = self._make_grid(nx, ny).to(x[i].device)
y = x[i].sigmoid()
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
z.append(y.view(bs, -1, self.no))
return x if self.training else (torch.cat(z, 1), x)
②修改export.py脚本(此脚本是导出onnx模型使用的)
由于地平线AI工具链支持的ONNX opset版本为 10 和 11, 请将 torch.onnx.export 的 opset_version 参数根据您要使用的版本进行修改。
将 torch.onnx.export 部分的默认输入名称参数由 ‘images’ 改为 ‘data’,与模型转换示例包的YOLOv5示例脚本保持一致。
将 parser.add_argument 部分中默认的数据输入尺寸640x640改为模型转换示例包YOLOv5示例中的672x672。
具体详见地平线官方手册:模型转换说明
修改完成执行上述指令即可完成onnx模型的转换,转换完成之后生成如下文件:
以下在转换过程中需注意:
①、报错:报错AttributeError: Can’t get attribute ‘SPPF’
解决方法
②、在模型转换时需要将export.py从models文件夹中复制到源码的根目录下使用。或者将转换指令修改为:
python models/export.py --weights yolov5s.pt --img 672 --batch 1 --opset=11
转换过程详见官方手册:模型转换手册
或者使用官方的DDK包中已经写好的转换脚本进行意见模型转换
路径如下:/ddk/samples/ai_toolchain/horizon_model_convert_sample/04_detection/03_yolov5/mapper
下面说一下在转换BIN模型过程中我遇到的问题:
在使用地平线模型转换工具转换完成模型之后,进行单张图片推理的过程中遇到的问题。
按照官方5.0版本修改相应文件后,进行第四步sh 04_inference.sh推理时报错如下:
原因是在进行第四步模型推理时是使用的上一步模型转换后生成的以下两种模型的其中一个来进行推理(此模型是专门用作推理使用的)
使用Netron工具查看最终模型的数据数据(或者直接查看.pt模型转换完成onnx模型后的结果)最后是33
根据这一项数据更改后处理文件postprocess.py,存放在如下图所在位置:
在后处理文件postprocess.py中修改标签数量、reshape形状
因为转换完成的onnx模型最后以为数据是33,因此再此处需要修改为一致的但维度不能减少,所以修改为11,11x3=33
coco_metric.py中修改:(根据自己模型标签情况而定)
coco_classes.names中修改:(根据自己模型标签情况而定)
到此修改后处理操作完成,再次运行sh 04_inference.sh进行推理得到结果:
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-A1J8C5KH-1662096991377)(_v_images/20220831142234717_3222.png)]