onnx2rknn.py 1.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344
  1. from rknn.api import RKNN
  2. # ONNX model uses NCHW format with float32 input normalized to 0-1
  3. # mean_values=0, std_values=1 means the RKNN will NOT apply any normalization
  4. # during inference — the user must provide pre-normalized (0-1) float32 NHWC input
  5. #
  6. # IMPORTANT: RKNN inference always expects NHWC input (1,H,W,C) regardless of
  7. # the ONNX model's NCHW layout. The toolkit handles the transpose internally.
  8. rknn = RKNN(verbose=True)
  9. rknn.config(
  10. target_platform='rk3588',
  11. mean_values=[[0, 0, 0]],
  12. std_values=[[1, 1, 1]],
  13. quantized_dtype='w8a8',
  14. optimization_level=3
  15. )
  16. print('Loading ONNX model...')
  17. ret = rknn.load_onnx(model='yolo11m_safety.onnx', input_size_list=[[3, 640, 640]])
  18. if ret != 0:
  19. print("load_onnx failed")
  20. exit(1)
  21. print('Building RKNN model (do_quantization=False, float32)...')
  22. ret = rknn.build(dataset='dataset.txt', do_quantization=False)
  23. if ret != 0:
  24. print("build failed")
  25. exit(1)
  26. print('Exporting RKNN model...')
  27. rknn.export_rknn('./yolo11m_safety.rknn')
  28. # Optional: hybrid quantization (uncomment to enable int8 weights with float outputs)
  29. # This can improve NPU performance while keeping output precision
  30. # ret = rknn.build(dataset='dataset.txt', do_quantization=True, quantized_dtype='asymmetric_quantized-u8')
  31. # if ret != 0:
  32. # print("build with quantization failed")
  33. # exit(1)
  34. # rknn.export_rknn('./yolo11m_safety_quantized.rknn')
  35. rknn.release()
  36. print("Done: yolo11m_safety.rknn")