RK3588部署rknntoolkit2进行模型转换

  • Post author:
  • Post category:其他




任务目标

  1. 将训练好的yolov7.pt转换为yolov7.onnx。
  2. 将yolov7.onnx转换为yolov7.rknn。



实验过程

说明:本次实验仅为模型转换步骤的验证,所以不涉及模型训练部分。

  1. 首先去官网下载源码和yolov7.pt文件。

源码地址:https://github.com/WongKinYiu/yolov7.git

模型地址:https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7.pt



  1. models/yolo.py

    中将Detect类中的前向传播过程由:
    def forward(self, x):
        # x = x.copy()  # for profiling
        z = []  # inference output
        self.training |= self.export
        for i in range(self.nl):
            x[i] = self.m[i](x[i])  # conv
            bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
            x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()

            if not self.training:  # inference
                if self.grid[i].shape[2:4] != x[i].shape[2:4]:
                    self.grid[i] = self._make_grid(nx, ny).to(x[i].device)
                y = x[i].sigmoid()
                if not torch.onnx.is_in_onnx_export():
                    y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy
                    y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
                else:
                    xy, wh, conf = y.split((2, 2, self.nc + 1), 4)  # y.tensor_split((2, 4, 5), 4)  # torch 1.8.0
                    xy = xy * (2. * self.stride[i]) + (self.stride[i] * (self.grid[i] - 0.5))  # new xy
                    wh = wh ** 2 * (4 * self.anchor_grid[i].data)  # new wh
                    y = torch.cat((xy, wh, conf), 4)
                z.append(y.view(bs, -1, self.no))

        if self.training:
            out = x
        elif self.end2end:
            out = torch.cat(z, 1)
        elif self.include_nms:
            z = self.convert(z)
            out = (z, )
        elif self.concat:
            out = torch.cat(z, 1)
        else:
            out = (torch.cat(z, 1), x)

        return out


改为:

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

        return x


注:这一步一定要改,否则在onnx转rknn时会报错



修改后将yolov7.pt移动至export.py同一文件夹下,运行export.py便可以得到yolov7.onnx。

3. 去官网下载rknntoolkit2,并在

Linux

上解压,官网地址如下:

https://github.com/rockchip-linux/rknn-toolkit2

使用Miniconda创建rknn环境:

conda create -n rknn python=3.8
conda activate rknn

环境安装:

#安装python工具,ubuntu20.04默认是安装了python3.8.10
sudo apt update
sudo apt-get install python3-dev python3-pip python3.8-venv gcc

#安装相关库和软件包
sudo apt-get install libxslt1-dev zlib1g-dev libglib2.0 libsm6 \
libgl1-mesa-glx libprotobuf-dev gcc

安装rknntoolkit2:

#安装依赖库,根据rknn-toolkit2\doc\requirements_cp38-1.4.0.txt
pip3 install numpy
pip3 install -r doc/requirements_cp38-1.4.0.txt

#安装rknn_toolkit2
pip3 install packages/rknn_toolkit2-1.4.0_22dcfef4-cp38-cp38-linux_x86_64.whl

检测是否安装成功:

(.toolkit2_env) llh@-:~/project-Toolkit2$ python
Python 3.8.10 (default, Jun 22 2022, 20:18:18)
[GCC 9.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> from rknn.api import RKNN
>>>
#没有报错就安装成功了
  1. 在rknn-toolkit2工程文件夹中浏览至./examples/onnx/yolov5,将我们转换得到的yolov7.onnx复制到该文件夹下,修改该文件夹下的test.py中的内容为自己模型的名字,要修改的地方如下:

    在这里插入图片描述
    在这里插入图片描述

    此处我们只需要修改前两项和搭载平台的名字即可。

    最后运行test.py,即可得到rknn模型。



版权声明:本文为xingman510原创文章,遵循 CC 4.0 BY-SA 版权协议,转载请附上原文出处链接和本声明。