当前位置: 首页 > news >正文

训练自己的yolo模型,并部署到rk3588上

文章目录

      • 1. 训练自己的模型
      • 2. pt转onnx
      • 3. onnx转rknn
      • 4. 后续……

1. 训练自己的模型

如何训练自己的yolo模型,网络上已经有足够多的教程,这里只简单的描述一下。如果已经有了自己的.pt模型,那么可以直接跳转到第二节。

此处是以检测模型的训练为示例进行说明,其他的比如分割、姿态的模型搜索一下教程即可,一通百通。

(如何安装cuda,以及安装pytorch,这里不再进行说明,请参考网络上的教程)

假设现在你已经使用conda创建了一个deep的环境,并在其中已经安装了pytorch,然后现在在某一目录下,进行下述的操作。

首先拉取ultralytics的代码:

git clone https://github.com/ultralytics/ultralytics.git

然后进入到目录中,进行操作:

cd ultralytics
pip install -e .  # 这里是自动安装ultrlytics的依赖

如果上述pip安装过于缓慢,可以考虑换源,这里以Linux为例,其他的需要搜索解决一下。

vim ~/.config/pip/pip.conf

然后将其中的内容更换为:

[global]
index-url = https://repo.huaweicloud.com/repository/pypi/simple/[install]
trusted-host = repo.huaweicloud.com/repository/pypi

这样ultralytics就准备完成了。

然后在一个目录下专门存放你的数据集,比如在~/datasets下,并且你的数据集文件夹假设为mydatasets,然后有对应的imageslabels子目录,有如下的目录结构:

└── mydatasets├── images│   ├── test│   ├── train│   └── val└── labels├── test├── train└── val

然后我们在刚才git下载的目录下进行操作:

进入到ultralytics/ultralytics/cfg/datasets下,创建一个新的data_my.yaml文件:

path: /home/xxx/datasets/mydatasets  # 刚刚放数据集的目录
train: images/train  # 训练图片地址,是对于path的相对地址
val: images/val  # 验证图片地址,是对于path的相对地址
test: images/test # 测试图片地址,是对于path的相对地址(这个可选,如果只是训练,这个路径可以为空)nc: 2 # 这个是类别的数量,根据自己的数据集进行处理!
names:0: person1: bicycle

上面的names这一项,写法也有其他的形式,比如:

names: ['person', 'bicycle']

使用哪种方式,看你的喜好。

然后我们从git上下载预训练模型(多数时候建议使用预训练模型),访问git地址,然后在下面看到这些进行下载:

在这里插入图片描述

直接点击Model下方的蓝色字体即可下载。假设下载了s模型,然后将其放到目录中,那么可以得到如下的目录结构:

├── CITATION.cff
├── CONTRIBUTING.md
├── docker
├── docs
├── examples
├── LICENSE
├── mkdocs.yml
├── pyproject.toml
├── README.md
├── README.zh-CN.md
├── tests
├── ultralytics
└── yolo11s.pt     # 这就是刚刚下载的预训练模型

在上述的预训练模型的同级目录下新建一个文件train_my.py,然后有如下的内容:

