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

文章目录

- - [1. 训练自己的模型](#1. 训练自己的模型)
  - [2. pt转onnx](#2. pt转onnx)
  - [3. onnx转rknn](#3. onnx转rknn)
  - [4. 后续......](#4. 后续……)

1. 训练自己的模型

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

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

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

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

首先拉取ultralytics的代码：

shell 复制代码

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

然后进入到目录中，进行操作：

shell 复制代码

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

如果上述pip安装过于缓慢，可以考虑换源，这里以Linux为例，其他的需要搜索解决一下。
shell 复制代码
vim ~/.config/pip/pip.conf
然后将其中的内容更换为:
ini 复制代码
[global]
index-url = https://repo.huaweicloud.com/repository/pypi/simple/

[install]
trusted-host = repo.huaweicloud.com/repository/pypi

这样ultralytics就准备完成了。

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

shell 复制代码

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

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

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

yaml 复制代码

path: /home/xxx/datasets/mydatasets  # 刚刚放数据集的目录
train: images/train  # 训练图片地址，是对于path的相对地址
val: images/val  # 验证图片地址，是对于path的相对地址
test: images/test # 测试图片地址，是对于path的相对地址（这个可选，如果只是训练，这个路径可以为空）

nc: 2 # 这个是类别的数量，根据自己的数据集进行处理！
names:
  0: person
  1: bicycle

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

yaml 复制代码

names: ['person', 'bicycle']

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

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

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

shell 复制代码

├── 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，然后有如下的内容：

python 复制代码

# 加载预训练模型
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文件的目录下打开终端，切换环境，就可以开始训练了。

shell 复制代码

conda activate deep
python ./train_my.py

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

2. pt转onnx

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

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

shell 复制代码

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

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

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

shell 复制代码

├── 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的目录下，打开终端，运行如下的命令：

shell 复制代码

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

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

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

shell 复制代码

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即可，这里直接给出创建方式：

shell 复制代码

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

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

shell 复制代码

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目录下直接打开终端进行处理：

shell 复制代码

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官方的运行模型的代码进行优化整理，提高识别的帧率、优化前后处理的代码等等......