ITX-3588J开发（五）

YOLOv8部署

准备工作

板端

见我的ITX-3588J开发（四），升级开发板toolkit及驱动，本文默认使用rknn-toolkit2-lite的2.0版本。

PC端

我记得好像PC端的工具包不支持windows系统，建议使用Ubuntu系统的电脑或服务器，虚拟机应该也行。
网址：https://github.com/airockchip/rknn-toolkit2/tree/master
切换Tags到2.0版本，并到此处

下载与自己环境python版本对应的wheel和requirements.txt文件。以python版本3.9为例。
首先构建虚拟环境并激活。

conda create -n rknn python=3.9
conda activate rknn

安装依赖包。个人建议把txt里的torch有关的那一行改为torch==1.11.0

pip install -r requirements_cp39-2.0.0b0.txt

然后

pip install rknn_toolkit2-2.0.0b0+9bab5682-cp39-cp39-linux_x86_64.whl

PC端准备完成。

权重文件

常规需求直接前往rknn_model_zoo下载对应的onnx文件即可。
接下来讲解如何使用个人数据集训练yolov8并得到onnx文件。
首先必须使用airockchip官方指定的版本进行训练。具体的训练流程和常规的yolov8一样，环境要求也相同，直接使用之前的环境应该没问题。参考YOLOV8导出onnx模型，再导出RKNN模型，对yolov8的代码进行一定的修改后可以提高推理效率。
修改位置：
在conv.py文件中将Conv类的SiLU替换为ReLU。

class Conv(nn.Module):
    # Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation)
    #default_act = nn.SiLU()  # 默认激活函数
    default_act = nn.ReLU()   # 改动，silu在rknn上不能和conv融合
    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
        super().__init__()
        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)
        self.bn = nn.BatchNorm2d(c2)
        self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()

在head.py文件中在Detect类中添加onnx导出代码。

class Detect(nn.Module):
    """YOLOv8 Detect head for detection models."""
    dynamic = False  # force grid reconstruction
    export = False  # export mode
    shape = None
    anchors = torch.empty(0)  # init
    strides = torch.empty(0)  # init

    ###  添加适用于RKNN导出onnx的部分代码
    conv1x1 = nn.Conv2d(16, 1, 1, bias=False).requires_grad_(False)
    aaa = torch.arange(16, dtype=torch.float)
    conv1x1.weight.data[:] = nn.Parameter(aaa.view(1, 16, 1, 1))
    ###
    
    def __init__(self, nc=80, ch=()):  # detection layer
        super().__init__()
        self.nc = nc  # number of classes
        self.nl = len(ch)  # number of detection layers
        self.reg_max = 16  # DFL channels (ch[0] // 16 to scale 4/8/12/16/20 for n/s/m/l/x)
        self.no = nc + self.reg_max * 4  # number of outputs per anchor
        self.stride = torch.zeros(self.nl)  # strides computed during build 
        c2, c3 = max((16, ch[0] // 4, self.reg_max * 4)), max(ch[0], min(self.nc, 100))  # channels
        self.cv2 = nn.ModuleList(
            nn.Sequential(Conv(x, c2, 3), Conv(c2, c2, 3), nn.Conv2d(c2, 4 * self.reg_max, 1)) for x in ch)
        self.cv3 = nn.ModuleList(nn.Sequential(Conv(x, c3, 3), Conv(c3, c3, 3), nn.Conv2d(c3, self.nc, 1)) for x in ch)
        self.dfl = DFL(self.reg_max) if self.reg_max > 1 else nn.Identity()
        # dfl:  DFL（Distribution Focal Loss）

    def forward(self, x):
        """Concatenates and returns predicted bounding boxes and class probabilities."""
        shape = x[0].shape  # BCHW
         ###  添加适用于RKNN导出onnx的部分代码
        if self.export and self.format in ('onnx'):
            y = []
            for i in range(self.nl):
                t1 = self.cv2[i](x[i])
                t2 = self.cv3[i](x[i])
                y.append(self.conv1x1(t1.view(t1.shape[0], 4, 16, -1).transpose(2, 1).softmax(1)))
                y.append(t2)
            return y
        ###
        for i in range(self.nl):
            x[i] = torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)
        if self.training:
            return x

新建一个pt2onnx.py文件。记得按照自己需求修改model_path和imgsz，opset必须为12，不然之后转rknn大概率出问题。

from ultralytics import YOLO
model_path = "runs/detect/train/weights/best.pt" 
model = YOLO(model_path)
success = model.export(format='onnx', imgsz=(640, 640), opset=12, simplify=True)  # export the model to onnx format
assert success
print("转换成功")

所有准备工作完成。

onnx转rknn

官方的rknn_model_zoo提供了蛮多转换代码的。找到yolov8的转换代码。运行就行。