# 加载预训练模型
from ultralytics import YOLO# 加载预训练模型
model = YOLO("yolo11s.pt")  # 使用YOLOv8 Nano版本作为起点# 训练模型
results = model.train(data="data_my.yaml",  # 使用自定义数据集进行训练epochs=500,  # 训练500个周期imgsz=640,  # 调整图像大小batch=16,  # 减少批次大小)  

然后在py文件的目录下打开终端,切换环境,就可以开始训练了。

conda activate deep
python ./train_my.py

然后就可以在runs/detect下面找到对应的训练结果,得到权重的文件夹,然后有best.pt文件了。

2. pt转onnx

得到了我们想要的best.pt文件,我们就需要将其转换为best.onnx文件,但是需要使用rk的官方代码来进行转换,否则将导致模型的裁剪不正确的问题。

在合适的目录下,打开终端,拉取rk的转换代码:

git clone https://github.com/airockchip/ultralytics_yolo11.git

注意看到这个比之前的仓库名字多了_yolo11字样。

进入到ultralytics_yolo11目录之后,我们新建一个models的文件夹(这里只是为了方便,没有要求),然后将我们的best.pt拷贝进来。大致的结构如下:

├── LICENSE
├── mkdocs.yml
├── models
│   └── best.pt # 这里就是我们的模型
├── pyproject.toml
├── README.md

将我们自己数据集的yaml文件(对应第一节的data_my.yaml)拷贝到目录: /home/xxx/rknn/ultralytics_yolo11/ultralytics/cfg/datasets

这里就是示意的一个目录,具体参考自己的目录。

然后我们修改/home/xxx/rknn/ultralytics_yolo11/ultralytics/cfg/default.yaml文件:

在这里插入图片描述

将这里的两个路径对应修改。

然后回到ultralytics_yolo11的目录下,打开终端,运行如下的命令:

conda activate deep
export PYTHONPATH=./
python ./ultralytics/engine/exporter.py

然后就可以在models目录下得到best.onnx文件了。

如果运行报错没有onnx模块,使用命令安装:

pip install onnx

另外,可以验证一下onnx模型是否转换正确,可以打开netron,然后导入这个onnx模型,然后点击images,需要出现类似这样的output才是正确的!

在这里插入图片描述

3. onnx转rknn

得到上述的best.onnx文件之后,我们需要进行转换为合适的rknn模型。

准备一个rknn的环境:

这里需要参考rknn官方文档,具体看其中的2.1.2.3 创建RKNN-Toolkit2 Conda环境2.1.2.4 安装 RKNN-Toolkit2即可,这里直接给出创建方式:

conda create -n rknn python=3.8
conda activate rknn
pip install rknn-toolkit2

在合适的目录下拉取代码:

git clone https://github.com/airockchip/rknn_model_zoo.git

然后我们将刚才的best.onnx文件的绝对路径记住,我们进入到rknn_model_zoo目录下,然后修改rknn_model_zoo/examples/yolo11/python/yolo11.py文件的内容:

在这里插入图片描述

然后在这个yolo11.py目录下直接打开终端进行处理:

conda activate rknn
python convert.py /home/xxx/rknn/ultralytics_yolo11/models/best.onnx rk3588 i8 /home/xxx/rknn/ultralytics_yolo11/models/best.rknn

这样就能在刚刚的models文件夹下看到best.rknn文件了,这个就可以在rk3588上进行使用了。

4. 后续……

后续会考虑将rknn官方的运行模型的代码进行优化整理,提高识别的帧率、优化前后处理的代码等等……

相关文章:

  • HTTPS 协议:数据传输安全的坚实堡垒
  • Prompt工程:解锁大语言模型的终极密钥
  • vscode的Embedded IDE创建keil项目找不到源函数或者无法跳转
  • 服务器中分布式存储数据技术都包含哪些内容?
  • 攻防世界逆向刷题笔记(新手模式9-1?)
  • Android15 Camera Hal设置logLevel控制日志输出
  • Android 15 控制亮屏灭屏接口实现
  • Cherry Studio连接配置MCP服务器
  • Android 启动优化
  • Python os 模块简介及基础使用示例
  • 在PyTorch中,对于一个张量,如何快速为多个元素赋值相同的值
  • 【笔记】解决ImportError: cannot import name ‘Iterable‘ from ‘collections‘
  • 【会议推荐】2025年模式识别与大数据国际会议(PRBD 2025)
  • 华为认证中HCIA/HCIP/HCIE是什么等级?怎么考试?
  • #跟着Lucky学鸿蒙# HarmonyOS NEXT 工程介绍
  • ES 在大查询场景下导致 GC 频繁,如何定位和解决?
  • 用 Python 打造你的专属虚拟试衣间!——AI+AR 如何改变时尚体验
  • 模型评价指标介绍
  • emqx、MongoDB或者java程序,出现 Too many open files 问题
  • Flink系列文章列表
  • 青海微网站建设/模板之家官网
  • 设计专业考研平面素材网站/什么是核心关键词
  • 酒店网站建设设计/上海seo网站排名优化公司
  • 深圳建网建网站/谷歌seo外链
  • 有关建设旅行网站的建设/深圳网络营销外包公司推荐
  • 现代著名设计师及作品/手机优化软件哪个好