51 lines
1.4 KiB
Python
51 lines
1.4 KiB
Python
from rknn.api import RKNN
|
|
import os
|
|
|
|
ONNX_MODEL = 'runs/train/YNDC_20251010_yolov5_optimize/weights/best.onnx'
|
|
RKNN_MODEL = 'runs/train/YNDC_20251010_yolov5_optimize/weights/best_rk3576.rknn'
|
|
DATASET_TXT = '/data6/zhangw/yolov5/data/images/lianghua.txt'
|
|
|
|
def main():
|
|
print("==== 1. 创建 RKNN 对象 ====")
|
|
rknn = RKNN(verbose=True)
|
|
|
|
print("==== 2. RKNN 配置 ====")
|
|
rknn.config(
|
|
target_platform='rk3576',
|
|
mean_values=[[0, 0, 0]],
|
|
std_values=[[255, 255, 255]],
|
|
quantized_dtype='w8a8',
|
|
optimization_level=3
|
|
)
|
|
|
|
print("==== 3. 写入量化校准文件 ====")
|
|
os.makedirs(os.path.dirname(DATASET_TXT), exist_ok=True)
|
|
with open(DATASET_TXT, 'w') as f:
|
|
f.write('/data6/zhangw/yolov5/data/images/bus.jpg\n')
|
|
f.write('/data6/zhangw/yolov5/data/images/zidane.jpg\n')
|
|
|
|
print("==== 4. 加载 ONNX ====")
|
|
ret = rknn.load_onnx(model=ONNX_MODEL)
|
|
if ret != 0:
|
|
print("❌ 加载 ONNX 失败")
|
|
exit(ret)
|
|
|
|
print("==== 5. 构建 RKNN（量化） ====")
|
|
ret = rknn.build(do_quantization=True, dataset=DATASET_TXT)
|
|
if ret != 0:
|
|
print("❌ 构建 RKNN 失败")
|
|
exit(ret)
|
|
|
|
print("==== 6. 导出 RKNN ====")
|
|
ret = rknn.export_rknn(RKNN_MODEL)
|
|
if ret != 0:
|
|
print("❌ 导出失败")
|
|
exit(ret)
|
|
|
|
print(f"🎉 RKNN 模型导出成功: {RKNN_MODEL}")
|
|
|
|
rknn.release()
|
|
|
|
if __name__ == '__main__':
|
|
main()
|