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

一、任务介绍

        我的任务是在服务器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目录下

        运行命令：python export.py --weights best.pt --include onnx

        这时会在yolov5目录下获得由best.pt转换得到的best.onnx

2.2、第二步、使用虚拟机，安装rknn-toolkit2并完成格式转换

2.2.1、安装rknn-toolkit2工具包

        onnx格式只是我们需要的中间转换格式，这一步我们要完成onnx格式到rknn格式的转换

        首先我们需要使用的转换工具是rknn-toolkit2

        github地址：https://github.com/rockchip-linux/rknn-toolkit2/tree/master

                

        由于转换工具需要再linux系统上运行，所以我创建了一个虚拟机来配置环境。

        在这一部分主要环境配置上，我使用的是ubuntu-22.04.2系统，和anaconda3-2023.07-1-Linux-x86_64.sh.

        anaconda的使用方法和window环境下一致，先创建一个python3.6的基础环境，然后把对应版本的rknn-toolkit安装包：rknn-toolkit2/packages/rknn_toolkit2-1.5.2+b642f30c-cp36-cp36m-linux_x86_64.whl 和环境依赖的包的清单TXT文件rknn-toolkit2/doc/requirements_cp36-1.5.2.txt放到虚拟机系统中

        然后激活环境 conda activate py36

        安装环境包 pip install -r requirements_cp36-1.5.2.txt

        安装rknn-toolkit2工具包本身 pip install rknn_toolkit2-1.5.2+b642f30c-cp36-cp36m-linux_x86_64.whl

我写这篇笔记的时候是使用相同的文件安装，所以会提示有相同的版本。在您安装完成后可以进入python中import一下这个包，过程如下，如果没有报错，说明安装完成

2.2.2、开始转换模型格式，onnx转rknn

        首先需要把rknn-toolkit2/examples/onnx/yolov5这个文件夹中的文件下载到我们的虚拟机系统中，再把我们之前转换得来的best.onnx也放到一起

        我是用的下面这个onnx2rknn文件夹

        然后还需要修改一下test.py文件，就三个位置，1、需要转换的模型路径和测试图片的路径 2、自定义的类别 3、你要搭载的平台具体型号

        这时候再运行python test.py

        运行程序，就可以得到转换结果best.rknn和测试结果图像

然后把best.rknn复制出来，再进行最后一步，进入平台（我是用的是rk3588）实际使用rknn模型。

2.3、RK3588部署

2.3.1、基本步骤

                0, 基本信息 芯片型号：RK3588s Linux架构：aarch64

                        

                1, 获得代码 git clone GitHub - rockchip-linux/rknpu2

                2, 进入命令行安装依赖的包 sudo apt install cmake imagemagick

                3, Build example。运行代码前先生成可执行文件，进入rknpu2/examples/rknn_yolov5_demo，把我们best.rknn也复制到这个文件夹。

        在生成demo之前还需要修改两个文件，一个是include/postprocess.h 把其中的

#define OBJ_CLASS_NUM 3 #修改为你的数据集有几个类型，我的是有3个

        还有model/coco_80_labels_list.txt这个文件是写你的类型分别叫什么名字，我直接写的0、1、2，每行一个数字

这时候运行./build-linux_RK3588.sh就可以生成一个demo文件了

        ./build-linux_RK3588.sh

4, 运行demo。

./build/build_linux_aarch64/rknn_yolov5_demo best.rknn 00603.jpg

        注意一下目录即可，命令就三个部分

        1、 demo文件的位置 ./build/build_linux_aarch64/rknn_yolov5_demo

        2、rknn文件位置，我放到的是当前路径，所以直接写的名字，对应即可

        3、被测图像位置，同理，对应即可

2.3.2、修改添加labels保存功能

        这一步我是通过修改src/main.cc实现的，当然修改前一定记得备份

        这里我先分享一下我修改后的文件：

                        链接：https://pan.baidu.com/s/139a9Vcglj7sweIAkG3ZMsw 
                        提取码：knrq 

        首先我们定位到315行，可以看到这里是打印检测结果的那个句子，那么我们只要把他打印的角点坐标转换成yolo的标签格式即可，然后用outfile这句把信息写入到输出文件中，当然还有一个文件名管理和输出文件创建的过程，这个我写在前面了                   

        169行到181行就是管理输出文件名的代码，实现的是由图像名123.jpg得到输出标签名123.txt，在这里我把所有的输出都放到了一个out_labels文件夹中，我是手动创建的这个文件夹，拿测试结果的时候就是把这个文件夹复制走，再手动把里面清空

        修改过main.cc文件后还需要再按照新的main.cc生成一个新的demo

        也就是再运行一下   ./build-linux_RK3588.sh

2.3.3、批量运行命令

        这是一个待优化的点：每次都在调用demo检测一张图实现对数据集的评估，只是建立在不关注运行速度的前提下的，这可能是下一阶段的重点工作

                这样再拿测试结果去其他地方计算ap等性能评价指标即可了。