训练自己的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
,然后有对应的images
和labels
子目录,有如下的目录结构:
└── 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
官方的运行模型的代码进行优化整理,提高识别的帧率、优化前后处理的代码等等……