YOLO V5
训练模型部署到瑞芯微的板子上面,官方是有给出案例和转过详情的。并且也提供了Python
版本的推理代码,以及C语言
的代码。
但是,对于转换过程中的细节,哪些需要改?怎么改?如何改,和为什么这样改的问题,并没有给出详细的介绍。于是,本文就是对官方给出部分外的一个补充。这部分都是踩过坑的总结,相信会对你的操作会有较大帮助的。
一、从pytorch
的pt
到rknn
转换
- 第一步: 使用
yolov5
提供的export.py
函数导出yolov5.onnx
模型
python3 export.py --weights yolov5s.pt --img-size 640 --include onnx
- 第二步:使用
onnxsim
简化导出的yolov5.onnx
模型
onnxsim是一个基于ONNX规范的工具,通过简化ONNX模型和优化ONNX模型,帮助用户减小模型大小、提高模型的推理速度和减少推理过程中的内存开销。
onnxsim的工作原理是将一个ONNX模型简化成最少的节点,并优化这些节点,以最小化推理过程中的开销。
同时,onnxsim还可以处理支持的神经网络层类型,支持多个平台,例如:CPU,GPU, FPGA等。
onnxsim
安装和使用:onnx-simplifier
pip3 install onnxsim
Then:
onnxsim input_onnx_model output_onnx_model
- 第三步:要完全使用rknn提供的部署转换代码,需要根据简化后的onnx模型,选取合适层的输出,以替代以下代码中的
‘378’,‘439’和‘500’
,如下图onnx例子中的'onnx::Reshape_446',‘onnx::Reshape_484’,‘onnx::Reshape_522’
。(这三个name,可能都是不一样的,是什么就填什么即可)
# Load ONNX model
print('--> Loading model')
ret = rknn.load_onnx(model=ONNX_MODEL, outputs=['onnx::Reshape_446', 'onnx::Reshape_484', 'onnx::Reshape_522'])
if ret != 0:
print('Load yolov5 failed!')
exit(ret)
print('done')
采用Netron
打开的onnx
文件,如下:
疑问:为什么不用最后合并后的输出结果?
因为,最后的形状不固定导致的,有可能5个框,有可能10个框
。输出模型到固定大小,后续操作放到后处理,目的是为了加快模型的npu上
的推理速度(这里是我的理解,不一定正确,欢迎补充)
在
PyTorch
中,神经网络的输出形状通常是根据输入形状来自动计算的,而在ONNX
中,输出形状需要在转换时进行显式指定,这是由于ONNX
的静态图执行模型与PyTorch
的动态图执行模型不同所致。
当你将PyTorch
模型转换为ONNX
模型时,你需要为ONNX
模型中的每个输出定义固定的形状,以便在模型执行时为其分配正确的内存空间。如果输出形状不固定,那么ONNX
运行时就需要在运行时动态调整输出形状,这将使得模型在部署时的性能受到影响。
因此,在转换PyTorch
模型为ONNX
模型时,你需要手动指定每个输出的固定形状,以便在执行时能够顺利运行。
Yolo v5
的输出格式一般为a × b × c × 85
的形式,其中:
a*b*c
表示框的数目85
则涵盖框的位置信息(xc,yc,w,h)
、前景的置信度Pc
和80
个类别的预测条件概率c1,...,c80。(4+1+80
,无背景类)
如果是你自己的模型,可能是只有3个目标类别,那么最后就是4+1+3=8,这个值记得在onnx
模型中查看到。
二、需要注意事项
2.1、 设定anchor
值
anchor
的设定,在训练yolo v5
模型时候,是可以设定自动适应,采用聚类的方式,通过标注的目标框的大小,给出anchor
的值。在train.py
中,noaotoanchor
的默认为False
,如果设定为True
,则会使用默认的anchor
设定。
所以,如果经过autoanchor
,给出了新的anchor
设定,那么在推理和转完rknn
后的设定,都需要与之相匹配的anchor
,这个很重要。
为什么官方和很多博客,都没有注意到这个问题呢?因为大多数情况下,aotoanchor
并没有发挥作用。都是使用了默认的,导致很多人即便没有注意到这个问题,最后的结果也不差。
但是,如果是不一样的,结果就会比较差,这个值就需要对应的做修改了。
2.1.1、训练阶段记录
如果在训练阶段,你已经关注到autoAnchor
的输出结果,可以在这里直接进行记录,在terminal
打印的内容,大致如下:
AutoAnchor: 3.60 anchors/target, 0.974 Best Possible Recall (BPR). Anchors are a poor fit to dataset ⚠, attempting to improve...
AutoAnchor: WARNING ⚠ Extremely small objects found: 764 of 27545 labels are <3 pixels in size
AutoAnchor: Running kmeans for 9 anchors on 27522 points...
AutoAnchor: Evolving anchors with Genetic Algorithm: fitness = 0.8052: 100%|██████████| 1000/1000 00:10
AutoAnchor: thr=0.25: 0.9996 best possible recall, 5.11 anchors past thr
AutoAnchor: n=9, img_size=640, metric_all=0.358/0.805-mean/best, past_thr=0.532-mean: 5,5, 7,8, 11,11, 17,17, 28,28, 41,37, 56,56, 79,82, 143,140
2.1.2、pt文件查询记录
查询autoAnchor
记录到.pt
文件内的anchor
设定,如下:
import torch
import sys
sys.path.append("path/yolov5-master")
weights = 'best.pt'
model = torch.load(str(weights[0] if isinstance(weights, list) else weights), map_location='cpu')
model1 = model['ema' if model.get('ema') else 'model']
model2 = model1.float().fuse().model.state_dict()
for k,v in model2.items():
if 'anchor' in k:
# print(k)
# print(v)
print(v.numpy().flatten().tolist())
打印结果:
Fusing layers...
[0.54345703125, 0.58251953125, 0.8525390625, 0.88818359375, 1.353515625, 1.318359375, 1.0859375, 1.0380859375, 1.75390625, 1.705078125, 2.38671875, 2.462890625, 1.7421875, 1.6787109375, 2.578125, 2.458984375, 3.904296875, 3.75]
[4.34765625, 4.66015625, 6.8203125, 7.10546875, 10.828125, 10.546875, 17.375, 16.609375, 28.0625, 27.28125, 38.1875, 39.40625, 55.75, 53.71875, 82.5, 78.6875, 124.9375, 120.0]
YOLOv5m summary: 308 layers, 21037791 parameters, 0 gradients
第二行是真的,需要取整。第一行…
经过我的发现,如果你打印的anchor
就一行,那么可能是默认的anchor
(默认使用COCO
数据集的anchor
),就是good fit to dataset
,也就是默认的:
[[10, 13], [16, 30], [33, 23],
[30, 61], [62, 45],[59, 119],
[116, 90], [156, 198], [373, 326]]
2.2、rk3588推理性能
yolo v5m
量化前性能:
推理性能:
Performance
Total Time(us): 194162
FPS: 5.15
占用内存:
Memory Profile Info Dump
NPU model memory detail(bytes):
Total Weight Memory: 39.83 MiB
Total Internal Tensor Memory: 19.50 MiB
Total Memory: 59.33 MiB
量化后性能
推理性能:
Performance
Total Time(us): 137508
FPS: 7.27
占用内存:
Memory Profile Info Dump
NPU model memory detail(bytes):
Total Weight Memory: 20.03 MiB
Total Internal Tensor Memory: 8.75 MiB
Total Memory: 28.78 MiB
总的来说:
- 模型时间效率上,量化后能降低
30%
,194ms
到137ms
; - 占用内存上,量化后减少
50%
,59Mib
到29Mib
;
三、C/C++ API部署
-
目标检测 YOLOv5 - 基于 瑞芯微 Rockchip RKNN C API 实现 ----------- github代码
-
yolov8 瑞芯微 RKNN 的 C++部署------------- github代码
上述两个参考链接,基本囊括了一下几个部分:
rknn
模型转换Python rknn
推理c/c++ rknn
推理(YOLO v5
部分是瑞芯微官方开放的代码)
如果你也是参考瑞芯微官方的C API
代码,那么替换上你的模型后,有几个地方需要修改:
- 输入图像大小要改
anchor
尺寸要改
const int anchor0[6] = {4, 5, 7, 7, 11, 11};
const int anchor1[6] = {17, 17, 28, 27, 38, 39};
const int anchor2[6] = {56, 54, 83, 79, 125, 120};
- 前景
box
阈值修改
const float box_conf_threswin = 0.25;
nms
阈值修改
const float nms_threswin = 0.1;
- 类别置信度重新调整
objProbs.push_back(current_prob*box_confidence);
- 针对各个类,采用不同的阈值(待补充,这部分瑞芯微未采用这种二次过滤方式)
尤其是anchor
这里,如果设定的不对,那么输出的结果就会非常的奇怪。如果是对的,那么差异性相对会小很多(和本地pt
测试结果对比)。
四、总结
本文是对YOLO V5
模型部署到瑞芯微板子上遇到的问题汇总。当然可能还会存在其他的更多问题,但是暂时还没有遇到,所以后面如果还会遇到什么问题,还会补充到这里。
如果你也正在做这块,并且遇到了问题,可以评论交流。目前还发现就是转模型后的评估问题,这个后面也会按照官方教程进行测试,这是下一篇的预告,期待。