前言
Hello,小伙伴们~~我最近做了一个比较有意思的东西,想起来也好久没有写博客了,就记录一下吧。希望和大家一起学习,一起进步!
我简单介绍一下我最近做的这个东西的经过哈~上个月在B站上看到了一个博主发了一条视频关于边缘计算相关的,他仅仅用了一块巴掌不到的开发板完成了yolo的目标检测功能,而且延时还比较低,成本也仅有100元不到!倍感神奇哇!然后我就去学习了一下,关于NPU计算相关的知识,刚好有小伙伴手里有一块类似的开发板,于是就有了这篇博客。
PS:我用的这块开发板的NPU芯片的算力达到了6TOPS,而上述博主的开发板芯片只有0.5TOPS。但是我尽力啦~~
设备
板卡
为了避免打广告,我这里就简单放一下板卡的基本信息的介绍吧。
然后,板卡长这样:
这个板卡是自带系统的,我使用的是Ubuntu20.04的操作系统。
电脑
这个电脑需要有Ubuntu操作系统。我使用的是windows10,然后安装了Ubuntu20.04虚拟机。这台电脑主要是将pytorch生成的.pt文件转换成.onnx文件,最后转成.rknn文件。
权重转换
yolov8训练
在权重转换前,需要创建自己训练集,然后使用yolov8训练出自己训练集的权重。这里就不做过多的解释了,简单贴一下训练集中的标注文件和训练用的比较重要的两个配置文件。
标注文件解释
训练集配置(myDetData.yaml)
# 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: 训练集路径 # dataset root dir
train: images/train # train images (relative to 'path') 118287 images
val: images/val # val images (relative to 'path') 5000 images
#test: test-dev2017.txt # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794
# Classes
names:
0: hole #种类索引和具体类别名称
nc: 1 #种类数量
网络配置(yolov8s.yaml)
# Ultralytics YOLO 🚀, AGPL-3.0 license
# YOLOv8 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect
# Parameters
nc: 1 # number of classes
scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n'
# [depth, width, max_channels]
# n: [0.33, 0.25, 1024] # YOLOv8n summary: 225 layers, 3157200 parameters, 3157184 gradients, 8.9 GFLOPs
s: [0.33, 0.50, 1024] # YOLOv8s summary: 225 layers, 11166560 parameters, 11166544 gradients, 28.8 GFLOPs
# m: [0.67, 0.75, 768] # YOLOv8m summary: 295 layers, 25902640 parameters, 25902624 gradients, 79.3 GFLOPs
# l: [1.00, 1.00, 512] # YOLOv8l summary: 365 layers, 43691520 parameters, 43691504 gradients, 165.7 GFLOPs
# x: [1.00, 1.25, 512] # YOLOv8x summary: 365 layers, 68229648 parameters, 68229632 gradients, 258.5 GFLOPs
# YOLOv8.0n backbone
backbone:
# [from, repeats, module, args]
- [-1, 1, Conv, [64, 3, 2]] # 0-P1/2
- [-1, 1, Conv, [128, 3, 2]] # 1-P2/4
- [-1, 3, C2f, [128, True]]
- [-1, 1, Conv, [256, 3, 2]] # 3-P3/8
- [-1, 6, C2f, [256, True]]
- [-1, 1, Conv, [512, 3, 2]] # 5-P4/16
- [-1, 6, C2f, [512, True]]
- [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32
- [-1, 3, C2f, [1024, True]]
- [-1, 1, SPPF, [1024, 5]] # 9
# YOLOv8.0n head
head:
- [-1, 1, nn.Upsample, [None, 2, 'nearest']]
- [[-1, 6], 1, Concat, [1]] # cat backbone P4
- [-1, 3, C2f, [512]] # 12
- [-1, 1, nn.Upsample, [None, 2, 'nearest']]
- [[-1, 4], 1, Concat, [1]] # cat backbone P3
- [-1, 3, C2f, [256]] # 15 (P3/8-small)
- [-1, 1, Conv, [256, 3, 2]]
- [[-1, 12], 1, Concat, [1]] # cat head P4
- [-1, 3, C2f, [512]] # 18 (P4/16-medium)
- [-1, 1, Conv, [512, 3, 2]]
- [[-1, 9], 1, Concat, [1]] # cat head P5
- [-1, 3, C2f, [1024]] # 21 (P5/32-large)
- [[15, 18, 21], 1, Detect, [nc]] # Detect(P3, P4, P5)
训练配置(train.py)
from ultralytics import YOLO
if __name__ == '__main__':
# 加载模型
model = YOLO(r'yolov8s.yaml').load("yolov8s.pt") # 使用预训练权重训练
# 训练参数 ----------------------------------------------------------------------------------------------
model.train(
data=r'myDetectData.yaml',
epochs=300, # (int) 训练的周期数
patience=50, # (int) 等待无明显改善以进行早期停止的周期数
batch=32, # (int) 每批次的图像数量(-1 为自动批处理)
imgsz=640, # (int) 输入图像的大小,整数或w,h
save=True, # (bool) 保存训练检查点和预测结果
save_period=-1, # (int) 每x周期保存检查点(如果小于1则禁用)
cache=False, # (bool) True/ram、磁盘或False。使用缓存加载数据
device=0, # (int | str | list, optional) 运行的设备,例如 cuda device=0 或 device=0,1,2,3 或 device=cpu
workers=8, # (int) 数据加载的工作线程数(每个DDP进程)
project='runs/train', # (str, optional) 项目名称
name='exp', # (str, optional) 实验名称,结果保存在'project/name'目录下
exist_ok=False, # (bool) 是否覆盖现有实验
pretrained=True, # (bool | str) 是否使用预训练模型(bool),或从中加载权重的模型(str)
optimizer='SGD', # (str) 要使用的优化器,选择=[SGD,Adam,Adamax,AdamW,NAdam,RAdam,RMSProp,auto]
verbose=True, # (bool) 是否打印详细输出
seed=0, # (int) 用于可重复性的随机种子
deterministic=True, # (bool) 是否启用确定性模式
single_cls=False, # (bool) 将多类数据训练为单类
rect=False, # (bool) 如果mode='train',则进行矩形训练,如果mode='val',则进行矩形验证
cos_lr=False, # (bool) 使用余弦学习率调度器
close_mosaic=0, # (int) 在最后几个周期禁用马赛克增强
resume=False, # (bool) 从上一个检查点恢复训练
amp=True, # (bool) 自动混合精度(AMP)训练,选择=[True, False],True运行AMP检查
fraction=1.0, # (float) 要训练的数据集分数(默认为1.0,训练集中的所有图像)
profile=False, # (bool) 在训练期间为记录器启用ONNX和TensorRT速度
# freeze= None, # (int | list, 可选) 在训练期间冻结前 n 层,或冻结层索引列表。
# 分割
overlap_mask=True, # (bool) 训练期间是否应重叠掩码(仅适用于分割训练)
mask_ratio=4, # (int) 掩码降采样比例(仅适用于分割训练)
# 分类
dropout=0.0, # (float) 使用丢弃正则化(仅适用于分类训练)
# 超参数 ----------------------------------------------------------------------------------------------
lr0=0.01, # (float) 初始学习率(例如,SGD=1E-2,Adam=1E-3)
lrf=0.01, # (float) 最终学习率(lr0 * lrf)
momentum=0.937, # (float) SGD动量/Adam beta1
weight_decay=0.0005, # (float) 优化器权重衰减 5e-4
warmup_epochs=3.0, # (float) 预热周期(分数可用)
warmup_momentum=0.8, # (float) 预热初始动量
warmup_bias_lr=0.1, # (float) 预热初始偏置学习率
box=7.5, # (float) 盒损失增益
cls=0.5, # (float) 类别损失增益(与像素比例)
dfl=1.5, # (float) dfl损失增益
pose=12.0, # (float) 姿势损失增益
kobj=1.0, # (float) 关键点对象损失增益
label_smoothing=0.0, # (float) 标签平滑(分数)
nbs=64, # (int) 名义批量大小
hsv_h=0.015, # (float) 图像HSV-Hue增强(分数)
hsv_s=0.7, # (float) 图像HSV-Saturation增强(分数)
hsv_v=0.4, # (float) 图像HSV-Value增强(分数)
degrees=0.0, # (float) 图像旋转(+/- deg)
translate=0.1, # (float) 图像平移(+/- 分数)
scale=0.5, # (float) 图像缩放(+/- 增益)
shear=0.0, # (float) 图像剪切(+/- deg)
perspective=0.0, # (float) 图像透视(+/- 分数),范围为0-0.001
flipud=0.0, # (float) 图像上下翻转(概率)
fliplr=0.5, # (float) 图像左右翻转(概率)
mosaic=1.0, # (float) 图像马赛克(概率)
mixup=0.0, # (float) 图像混合(概率)
copy_paste=0.0, # (float) 分割复制-粘贴(概率)
)
注意事项
在后续导出onnx权重时,有一点要特别注意,就是yolov8的版本问题,使用最新版的yolov8版本可能会报错,AttributeError: 'Segment' object has no attribute 'detect'。我使用的是yolov8的8.0.151版本。(注意:该文件最好仅用于yolov8训练)
pt转onnx
这个地方我使用的是这个代码:”pt转onnx“。这个里面对应的yolov8版本就是8.0.151的,如果用的版本对不上很有可能会报 AttributeError: 'Segment' object has no attribute 'detect'的错误。另外,这个链接里面的代码还有一部分需要修改。(注意:该链接代码最好仅用于pt转onnx)
在./ultralytics/nn/modules/head.py中需要修改为:
# 导出 onnx 增加
y = []
for i in range(self.nl):
sigmoid = nn.Sigmoid()
t1 = self.cv2[i](x[i])
# 这里没有加sigmoid会出现置信度大于1的情况
t2 = sigmoid(self.cv3[i](x[i]))
y.append(t1)
y.append(t2)
a = torch.sum(t2, dim=(1, ), keepdim=True).clip(min=0, max=1)
y.append(a)
return y
如果要导出yolov8s-seg对应的onnx文件的小伙伴还需要改一个地方:
p = self.proto(x[0]) # mask protos
bs = p.shape[0] # batch size
# 导出 onnx 增加(修改)
# mc = torch.cat([self.cv4[i](x[i]).view(bs, self.nm, -1) for i in range(self.nl)], 2) # mask coefficients
mc = [self.cv4[i](x[i]) for i in range(self.nl)]
x = self.detect(self, x)
return x[0], x[1], x[2], mc[0], x[3], x[4], x[5], mc[1], x[6], x[7], x[8], mc[2], p
另外,就是导出部分需要再改一下代码。
在./ultralytics/engine/model.py文件中
if self.task == "detect":
print("=========== onnx =========== ")
import torch
dummy_input = torch.randn(1, 3, 640, 640)
input_names = ["data"]
output_names = ["reg1", "cls1", "scc1", "reg2", "cls2", "scc2", "reg3", "cls3", "scc3"]
torch.onnx.export(self.model, dummy_input, "./yolov8s_rknn.onnx", verbose=False, input_names=input_names, output_names=output_names, opset_version=11)
print("======================== convert onnx Finished! .... ")
return
if self.task == "segment":
print("=========== onnx =========== ")
import torch
dummy_input = torch.randn(1, 3, 640, 640)
input_names = ["data"]
output_names = ["cls1", "reg1", "scc1", "mc1", "cls2", "reg2", "scc2", "mc2", "cls3", "reg3", "scc3", "mc3", "seg"]
torch.onnx.export(self.model, dummy_input, "./yolov8s_seg_rknn.onnx", verbose=False, input_names=input_names, output_names=output_names, opset_version=12)
print("======================== convert onnx Finished! .... ")
return
修改完上述代码后,就可以新建一个文件(export.py)了。文件内容如下:
from ultralytics import YOLO
# 推理
model = YOLO('./best.pt')
results = model(task='detect', mode='predict', source='./20240823092916.png', line_width=3, show=True, save=True, device='0')
# results = model(task='segment', mode='predict', source='./20240823092916.png', line_width=3, show=True, save=True, device='0')
运行上面的文件就会出现下面的报错,如果报错和下方图片一样,则表明转换成功。
这里我按照官方给的文件进行修改的。
使用netron查看之后可以发现经过上述修改之后导出的文件结果和官方给出的结果是一模一样的。
detect
segment
到这,基本上就完成一半的工作了。后面就是onnx转rknn了,这一部分需要在Ubuntu下完成。
onnx转rknn
完成这一步需要先配置一下python的环境,然后下载rknn-toolkit2文件。这个文件有点大,小伙伴也可以直接下载对应的requirements_cp38-1.6.0.txt和rknn_toolkit2-1.6.0+81f21f4d-cp38-cp38-linux_x86_64.whl,这样会小很多。然后使用下面命令
cd 到requirements_cp38-1.6.0.txt路径下
pip install -r requirements_cp38-1.6.0.txt -i https://pypi.tuna.tsinghua.edu.cn/simple/
等到上述指令安装完成后
cd 到rknn_toolkit2-1.6.0+81f21f4d-cp38-cp38-linux_x86_64.whl路径下
pip install rknn_toolkit2-1.6.0+81f21f4d-cp38-cp38-linux_x86_64.whl -i https://pypi.tuna.tsinghua.edu.cn/simple/
到这一步,基本环境就配置完成了,如果后面还需要其他的包,再继续安装就好了。
下面再下载rknn_model_zoo,就可以进行onnx转rknn的模型转换了。
下载完成后,先进入到./rknn_model_zoo/examples/yolov8/python目录下,将上一步生成的onnx文件复制到该路径下。
先运行下面的指令,验证导出模型的是否正确。
python3 yolov8.py --model_path yolov8s_rknn.onnx --img_show
这里如果是使用自己的训练集的话,需要改一下yolov8.py中的CLASSES的种类。上面这种情况就表明onnx权重导出没问题。
python3 convert.py yolov8s_rknn.onnx rk3588s(或者rk3588)
出现--> Export rknn model done则表明导出成功。
rknn模型在./rknn_model_zoo/examples/yolov8/model中