系列文章目录

RK3566 RK3568 RK3588 使用RKNN NPU进行推理
系列（一），开发环境的搭建

RK3566 RK3568 RK3588 使用RKNN NPU进行推理
系列（二），板端推理，板上部署

RK3566 RK3568 RK3588 使用RKNN NPU进行推理
系列（三），训练自己的数据集并转成rknn

RK3566 RK3568 RK3588 使用RKNN NPU进行推理
系列（四），RKNN 推理速度和效率

RK3566 RK3568 RK3588 使用RKNN NPU进行推理
系列（五），RKNN NPU推理速度和CPU 、内存、NPU使用情况

前言

本次我们训练自己的数据而不是用yolo公版demo的数据，把自己的数据集训练成pt转成onnx后再通过rknn-toolkit2转成rknpu 可以识别的RKNN 模型文件。

---

一：训练环境搭建

1.1 Win11/Win10 +pytorch+conda，电脑配置了一张英伟达RTX4060的显卡。

1.2 训练时间和电脑配置相关，训练的数据集越大，img size越大，训练所需时间越久，运用显卡来训练比CPU快很多。

1.3本次只是搭建利于在rknpu 上跑的ONNX环境，其他平台也类似。

二：下面我列出我电脑已经安装好的环境，环境配置方法请参考其他链接或者百度搜 索下，基本上大差不差.
注意：以下能打开的前提是conda等环境已经安装好了，如没有安装请先安装。

2.1 打开Anaconda PowerShell Prompt
Conda 环境，win64位，23.11.3版本，

2.2 英伟达显卡
输入nvidia-smi 查看英伟达显卡信息,
查看版本为12.7

2.3 torch版本
import torch
print(torch.version)
2.6.0+cu118 带cuda显卡的版本

2.4 python 版本
python -V
Python 3.9.0

2.5 labelImg 已经安装好

三：我们从官网下载Yolov5 版本
注意Yolov5 也有好几个版本，其中Yolov5 6.0及以上和6.0以下修改的方式不同
否则训练出来的onnx在转成rknn报错，或者在rknpu 跑不通。
这里最好下载6.0以上的版本。

数据集整理
3.1我们为了区别和yolo5 80种类不一样，我们在这里只标注1个类别，而且图片是重新采样过的。
3.2这里我们已经用lableImg工具标注好，如下，标注了200多张图片

3.3 把标注好的数据集放入,其中一个是图片另外一个是标签，图片和标签一一对应
训练用的，放入如下路径
Yolov5\data\images\train\images
Yolov5\data\images\train\labels
评估用的，放入如下路径
Yolov5\data\images\val\images
Yolov5\data\images\val\labels

四:接着修改代码：
我们可以按720P或者640分辨率进行训练，根据实际情况修改，分辨率设置越小推理速度越快
4.1修改data\VOC.yaml
修改我们刚刚替换的路径
path: ./data
train: # train images (relative to ‘path’) 16551 images

• images/train/images

val: # val images (relative to ‘path’) 4952 images

• images/val/images
  4.2 拷贝train.py Mytrain.py并根据情况修改相关指令，这里不用修改代码

根据电脑配置不同有可能报错的，默认为8改小点
–workers’, type=int, default=4,

4.3 拷贝train.py Mytrain.py并修改相关指令
我这里修改epochs为50次，为了节省时间
4.4 python ./Mytrain.py 输入指令开始训练
我用了198张图片进行训练
che’ images and labels… 198 found, 1 missing, 0 empty, 0 corrupt: 100%|██

50 epochs completed in 0.170 hours.
//导出路径，此时路径下生成pt
Results saved to runs\train\exp18

五：pt导出onnx
5.1导出的时候需要修改yolo.py，训练的时候无需修改，
否则生成的onnx无法转成正确的rknn。

打开yolo.py  class Detect(nn.Module):
def forward(self, x):
   z = []  # inference output
        #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.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))

        # return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(z, 1), x)
//注释掉，这个修改成
   z.append(torch.sigmoid(self.m[i](x[i])))
   return z

5.2 执行如下命令导出 onnx
python export.py --weights BlackWhite_best.pt --include onnx
此时在 runs\train\exp18 会生成 BlackWhite_best.onnx

六：把ONNX转成RKNN
需要再ubuntu中转，我们需要把BlackWhite_best.onnx放到如下路径
rknn-toolkit2\examples\onnx\yolov5
6.1修改test.py
ONNX_MODEL = ‘BlackWhite_best.onnx’
RKNN_MODEL = ‘BlackWhite_best.rknn’
IMG_PATH = ‘./420.jpg’
DATASET = ‘./mydataset.txt’

6.2执行python ./test.py
通过export_rknn 把onnx转成rknn并进行推理，
推理函数rknn.inference，
boxes, classes, scores = yolov5_post_process(input_data)进行后处理

生成的结果保存在本地中,img_1 = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

这里用RKNN+Opencv结合方式进行处理。