pt-＞onnx-＞rknn(量化) step by step & FAQ

原创

已于 2025-02-14 17:34:30 修改 · 1.7k 阅读

21 ·

CC 4.0 BY-SA版权

文章标签：

#python #深度学习 #YOLO #ai

于 2025-02-14 11:29:07 首次发布

文档修订中...

1.pt->onnx

这个转换是在yolov11的docker环境做的转换。非常简单。

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# 获取当前脚本文件所在目录的父目录，并构建相对路径
import os
import sys
current_dir = os.path.dirname(os.path.abspath(__file__))
project_path = os.path.join(current_dir, '..')
sys.path.append(project_path)
sys.path.append(current_dir)
#based: https://docs.ultralytics.com/modes/export/#key-features-of-export-mode
from ultralytics import YOLO

# Load a model
#model = YOLO("yolo11n.pt")  # load an official model
model = YOLO(r"./best.pt")  # load a custom trained model
# Export the model
model.export(format="onnx")

2.onnx->rknn

这个转换需要在rknn-toolkit2所在docker环境转换。这个过程中，注意事项有几个：

需要引入一个校验图片，一个或者一组。
图片的索引文件是：dataset.txt，内容是一个或者多个图片文件。
命令行参数中可以指定最终的量化级别：

python onnx2rknn.py ./yolo11n.onnx rk3588 i8

#onnx2rknn.py
import cv2
import numpy as np

from rknn.api import RKNN
import os

if __name__ == '__main__':

    platform = 'rk3588'
    exp = 'yolov11'
    Width = 640
    Height = 640
    # Model from https://github.com/airockchip/rknn_model_zoo
    MODEL_PATH = './best.onnx' 
    NEED_BUILD_MODEL = True
    # NEED_BUILD_MODEL = False
    im_file = './frame_0127.png'

    # Create RKNN object
    rknn = RKNN()

    OUT_DIR = "rknn_models"
    RKNN_MODEL_PATH = './{}/{}_{}.rknn'.format(
        OUT_DIR, exp+'-'+str(Width)+'-'+str(Height), platform)
    if NEED_BUILD_MODEL:
        DATASET = './dataset.txt'
        rknn.config(mean_values=[[0, 0, 0]], std_values=[
                    [255, 255,