yolov5的pt模型转化为rk3588的rknn，并在rk3588上调用api进行前向推理

诶尔法Alpha

已于 2024-11-25 16:52:35 修改

阅读量1.7k

点赞数 25

CC 4.0 BY-SA版权

分类专栏： yolo项目文章标签： YOLO rknn 神经网络深度学习目标检测

于 2024-11-25 16:37:02 首次发布

本文链接：https://blog.youkuaiyun.com/weixin_45354497/article/details/144021373

当使用yolov5进行目标检测且进行边缘计算的场景时，要考虑性价比或者国产化的话，rk3588板子是个不错的选择。

本篇介绍yolov5的pytorch模型转化为rknn的流程，并展示在rk板子上如何调用相关api来使用转好的rknn模型进行前向推理。

pt转rknn流程

pt转onnx

首先将训练好的pt模型转为onnx中间模型，在转之前需要先修改主目录底下的models下的yolo.py的部分代码，将如图的forward推理部分进行注释。
在这里插入图片描述
替换为：

    def forward(self, x):
        z = []
        for i in range(self.nl):
            x[i] = self.m[i](x[i])
 
        return x

如果识别出现乱框的现象，则替换为：

    def forward(self, x):
        z = []
        for i in range(self.nl):
            x[i] = torch.sigmoid(self.m[i](x[i]))
 
        return x

记得在训练时再给他改回去，否则会报错，在转模型时才需要改这部分代码。
接下来运行export.py来转onnx：

python export.py --weights runs/train/exp/weights/best.pt --img 320 --batch 1 --include onnx

onnx转rknn

首先创建一个虚拟环境，我创建的python版本为3.8。可以直接使用pip install -r requirements.txt来创建环境；
如果出现报错，则使用conda env create -f environment.yml来创建。如果pip安装可以的话，可以直接-i换源比较方便。
这里rknn-toolkit2安装会失败是正常的，下面会手动进行安装。

*以上为直接复制我的环境，也可以根据下面的步骤来自己安装相应的包。

接下来下载RKNN-Toolkit2，可以通过官网下载，也可以通过我上传的资源下载。
下载完先进入刚才创建的环境安装rknn-toolkit2包，如图，进入packages底下，安装相关依赖：

pip install -r requirements_cp38-1.6.0.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

在这里插入图片描述
然后安装相关rknn-toolkit2包：

pip install rknn_toolkit2-1.6.0+81f21f4d-cp38-cp38-linux_x86_64.whl

*接下来进行onnx转rknn
在如图所示路径下打开test.py
在这里插入图片描述
修改相关参数

最后一个参数要改为rk3588，默认为rk3566。改完后直接运行test.py文件即可转为rknn。

rk板子上进行前向推理

因为rk3588是aarch64架构的，所以不能用rknn-toolkit2包，而是要用rknn-toolkit-lite2包，在rk上安装对应的whl：
在这里插入图片描述
代码如下：

from copy import copy
import time
import numpy as np
import cv2
from rknnlite.api import RKNNLite

RKNN_MODEL = 'best.rknn'
# IMG_PATH = './input/240513_00000741.jpg'
OBJ_THRESH = 0.25
NMS_THRESH = 0.45
IMG_SIZE = (640, 640)
OUTPUT_VIDEO_PATH = 'output_1.mp4'
BOX = (450, 150, 1100, 550)
CLASSES = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', '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', 'couch',
           'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
           'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
           'hair drier', 'toothbrush']

anchors = [[[10, 13], [16, 30], [33, 23]],
           [[30, 61], [62, 45], [59, 119]],
           [[116, 90], [156, 198], [373, 326]]]

class Letter_Box_Info():
    def __init__(self, shape, new_shape, w_ratio, h_ratio, dw, dh, pad_color) -