RK3568笔记十九：Yolov5-seg训练部署

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

根据airockchip/yolov5 仓库训练并部署全过程记录。

1、创建环境

conda create -n yolov5_seg_env python=3.8

2、激活

 conda activate yolov5_seg_env

3、安装

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

4、下载源码

https://github.com/airockchip/yolov5

这个仓库是 rknn 官方的 airockchip/yolov5 仓库，改版本是基于 GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite （v7.0），提交 ID 为 915bbf294bb74c859f0b41f1c23bc395014ea679

该仓库的 yolov5 对 rknpu 设备进行了部署优化:

• 优化 focus/SPPF 块，以相同的结果获得更好的性能；

• 更改输出节点，从模型中删除 post_process（检测和分割后处理量化方面不友好）；

• 使用 ReLU 代替 SiLU 作为激活层（仅在使用该仓库训练新模型时才会替换）

使用该仓库直接训练导出导出适配 RKNPU 的模型，或者使用 ultralytics/yolov5 训练模

型，使用该仓库导出适配 RKNPU 的模型

这里使用该仓库直接训练部署测试。也可以自定义数据集，操作流程一样。

5、下载数据

训练的是COCO128数据集，直接下载官网的

download: https://ultralytics.com/assets/coco128-seg.zip

下载后在yolov5-master目录下创建一个datasets并解压

6、修改coco128-seg.ymal配置文件

# 1、修改数据集路径
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ./datasets/coco128-seg  # dataset root dir
train: images/train2017  # train images (relative to 'path') 128 images
val: images/train2017  # val images (relative to 'path') 128 images
test:  # test images (optional)
​
# 2、屏蔽下载链接
# Download script/URL (optional)
#download: https://ultralytics.com/assets/coco128-seg.zip

7、下载yolov5s-seg.pt文件测试

下载地址：

 https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s-seg.pt

实例分割简单测试

python segment/predict.py --weights yolov5s-seg.pt --source ./data/images/bus.jpg

导出onnx模型测试

测试的目的是看导出ONNX看模型是什么样的

python export.py --rknpu --weights G:/资料/嵌入式/【正点原子】RK3568开发板资料（A盘）-基础资料/01、AI例程/Ai_Code/31_yolov5_seg/yolov5-master/yolov5s-seg.pt --include onnx

会在yolov5-master目录下生成onnx模型

使用 Netron打开ONNX模型

输出主要分为两部分，一部分是检测结果，另一部分是分割结果。由于后处理部分量化方面不是

很好，其模型后处理部分放到 CPU 处理，模型直接输出三种尺寸的特征图。

以特征图大小 80*80 为例，相关输出节点有两个（实例分割掩码数据和目标检测分开） ：

12558080 和 1968080，前一个节点 1[3(80+5)]8080 中，3 代表了每个网格的 3 个 anchor，（5+80）代表每个 anchor 包含的信息。其中 5 表示 4 个位置坐标（x，y，w，h）和一个置信度，置信度表示这个网格中可能存在物体的概率，其中 80 表示 coco 数据集 80 个类别的分数；

后一个节点 1[332]8080 表示检测框 mask 系数。

最后的“371”节点是模型分割结果，后处理通过与目标框里面的 mask 系数矩阵乘法，处理后得

到最终的分割效果。

8、训练coco128-seg数据集

python segment/train.py --weights yolov5s-seg.pt --data coco128-seg.yaml --epochs 300 --img 640 --batch-size 16

训练300轮，租的Autodl服务器，不到1小时。

9、导出

通过 airockchip/yolov5 仓库的 export.py导出 onnx 模型。

python export.py --rknpu --weights runs/train-seg/exp3/weights/best.pt --include onnx

best.onnx 模型保存在当前目录下，并且生成了一个 RK_anchors.txt 文件，该文件保存了

anchors 的值，用在后面板卡 C++ 部署。

我们可以使用 Netron 查看下 best.onnx 模型，其模型输出如下图所示：

10、转换成 rknn 模型以及简单测试

onnx2rknn.py

import os
import sys
import numpy as np
from rknn.api import RKNN
​
# 导出rknn模型路径
RKNN_MODEL = './best.rknn'
​
# ONNX模型路径
MODEL_PATH = './yolov5s-seg.onnx'
​
DATASET_PATH = './dataset.txt'
DEFAULT_QUANT = True
​
# 默认是rk3588平台
platform = "rk3568"
​
if __name__ == '__main__':
​
    # Create RKNN object
    rknn = RKNN(verbose=False)
​
    # Pre-process config
    print('--> Config model')
    rknn.config(mean_values=[[0, 0, 0]], std_values=[
                    [255, 255, 255]], target_platform=platform)
    print('done')
​
    # Load model
    print('--> Loading model')
    ret = rknn.load_onnx(model=MODEL_PATH)
    #ret = rknn.load_pytorch(model=MODEL_PATH, input_size_list=[[1, 3, 640, 640]])
    if ret != 0:
        print('Load model failed!')
        exit(ret)
    print('done')
​
    # Build model
    print('--> Building model')
    ret = rknn.build(do_quantization=DEFAULT_QUANT, dataset=DATASET_PATH)
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('done')
​
    # Export rknn model
    print('--> Export rknn model')
    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print('Export rknn model failed!')
        exit(ret)
    print('done')
​
    # 精度分析,,输出目录./snapshot
    #print('--> Accuracy analysis')
    #ret = rknn.accuracy_analysis(inputs=['./subset/000000052891.jpg'])
    #if ret != 0:
    #    print('Accuracy analysis failed!')
    #    exit(ret)
    #print('done')
​
    # Release
    rknn.release()

转换需要在虚拟机下转换，环境根据正点原子的手册安装。

11、部署测试

将参考 rknn_model_zoo 仓库 提供的 C++ 部署例程

参考文章：RK3568笔记七：yolov5-seg实例分割测试验证_yolov5seg rknn-CSDN博客

测试正常

如有侵权，或需要完整代码，请及时联系博主。