栏目分类:
子分类:
返回
名师互学网用户登录
快速导航关闭
当前搜索
当前分类
子分类
实用工具
热门搜索
名师互学网 > IT > 软件开发 > 后端开发 > Python

yolov5模型转换为rknn模型

Python 更新时间: 发布时间: IT归档 最新发布 模块sitemap 名妆网 法律咨询 聚返吧 英语巴士网 伯小乐 网商动力

yolov5模型转换为rknn模型

目录

1.更改模型框架

2.重新训练模型

3.将重新训练的模型转换为onnx模型

4.将onnx模型转为rknn模型

5.pytorch模型保存注意事项


Yolov5官网模型训练框架为pytorch1.7.0,并且模型框架中部分算子在RKNN_Toolkit_V1.6.0中是不支持的,因此常规的模型转换过程是无法进行模型转换的。这里我们做一个详细的说明。

1.更改模型框架

更改模型框架,也就是替换掉不支持的模型算子,参考https://github.com/airockchip/yolov5/ 下面的代码,主要是对models/commom.py 进行以下修改:

第30行 将Conv定义中的激活函数进行修改self.act = nn.Hardswish() if act else nn.Identity() 替换为self.act = nn.ReLU() if act else nn.Identity()

第62行 将self.act = nn.LeakyReLU(0.1, inplace=True)替换为self.act =        nn.ReLU()

第104行,将self.conv = Conv(c1 * 4, c2, k, s, p, g, act) 替换为self.conv = Conv(c1, c2, k, 2, p, g, act)

第107行 将return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))替换为return self.conv(x)

2.重新训练模型

参考 train.py ,需要更换算法模型和训练用的图片和标签。

3.将重新训练的模型转换为onnx模型

转换onnx参考:2onnx.sh。

4.将onnx模型转为rknn模型

转换rknn参考:rknn/onnx2rknn.py。

注意要修改rknn/onnx2rknn.py里面的target_platform。

if __name__ == '__main__':

    # Create RKNN object
    rknn = RKNN(verbose=True)

    # pre-process config
    print('--> config model')
    rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], reorder_channel='0 1 2', target_platform=['rv1126'], output_optimize=1)
    print('done')

5.pytorch模型保存注意事项

在使用 load_pytorch 接口加载 pytorch 模型时, pytorch 模型只 能 用 torch.jit.trace 保 存 的 模 型。

用torch.save(net.state_dict())保存的模型,保存的只是模型的参数,并没有网络结构。只凭一个网络参数, RKNN Toolkit 无法构建相应的网络。

说明:更多模型转换注意事项见《Rockchip_Trouble_Shooting_RKNN_Toolkit_V1.6.0_CN.pdf》。

参考文献:

https://github.com/airockchip/yolov5

https://github.com/rockchip-linux/rknn-toolkit/tree/master/doc

转载请注明:文章转载自 www.mshxw.com
本文地址:https://www.mshxw.com/it/324203.html
我们一直用心在做
关于我们 文章归档 网站地图 联系我们

版权所有 (c)2021-2022 MSHXW.COM

ICP备案号:晋ICP备2021003244-6号