采用rknn-toolkit导出rknn模型并部署在rock3a-rk3568芯片上全流程

news2024/11/26 16:25:07

因工作需要，需要将目标检测模型部署在开发板上。在走了很多弯路后找到一个成功的案例并记载下来

这里说一下我现有的硬件设备。

我是购买的RADXA的rock3a开发板搭载的soc是rk3568

这是开发板的正面图，因为瑞芯微针对计算机视觉中的目标检测模型有一套自己的前向推理框架，所以我就着眼于搭载rockchip的开发板rock3a

目标检测模型这里采用的是yolo模型由于原生yolov5模型里面有一些算子可能在模型转换时不支持，这里采用瑞芯微官方推荐的yolov5，链接在下面：

GitHub - airockchip/yolov5: YOLOv5 in PyTorch > ONNX > CoreML > TFLite

下载好源代码后，需要修改一个地方：

修改的代码我这里贴出来

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

            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.onnx_dynamic or self.grid[i].shape[2:4] != x[i].shape[2:4]:
                    self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)

                y = x[i].sigmoid()
                if self.inplace:
                    y[..., 0:2] = (y[..., 0:2] * 2 + self.grid[i]) * self.stride[i]  # xy
                    y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
                else:  # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
                    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.grid[i]) * self.stride[i]  # xy
                    wh = (wh * 2) ** 2 * self.anchor_grid[i]  # wh
                    y = torch.cat((xy, wh, conf), 4)
                z.append(y.view(bs, -1, self.no))

        if os.getenv('RKNN_model_hack', '0') != '0':
            return z

        return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(z, 1), x)
        '''
        return x[0],x[1],x[2]

然后根据自己的数据集训练一版模型，训练过程这里不赘述。

这里将训练好的模型yolov5s.pt 转化为onnx模型

转化命令为：

python3 export.py --weights yolov5s.pt --img 640 --batch 1 --opset 12 --include onnx

转化完后需要用Netron软件打开模型的可视化界面

记下图中的三个框里面的节点名称，这三个节点是网络的输出节点

新建一个 python文件 onnx2rknn.py 该文件是仿制test.py文件编写的

文件内容为

from rknn.api import RKNN

ONNX_MODEL = 'yolov5s.onnx'
platform = "rk3568"
RKNN_MODEL = 'yolov5s_{}_out_opt.rknn'.format(platform)

if __name__ == '__main__':

    add_perm = False # 如果设置成True,则将模型输入layout修改成NHWC
    # 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=platform,
                )
    print('done')

    # Load tensorflow model
    print('--> Loading model')
    ret = rknn.load_onnx(model=ONNX_MODEL,outputs=['output', '327', '328'])# 这里一定要根据onnx模型修改
    if ret != 0:
        print('Load onnx model failed!')
        exit(ret)
    print('done')

    # Build model
    print('--> Building model')
    ret = rknn.build(do_quantization=True, dataset='./dataset.txt')
    if ret != 0:
        print('Build rkmodel failed!')
        exit(ret)
    print('done')

    # rknn.export_rknn_precompile_model(RKNN_MODEL)
    rknn.export_rknn(RKNN_MODEL)

    rknn.release()

并将该文件放置于rknn-toolkit工作目录下

我的目录是/home/rock/workspace/rknn-toolkit2-1.4.0/examples/onnx/yolov5

也就是example目录下的onnx专栏的yolov5子目录

注意需要在onnx2rknn.py文件中注明模型运行的目标平台我的目标平台是rk3568

另外需要将上一步 netron 查看的输出部分的三个节点在onnx2rknn.py文件中注明

也就是如下截图所示