yolov5训练出pt模型文件,转换成rknn格式并测试。

一、任务介绍

        我的任务是在服务器GPU上训练yolov5后,转换到使用NPU的平台上,对数据集进行测试,查看MAP等模型性能的损失,目前阶段还没有开展检测速度的评估。

        主要分为三个步骤

        1、在服务器GPU上训练出yolov5的模型文件,best.pt,然后使用yolov5中的export转换成onnx格式。

        2、在linux虚拟机上配置python环境并安装格式转换工具rknn-toolkit2,完成onnx格式转换rknn格式。

        3、在NPU上配置rknn_yolov5_demo,修改postrocess.h、main.cc等文件,添加label保存功能。编写sh脚本,实现对一个数据集文件夹检测。

二、具体步骤

2.1、第一步、使用GPU训练yolov5,并把pt文件转换成onnx格式

        训练过程不是格式转换的重点,这里不赘述,只说一些格式转换过程中的注意事项。

2.1.1、首选训练pt的python环境需要向使用的yolov5版本看齐,不必和后面rknn-toolkit2使用的python版本相同。

        比如:我使用的是yolov5-7.0版本、python版本使用的是python 3.10

        而配置rknn-toolkit2环境时使用的python版本是python 3.6

2.1.2、修改yolov5/models/yolo.py文件

        删掉class Detect(nn.Module)类的forward函数(yolo.py文件的第56行,当然yolov5版本不同行数可能会有变化)中一部分,修改结果:

                

        注意:训练的过程中不要使用这个修改的yolo.py训练,这个修改只是rknn格式转换过程中的要求。

2.1.3、运行yolov5/export.py文件

        将你训练yolov5得到的best.pt文件复制到yolov5目录下

评论 9
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值