YOLOv8部署在RKNN3588上

1. onnx导出

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

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

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

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代码

import os
import urllib
import traceback
import sys
import numpy as np

from rknn.api import RKNN

ONNX_MODEL = './yolov8n.onnx'
RKNN_MODEL = './yolov8n_2out_transpose_int8_best.rknn'
IMG_PATH = './bus.jpg'
DATASET = './dataset.txt'

# 是否打开量化
QUANTIZE_ON = True


# 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 ")
# 修改成自己的类即可
CLASSES = ("person", "bicycle", "car","dog")

if __name__ == '__main__':

    # Create RKNN object
    rknn = RKNN(verbose=True)

    # pre-process config
    print('--> Config model')
    rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]],             
                target_platform="rk3588")
    print('done')

    # Load ONNX model
    print('--> Loading model')
    ret = rknn.load_onnx(model=ONNX_MODEL)
    if ret != 0:
        print('Load model failed!')
        exit(ret)
    print('done')

    # Build model
    print('--> Build model.')
    ret = rknn.build(do_quantization=QUANTIZE_ON, dataset=DATASET)
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('done')
    """

    # Export RKNN model
    print('--> Export rknn model')
    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print('Export rknn model failed!')
        exit(ret)
    print('done')

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