RK3568笔记五：基于Yolov5的训练及部署

news2024/11/20 7:01:19

若该文为原创文章，转载请注明原文出处。

一. 部署概述

环境：Ubuntu20.04、python3.8

芯片：RK3568

芯片系统：buildroot

开发板：ATK-DLRK3568

开发主要参考文档：《Rockchip_Quick_Start_RKNN_Toolkit2_CN-1.4.0.pdf》、《Rockchip_User_Guide_RKNN_Toolkit2_CN-1.4.0.pdf》

二、yolov5模型训练

1、训练环境

训练是在云端训练的，平台AutoDL租了一台2080IT，配置如下：

2、环境搭建

1、创建conda环境

conda create -n rkyolov5 python=3.8       // 用于yolov5
conda create -n rknn2_env python=3.8 -y   // 用于rknn2

2、激活conda环境

conda activate rkyolov5

conda deactivate  // 退出环境

3、安装pytoch

pip install torch==1.7.1+cu101 torchvision==0.8.2+cu101 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html -i https://pypi.tuna.tsinghua.edu.cn/simple

4、下载yolov5-v6.0

首先需要在官网下载yolov5-6.0的项目我们打开yolov的官网，Tags选择6.0版本

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

5、训练

python train.py --img 640 --data coco128.yaml --cfg yolov5s.yaml --weights yolov5s.pt --epoch 300 --batch-size 16 --device 0

出错：The aliases was originally deprecated in NumPy 1.20; for more details and guidance see the original release note at:

原因numpy版本问题

处理，重新安装

pip uninstall numpy
 
pip install numpy==1.22

出错：AttributeError: module 'PIL.Image' has no attribute 'ANTIALIAS

ModuleNotFoundError: No module named 'PIL'

原因是pillow库版本不支持，降低版本

pip install pillow==9.5.0

6、pt转onnx

转换步骤：

修改models/yolo.py，修改class Detect(nn.Module):的forward函数

注意！！！仅在转换时修改，在训练时改回原状态！再训练时不要忘记哦！

# def forward(self, x):
#     z = []  # inference output
#     for 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()
#
#         if not self.training:  # inference
#             if self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic:
#                 self.grid[i] = self._make_grid(nx, ny).to(x[i].device)
#
#             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]  # wh
#             else:  # 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].view(1, self.na, 1, 1, 2)  # 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 output
    for i in range(self.nl):
        x[i] = self.m[i](x[i])  # conv

    return x

修改export.py函数的--opset为12

运行export.py

python export.py --weights best.pt --img 640 --batch 1 --opset 12

简化模型

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

三、安装rknn-toolkit2

根据正点原子的教程安装，先安装工具链，在创建环境，在安装rknn-toolkit2，之后转换模型

下载地址：GitHub - rockchip-linux/rknn-toolkit2

1、创建一个新的环境

conda create rknn2_env python=3.8 -y

2、安装

进入packages

pip install rknn_toolkit2-1.5.2+b642f30c-cp38-cp38-linux_x86_64.whl

等待安装

把转换好的onnx 拷贝到 rknn-toolkit2-master/examples/onnx/yolov5目录下，

3、修改test.py文件

ONNX_MODEL = 'yolov5s_relu.onnx'
RKNN_MODEL = 'yolov5s_relu.rknn'

4、运行

python test.py

出错： ImportError: /lib/x86_64-linux-gnu/libm.so.6: version `GLIBC_2.29' not found (required by /root/miniconda3/envs/rknn2_env/lib/python3.8/site-packages/rknn/api/lib/linux-x86_64/cp38/librknnc.so)

原因：工具链没安装，安装后测试正常。