当前位置:   article > 正文

YOLOv8部署在RKNN3588上_rk3588 yolov8

rk3588 yolov8

1. onnx导出

官方YOLOv8有提供export脚本,运行即可。

  1. yolo export model=yolov8n.pt format=onnx # export official model
  2. yolo export model=path/to/best.pt format=onnx # export custom trained model

注意,onnx要转成rknn模型,目前不支持太高版本的opset算子库,YOLOv8默认是opset=17,我们需要修改成12。

  1. yolo export model=path/to/best.pt format=onnx opset=12

2. onnx修改

以yolov8n和coco数据集为例,导出的onnx输出为[1,8400,85]的形式。

 

在转为rknn模型的过程中,如果打开了量化,类别得分会全部被量化为0,所以单输出的onnx转成rknn会有问题(如果有大佬有解决办法,并且成功对单输出的模型进行后处理的,欢迎讨论分享)

于是,我把onnx修改为两个输出,即bbox和cls两个分支。

并且我将bbox分支的reshape后面的算子全部去除,因为在板子上运行的时候老是报错,由于这个地方在源码中就是将bbox的预测输出变成xyhw的输出格式,于是干脆砍掉,在后处理的过程中修改即可。为了方便后处理,在输出之前添加了tanspose算子,将输出由[1,4,8400]变为[1,8400,4],我的类别是4类,所以我的两个输出都为1*8400*4。

 

3. onnx转rknn

官方有提供onnx转rknn的工具,rknn-toolkits。根据README搭建rknn的环境即可。

下面是yolov8,onnx转rknn的python代码

  1. import os
  2. import urllib
  3. import traceback
  4. import sys
  5. import numpy as np
  6. from rknn.api import RKNN
  7. ONNX_MODEL = './yolov8n.onnx'
  8. RKNN_MODEL = './yolov8n_2out_transpose_int8_best.rknn'
  9. IMG_PATH = './bus.jpg'
  10. DATASET = './dataset.txt'
  11. # 是否打开量化
  12. QUANTIZE_ON = True
  13. # CLASSES = ("person", "bicycle", "car", "motorbike ", "aeroplane ", "bus ", "train", "truck ", "boat", "traffic light","fire hydrant", "stop sign ", "parking meter", "bench", "bird", "cat", "dog ", "horse ", "sheep", "cow", "elephant","bear", "zebra ", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite","baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup", "fork", "knife ","spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza ", "donut", "cake", "chair", "sofa","pottedplant", "bed", "diningtable", "toilet ", "tvmonitor", "laptop ", "mouse ", "remote ", "keyboard ", "cell phone", "microwave ","oven ", "toaster", "sink", "refrigerator ", "book", "clock", "vase", "scissors ", "teddy bear ", "hair drier", "toothbrush ")
  14. # 修改成自己的类即可
  15. CLASSES = ("person", "bicycle", "car","dog")
  16. if __name__ == '__main__':
  17. # Create RKNN object
  18. rknn = RKNN(verbose=True)
  19. # pre-process config
  20. print('--> Config model')
  21. rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]],
  22. target_platform="rk3588")
  23. print('done')
  24. # Load ONNX model
  25. print('--> Loading model')
  26. ret = rknn.load_onnx(model=ONNX_MODEL)
  27. if ret != 0:
  28. print('Load model failed!')
  29. exit(ret)
  30. print('done')
  31. # Build model
  32. print('--> Build model.')
  33. ret = rknn.build(do_quantization=QUANTIZE_ON, dataset=DATASET)
  34. if ret != 0:
  35. print('Build model failed!')
  36. exit(ret)
  37. print('done')
  38. """
  39. # Export RKNN model
  40. print('--> Export rknn model')
  41. ret = rknn.export_rknn(RKNN_MODEL)
  42. if ret != 0:
  43. print('Export rknn model failed!')
  44. exit(ret)
  45. print('done')

关于推理代码,下次再分享。

声明:本文内容由网友自发贡献,转载请注明出处:【wpsshop】
推荐阅读
相关标签
  

闽ICP备14008679号