yoloV5-v3.0-模型转换,pt->onnx->ncnn

个人微信公众号

在这里插入图片描述

版本要求

yolov5的版本是v3.0

pytorch==1.5.0
torchvision==0.6.0

模型训练

网上有很多关于yoloV5训练的教程,可自行百度,如果实在不想找,可以参考下面的几个链接。

pt转onnx

训练结束以后,可以得到last.pt文件,接下来就是如何将pt文件转成为onnx文件了。

修改common.py

注意!YOLOv5在转onnx的时候,有个巨坑,就是在网络模型中有个Focus机制,其中有一个tensor的切片操作,这个操作在转ncnn的时候不被支持,需要对这个模型的切片操作进行修改!
具体的修改是在models文件夹下的common.py文件,其原始代码如下:

class Focus(nn.Module):
    # Focus wh information into c-space
    def __init__(self, c1, c2, k=1):
        super(Focus, self).__init__()
        self.conv = Conv(c1 * 4, c2, k, 1)

    def forward(self, x):  # x(b,c,w,h) -> y(b,4c,w/2,h/2)
        # return self.conv(torch.cat([x, x, x, x], 1)) 
        return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))

修改为

class Focus(nn.Module):
    # Focus wh information into c-space
    def __init__(self, c1, c2, k=1):
        super(Focus, self).__init__()
        self.conv = Conv(c1 * 4, c2, k, 1)

    def forward(self, x):  # x(b,c,w,h) -> y(b,4c,w/2,h/2)
        return self.conv(torch.cat([x, x, x, x], 1)) 
        # return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))

转换命令

 export PYTHONPATH="$PWD" && python models/onnx_export.py --img-size 320 320

到这里为止,就已经将pt模型转为了onnx模型,可以说已经脱离了pytorch了,但是转ncnn还需要对模型做一个Simplifier操作,因为转出的onnx模型还有许多冗余,这在ncnn里也是不支持的,避免转ncnn时各种报错,就先操作一下把!

python -m onnxsim  yolov5s.onnx yolov5s-sim.onnx

onnx转ncnn

这里直接用大佬开发好的转换工具直接转换

https://convertmodel.com/

在这里插入图片描述
如果没有报错,恭喜你转换成功,如果没有成功,那再返回到上面的步骤,多试几次吧~~~

参考

  • 7
    点赞
  • 36
    收藏
    觉得还不错? 一键收藏
  • 4
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 4
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值