yolov5 转换为rknn模型在3588上运行

为了把yolov5在rk3588上跑起来,在网上搜罗了一圈,踩了一些坑。由于瑞芯微的文档有升级,导致和网络的文章有出入,所以做个记录。

rknn-toolkit 转换文档:

瑞芯微的转换文档在 rknn-toolkit/example/pytorch/yolov5/REAME.md 里

当版本升级到 1.7.5的时候,文档被删除,说是要根据airockchip/rknn_model_zoo (github.com) 里的指令来完成。但是model_zoo里的文档也是没有说太清楚。

把rknn-toolkit版本退回到,在 1.7.3 版本中还是有的。下面就是按照1.7.3的说明来操作。

yolov-v50 的环境:

官方文档指定使用yolov5 v5.0,因此需要获取特定的版本。

需要安装 pytorch 1.10.0 以下版本,否则会报错:

AttributeError: ‘Upsample’ object has no attribute ‘recompute_scale_factor’

但是在python 3.10 版本下没有 pytorch 1.9.1 了,所以要装python 3.8的环境:

conda create -n YOLOV5-V50 python==3.8

然后把requirement.txt 中的torch 改成 torch<1.10.0

pip3 install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simpleSimple Indexicon-default.png?t=N6B9https://pypi.tuna.tsinghua.edu.cn/simple

权重文件:

在执行 python detect.py 默认下载的是最新版本的权重文件,我们需要v5.0的权重文件,所以需要手动下载

wget https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5s.pt

检测:

默认使用coco数据集,yolov5s.pt

python detect.py

修改模型,导出

重点是RKNN把最终的sigmoid从模型内部给拿到外部,把sigmoid交给cpu去计算了。具体的参见 rknn-toolkit/example/pytorch/yolov5/REAME.md

3. 直接使用pt模型转为rknn模型时,需要修改 yolov5/models/yolo.py文件的后处理部分,将class Detect(nn.Module) 类的子函数forward修改为:

def forward(self, x):

      z = [] # inference output

      for i in range(self.nl):

           x[i] = self.m[i](x[i]) # conv

      return x ```

导出

python export.py --weights yolov5s.pt --img 640 --batch 1 --include onnx

测试和转换

测试:

测试的时候生成的rknn文件是用于3566的,不知道怎么给弄成3588的,所以测试只能测试

然后把这个.onnx 文件放到

rknn-toolkit2/examples/onnx/yolov5/

修改 test.py,转换为rknn模型,并查看 result.jpg是不是正确

转换为3588的:

转换的脚本不在rknn-toolkit2下,而在 rknpu2/examples/rknn_yolov5_demo/convert_rknn_demo/yolov5 下,修改

platform = 'rk3588' 和 model path,然后运行脚本可以得到rknn

用这个rknn替换掉板子上的文件,可以运行,识别率似乎下降了一点,发现推理时间从20多ms变为30ms左右。

在测试的时候,发现官方提供的onnx是81层,而我们转化出来的是96层,可能延时时间长跟这个有关系。

你可能感兴趣的:(AI,YOLO,rknn)