当前位置:   article > 正文

RKNN3588——YOLOv8的PT模型转RKNN模型_yolov8 转 rknn

yolov8 转 rknn

一:PT转ONNX

1. 首先克隆rknn修改后的ultralytics版本项目到本地

https://github.com/airockchip/ultralytics_yolov8

  1. cd ultralytics-main
  2. pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
  3. pip install -e .

主要是修改了源码的ultralytics/nn/modules/head.py和ultralytics/engine/exporter.py两个文件。

2. 使用修改后的ultralytics对pt模型进行模型转换,此处的format=rknn代表支持rknn后续的转换,而不是用onnx,一定要注意!!!

yolo export model=/your_path/best.pt format=rknn

二:ONNX转RKNN

1. 在linux x86上配置Toolkit环境,具体步骤如下链接:

https://docs.radxa.com/rock5/rock5b/app-development/rknn_install

相关链接: 

# 下载 RKNN-Toolkit2 仓库
https://github.com/airockchip/rknn-toolkit2.git

# 下载 RKNN Model Zoo 仓库
https://github.com/airockchip/rknn_model_zoo.git

2. 模型的转换只能在rknn-toolkit2中,因此Linux系统只能是linux x86,配置环境时注意python版本问题

- RKNN-Toolkit2 is not compatible with [RKNN-Toolkit](https://github.com/airockchip/rknn-toolkit)
- Currently only support on:
  - Ubuntu 18.04 python 3.6/3.7
  - Ubuntu 20.04 python 3.8/3.9
  - Ubuntu 22.04 python 3.10/3.11

配置好环境后,可通过

  1. python
  2. import rknn

检查是否安装成功?!

3. 然后通过rknn_model_zoo里的YOLOv8的convert.py转换即可

  1. import sys
  2. from rknn.api import RKNN
  3. if __name__ == '__main__':
  4. # onnx模型位置
  5. model_path = "/home/zhangh/RKNN_Docker/ultralytics_yolov8-main/ultralytics_yolov8-main/yolov8n.onnx"
  6. # rknn导出位置
  7. output_path = "/home/zhangh/RKNN_Docker/ultralytics_yolov8-main/ultralytics_yolov8-main/weight/yolov8n.rknn"
  8. # Create RKNN object
  9. rknn = RKNN(verbose=False)
  10. # Pre-process config
  11. print('--> Config model')
  12. rknn.config(mean_values=[[0, 0, 0]], std_values=[
  13. [255, 255, 255]], target_platform='rk3588')
  14. print('done')
  15. # Load model
  16. print('--> Loading model')
  17. ret = rknn.load_onnx(model=model_path)
  18. if ret != 0:
  19. print('Load model failed!')
  20. exit(ret)
  21. print('done')
  22. # Build model
  23. print('--> Building model')
  24. ret = rknn.build(do_quantization=False)
  25. if ret != 0:
  26. print('Build model failed!')
  27. exit(ret)
  28. print('done')
  29. # Export rknn model
  30. print('--> Export rknn model')
  31. ret = rknn.export_rknn(output_path)
  32. if ret != 0:
  33. print('Export rknn model failed!')
  34. exit(ret)
  35. print('done')
  36. # Release
  37. rknn.release()

可视化rknn模型 

声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/人工智能uu/article/detail/933193
推荐阅读
相关标签
  

闽ICP备14008679号