RKNN：yolov8模型转换与板端推理流程

近期，在研究瑞芯微的RKNN模型推理时，遇到一些坑，现记录下来，以备忘，亦供同道者参考。

目录

1. 模型转换

1.1. 宿主机环境配置

1.2. onnx模型准备

1.3. onnx转rknn

2. 模型推理

2.1. 推理环境配置

2.2. 推理验证

2.2.1. 输入格式报错

 2.2.2. 需要使用torch

本文主要内容包括两块，一是在宿主机进行onnx-->rknn的模型转换，二是在SOC上利用转换好的rknn模型进行推理。

1. 模型转换

这一步需要在x86宿主机上进行，并配置相关环境。

1.1. 宿主机环境配置

主要就是在python环境中安装RKNN-Toolkit2。

构造环境这一步，建议在docker里面进行，这样可以避免和宿主机环境产生冲突，或者因为宿主机环境和需要的环境不匹配而导致的模型转换出现问题。

PS：之前博主在宿主机的python环境中直接pip安装的rknn-toolkit2，然而用一样的代码进行模型转换，最后推理结果总是出现错乱或者直接所有结果置信度特别低（接近0），最后放弃挣扎，转用docker才没有问题。

这里使用的版本是1.6.0，下载地址：

https://codeload.github.com/rockchip-linux/rknn-toolkit2/zip/refs/tags/v1.6.0https://codeload.github.com/rockchip-linux/rknn-toolkit2/zip/refs/tags/v1.6.0下载解压之后的目录如下所示：

.
├── CHANGELOG.md
├── doc
├── LICENSE
├── README.md
├── res
├── rknn-toolkit2
├── rknn_toolkit_lite2
└── rknpu2

进到“rknn-toolkit2/docker/docker_file/ubuntu_20_04_cp38”目录下，并利用提供的Dockerfile构建一个docker镜像：

docker build -t rknn_toolkit2_1.6.0 . -f Dockerfile_ubuntu_20_04_for_cp38

等待镜像构建完毕，就会在里面自动安装模型转换所需的所有软件环境。然后，就可以启动一个容器，进到docker系统中：

# 注意：我们上一步构造的镜像命名为rknn_toolkit2_1.6.0，需要改为你的镜像名字；
#      同时，这里作了一个主机和容器之间的路径映射，以便于文件共享，需要将“ /data/projects/other/rknn_test/share_dir”改为你自己的路径
docker run -v /data/projects/other/rknn_test/share_dir:/workspace/ -it rknn_toolkit2_1.6.0

至此，就构建好了所需的环境。

1.2. onnx模型准备

这一步是为了从我们训练的torch模型转为onnx中间格式的模型。这里不能直接使用ultralytics官方的库，需要作一定的改动，使其最后的detect部分改为rknn支持的方式:

 可以参考rknn修改后的ultralytics版本：

GitHub - airockchip/ultralytics_yolov8: NEW - YOLOv8 🚀 in PyTorch > ONNX > CoreML > TFLiteNEW - YOLOv8 🚀 in PyTorch > ONNX > CoreML > TFLite - airockchip/ultralytics_yolov8https://github.com/airockchip/ultralytics_yolov8也可以自己在ultralytics的源码上修改，基本就是改ultralytics/nn/modules/head.py和ultralytics/engine/exporter.py两个文件，这里不细说了，有空了可以单独聊聊这块内容。

使用修改后的ultralytics进行模型转换：

yolo export model=yolov8n.pt format=rknn

即可得到一个支持转rknn的onnx模型。

1.3. onnx转rknn

这一步需要利用rknn_model_zoo里面的代码，当然自己写一个转换的脚本也很简单。这里以rknn_model_zoo为例说明。先把该代码拉下来，地址在：

GitHub - airockchip/rknn_model_zooContribute to airockchip/rknn_model_zoo development by creating an account on GitHub.https://github.com/airockchip/rknn_model_zoo.git进入目录：rknn_model_zoo/examples/yolov8/python，然后使用模型转换脚本将onnx模型转为所需的rknn格式的模型： 

python convert.py yolov8n.onnx rk3588

OK，至此，我们在宿主机的工作就完毕了。接下来，需要到板子上进行操作。

2. 模型推理

2.1. 推理环境配置

首先配置一个python虚拟环境：

# 根据你使用python版本进行创建：
conda create -n py38 python=3.8

然后，在1.1中我们下载的文件里面有rknn-toolkit2-1.6.0/rknn_toolkit_lite2/packages这个文件夹，从里面找到对应的安装包，比如我这里Python用的3.8所以就选择rknn_toolkit_lite2-1.6.0-cp38-cp38-linux_aarch64.whl；

最后，将其拷贝到RK板子上，并在板子上打开命令行，进入Python虚拟环境，安装rknn_toolkit_lite2:

pip install rknn_toolkit_lite2-1.6.0-cp38-cp38-linux_aarch64.whl

2.2. 推理验证

将1.3中拉下来的demo代码考到板子上，进到rknn_model_zoo/examples/yolov8/python目录下，使用python脚本进行推理：

# 注意，../model/yolov8.rknn要改为你转换的rknn模型路径
python yolov8.py --model_path ../model/yolov8.rknn --img_show --img_save --target rk3588

这一步如果没问题的话，就会展示推理结果，不过可能会出现以下问题：

2.2.1. 输入格式报错

Traceback (most recent call last):
  File "yolov8.py", line 262, in <module>
    outputs = model.run([input_data])
  File "/home/linaro/pkgs/rknn_model_zoo/py_utils/rknn_executor.py", line 24, in run
    print(inputs.shape)
AttributeError: 'list' object has no attribute 'shape'

对于这个错误，可以将这行代码：input_data = img

修改为：

input_data = np.expand_dims(img, axis=3)

将这行代码：outputs = model.run([input_data]) 

修改为：

outputs = model.run(input_data)

 2.2.2. 需要使用torch

在脚本实现的dfl函数中，使用了torch库，我们在板子上肯定不可能再去装一个torch的。因此，这里需要改为numpy实现，将以下代码替换原来的dfl函数即可：

def dfl(postion):
    n, c, h, w = postion.shape
    p_num = 4
    mc = c // p_num
    y = postion.reshape(n, p_num, mc, h, w)
    y = softmax(y, 2)
    acc_metrix = np.arange(mc).reshape(1, 1, mc, 1, 1)
    y = (y * acc_metrix).sum(2)
    return y
    
def softmax(data, dim):
    max = np.max(data, axis=dim, keepdims=True).repeat(data.shape[dim], axis=dim)
    exps = np.exp(data - max)
    return exps / np.sum(exps, axis=dim, keepdims=True)

最后，没有其他意外的话，就可以得到推理结果了：