pt转onnx转rknn（yolov5部署RK3566）

最新推荐文章于 2025-03-15 01:06:18 发布

原创最新推荐文章于 2025-03-15 01:06:18 发布 · 782 阅读

6 ·

CC 4.0 BY-SA版权

文章标签：

#YOLO

笔者使用的rknn_yolov5_android_apk_demo实现检测的

顺便提一句，使用scrcpy-win64-v2.5通过笔记本远程连接rk3566的板子挺方便的。

该项目需要的yolov5-5.0版本的代码进行训练和模型，最后得到的模型输出头如下图所示

可以看到输出头有三个哦

且激活函数从sigmoid换成了relu（比较推荐，rk板子跑出的检测的id置信度和帧率都会提高）

训练遇到的问题：1.pytorch版本不能太高，笔者使用1.7.0

2.numpy版本不能太高，笔者使用1.22.0

3.每次开新的训练，数据集train和val文件生成的缓存文件进行删除

训练完成后进行pt到onnx模型的转换

模型转换需要修改代码/models/yolo.py，将Detect下的foward函数修改成：（必须）

import os
os.environ['RKNN_model_hack'] = 'npu_2'


def forward(self, x):
        z = []  # inference output
        for i in range(self.nl):
            if os.getenv('RKNN_model_hack', '0') != '0':
                x[i] = torch.sigmoid(self.m[i](x[i]))  # conv

        return x

然后使用export.py得到onnx

 python models/export.py --weights ./64500.pt --img 640 --batch 1

接下来将得到的onnx转换成rknn

使用（GitHub - airockchip/rknn-toolkit2）rknn-toolkit2-1.5.2（笔者的版本）工具进行转换

先安装转换环境
这里我们先conda create -n rknn python=3.8创建环境

在终端激活环境，在终端输入pip install -r requirements_cp38-1.5.2.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
然后再输入pip install rknn_toolkit2-1.5.2+708089d1-cp38-cp38-linux_x86_64.whl

这两个文件注意找对路径
然后，我们的转rknn环境就配置完成了。

报错可以换换虚拟环境python的版本

开始转换模型

在rknn-toolkit2-1.5.2\examples\onnx\yolov5\test.py中

将这些修改成自己模型的相关参数

注意IMG_PATH和IMG_PATH一定换成自己模型的，不然会出现id只识别第一个的情况！！

因为这对于模型量化有很大的影响

注意process函数要与下面代码一样

def process(input, mask, anchors):

    anchors = [anchors[i] for i in mask]
    grid_h, grid_w = map(int, input.shape[0:2])

    box_confidence = input[..., 4]
    box_confidence = np.expand_dims(box_confidence, axis=-1)

    box_class_probs = input[..., 5:]

    box_xy = input[..., :2]*2 - 0.5

    col = np.tile(np.arange(0, grid_w), grid_w).reshape(-1, grid_w)
    row = np.tile(np.arange(0, grid_h).reshape(-1, 1), grid_h)
    col = col.reshape(grid_h, grid_w, 1, 1).repeat(3, axis=-2)
    row = row.reshape(grid_h, grid_w, 1, 1).repeat(3, axis=-2)
    grid = np.concatenate((col, row), axis=-1)
    box_xy += grid
    box_xy *= int(IMG_SIZE/grid_h)

    box_wh = pow(input[..., 2:4]*2, 2)
    box_wh = box_wh * anchors

    box = np.concatenate((box_xy, box_wh), axis=-1)

    return box, box_confidence, box_class_probs

至此你的rknn模型就转换好了

部署到你的安卓项目享受npu的运行吧！