yolov5训练自己的pt文件,转onnx,再转成rknn,到RK3588开发板运行测试

一、训练自己的模型

  1. yolov5训练好自己的模型,例如训练完后,名称为best.pt,路径为runs/exp/weights/best.pt。

  1. 采用detect.py文件验证best.pt可以正常检测目标,再进行下一步工作。

二、pt转onnx

  1. 修改utils/yolo.py文件的后处理部分,将class Detect(nn.Module) 类的子函数forward由

def forward(self, x):        z = []  # inference outputfor 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()ifnot self.training:  # inferenceif self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic:                    self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)                y = x[i].sigmoid()if self.inplace:                    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]  # whelse:  # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953                    xy = (y[...,0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy                    wh = (y[...,2:4] * 2) ** 2 * self.anchor_grid[i]  # wh                    y = torch.cat((xy, wh, y[...,4:]), -1)                z.append(y.view(bs, -1, self.no))return x if self.training else (torch.cat(z,1), x)

修改为:

def forward(self, x):        z = []  # inference outputfor i in range(self.nl):            x[i] = self.m[i](x[i])  # convreturn x

注意:训练和检测的时候,yolo.py文件应保持原来的样子。

  1. 修改export.py文件

2.1修改export_onnx函数以下内容:

(1)将opset_version=opset 修改为:opset_version=12;

(2)修改outputnames为: output_names=['out378', 'out439', 'out500'])

(3)删除output_names后的dynamic段落;

2.2 修改parse_opt函数以下内容:

(1)将--data 的default修改为自己的yaml文件路径;

(2)将--weights的default修改为自己的pt文件路径:runs/exp/weights/best.pt。

(3)运行export.py文件,会在pt所在文件夹下生成onnx文件。

三、onnx转rknn

  1. 创建rknn虚拟环境进行转换。(已创建rknn虚拟环境,则跳至下一步)

需要的配置在rknn-toolkit2(-master)文件夹 的doc/requirements_cp38-1.4.0.txt中,按照以下步骤进行创建:

(1)进入rknn-toolkit2(-master)文件夹,右键打开终端,

(2)conda创建虚拟环境rknn:

conda create -n rknn python=3.8

(3)安装需要的环境:

pip install -r doc/requirements_cp38-1.4.0.txt -i https://pypi.douban.com/simple
  1. 进入rknpu2/examples/rknn_yolov5_demo/convert_rknn_demo/yolov5文件夹。

(1)将onnx文件复制到onnx_models文件夹下。

(2)将需要检测的图片例如mmj.jpg放入当前路径。

(3)将dataset.txt内容修改为mmj.jpg。

  1. 修改onnx2rknn.py文件:

(1)修改MODEL_PATH = './onnx_models/best.onnx'

(2)修改检测图片名称im_file = './mmj.jpg'

(3)修改生成的rknn文件名称RKNN_MODEL_PATH = './{}/best.rknn'.format(OUT_DIR,exp)

  1. 修改运行环境为rknn虚拟环境,运行onnx2rknn.py文件,生成的best.rknn文件在rknn_models文件夹下。

四、rknn编译

  1. 进入rknpu2/examples/rknn_yolov5_demo/convert_rknn_demo文件夹。

(1)将best.rknn文件复制到model/RK3588文件夹。

(2)将mmj.jpg 和 dataset.txt文件复制到model文件夹。

  1. 修改 /include/postprocess.h 文件中的 OBJ_CLASS_NUM 为自己训练的目标种类个数。

  1. 修改 /mode/coco_80_labels_list.txt文件内容为自己训练的种类名称。

  1. gcc编译器配置(如已喔诶之,则跳至下一步5)

(1)下载gcc编译器并解压至自定义路径(xxx/RK3588/gcc-arm-10.3-2021.07-x86_64-aarch64-none-linux-gnu)

(2)添加gcc编译器环境变量:在终端输入命令 vim /etc/profile。在文件末端添加gcc编译器路径/bin:

export PATH=$PATH:/home/yf01/ssd/RK3588/gcc-arm-10.3-2021.07-x86_64-aarch64-none-linux-gnu/bin

(3):wq 保存退出。

  1. 修改build-linux_RK3588.sh文件:

(1)修改TARGET_SOC="rk3588"

(2)修改gcc编译器路径为gcc编译器所在位置。

export TOOL_CHAIN=/home/yf01/ssd/RK3588/gcc-arm-10.3-2021.07-x86_64-aarch64-none-linux-gnu
  1. 运行环境为rknn虚拟环境,运行build-linux_RK3588.sh文件,生成install文件夹。文件夹中的rknn_yolov5_demo_Linux即为可以在RK3588开发板上运行的文件。

五、RK3588开发板上运行

  1. 将rknn_yolov5_demo_Linux文件夹复制到rk3588开发板上,进入rknn_yolov5_demo_Linux文件夹。

  1. 打开终端,输入指令:

./rknn_yolov5_demo ./model/RK3588/best.rknn ./model/mmj.jpg
  1. 运行结果为out.jpg,即mmj.jpg的检测结果。

注意:待检测图片尺寸应大于训练图片尺寸,否则可能报错。

你可能感兴趣的:(RK3588入门日志,深度学习,ubuntu,目标检测)