Yolo v5训练并移植到RK3588S平台

一、引言

        RK3588S支持NPU,提供高达6.0Tops的算力,可以用于部署深度学习项目。本篇文章介绍Yolo v5代码开发、模型转化(RK3588S只支持rknn模型文件)、部署。

        使用的RKNN-TooKit2,具体的环境搭建,请参考博文:RK3588(自带NPU)的环境搭建和体验(一)

本次使用Pytorch框架开发Yolo v5,在Window 10训练模型,在ubuntu下面转化模型。

二、Yolo v5

1、算法结构图

Yolo v5训练并移植到RK3588S平台_第1张图片图1 Yolo V5结构图

Yolo V5算法大致可以分为三个部分Backbone,FPN以及Yolo Head。Backbone是Yolo V5的主干网络,该层提取三个有效特征层,分别为C3(80,80,256)、C3(40,40,512)、C3(30,30,1024)。FPN是对有效层进行特征提取,进行特征融合,目的是结合不同尺度的特征信息。Yolo Head是Yolo V5的分类器和回归器,通过Backbone和FPN获得三个加强过的有效特征层,每个特征层都有宽、高、通道数。

1)Yolo V5使用了残差网络Residual,残差jua年纪可以分为两个部分,一个是1×1的卷积和一次3×3的卷积。残差边部分不做处理,直接与主干的输入、输出结合。

class Bottleneck(nn.Module):
    def __init__(self, c1, c2, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, shortcut, groups, expansion
        super(Bottleneck, self).__init__()
        c_ = int(c2 * e)  # hidden channels
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = Conv(c_, c2, 3, 1, g=g)
        self.add = shortcut and c1 == c2

    def forward(self, x):
        return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))

残差网络,容易进行优化,可以通过增加深度提高准确率,内部的残差块使用了跳跃连接,缓解了在深度神经网络中增加深度带来的梯度消失问题。

2)使用CSPnet网络结构,主干部分继续进行原来的残差块的堆叠,另一部分则是经过少量处理直接连接到最后。

class C3(nn.Module):
    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
        super(C3, self).__init__()
        c_ = int(c2 * e)
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = Conv(c1, c_, 1, 1)
        self.cv3 = Conv(2 * c_, c2, 1)  # optional act=FReLU(c2)
        self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))

    def forward(self, x):
        return self.cv3(torch.cat(
            self.m(self.cv1(x)),
            self.cv2(x),
        ), 1)

3)Focus网络结构,具体则是在一个图像中,每隔一个像素拿到一个值,这个时候获得四个独立的特征层,然后将四个独立的特征层进行堆叠。此时,宽高信息集中到了通道信息,输入通道扩大了四倍,拼接起来的特征层相对于原先的三通道变成了十二个通道。如下图所示。

Yolo v5训练并移植到RK3588S平台_第2张图片

class Focus(nn.Module):

    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
        super(Focus, self).__init__()
        self.conv = Conv(c1 * 4, c2, k, s, p, g, act)

    def forward(self, x):
        return self.conv(torch.cat((x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]), 1))

4)使用SiLU激活函数,SiLU激活是由Sigmoid和ReLU的改进版本。SiLU具有无上界、有下界、平滑、非单调的特点,公式为:

F(x) = x × sigmoid(x)

Yolo v5训练并移植到RK3588S平台_第3张图片

5) 使用了SPP结构,通过不同池化核大小的最大池化进行特征提取,提高网络的感受野大小。

class SPP(nn.Module):

    def __init__(self, c1, c2, k=(5, 9, 13)):
        super(SPP, self).__init__()
        c_ = c1 // 2
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
        self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])

    def forward(self, x):
        x = self.cv1(x)
        return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))

主干代码如下:

import torch
import torch.nn as nn


def autopad(k, p=None):  # kernel, padding
    # pad to 'same'
    if p is None:
        p = k // 2 if isinstance(k, int) else [x // 2 for x in k]  # auto-pad
    return p


class Conv(nn.Module):

    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
        super(Conv, self).__init__()
        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
        self.bn = nn.BatchNorm2d(c2)
        self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())

    def forward(self, x):
        return self.act(self.bn(self.conv(x)))

    def forward_fuse(self, x):
        return self.act(self.conv(x))


class Bottleneck(nn.Module):
    def __init__(self, c1, c2, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, shortcut, groups, expansion
        super(Bottleneck, self).__init__()
        c_ = int(c2 * e)  # hidden channels
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = Conv(c_, c2, 3, 1, g=g)
        self.add = shortcut and c1 == c2

    def forward(self, x):
        return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))


class C3(nn.Module):
    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
        super(C3, self).__init__()
        c_ = int(c2 * e)
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = Conv(c1, c_, 1, 1)
        self.cv3 = Conv(2 * c_, c2, 1)  # optional act=FReLU(c2)
        self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))

    def forward(self, x):
        return self.cv3(torch.cat(
            self.m(self.cv1(x)),
            self.cv2(x),
        ), 1)


class SPP(nn.Module):

    def __init__(self, c1, c2, k=(5, 9, 13)):
        super(SPP, self).__init__()
        c_ = c1 // 2
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
        self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])

    def forward(self, x):
        x = self.cv1(x)
        return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))


class Focus(nn.Module):

    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
        super(Focus, self).__init__()
        self.conv = Conv(c1 * 4, c2, k, s, p, g, act)

    def forward(self, x):
        return self.conv(torch.cat((x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]), 1))


class CSPDarknet(nn.Module):

    def __init__(self, base_channels, base_depth):
        super(CSPDarknet, self).__init__()

        # 载入图片是[640,640,3],base_channels为64

        self.stem = Focus(3, base_channels, k=3) # 640, 640, 3 -> 320, 320, 12 -> 320, 320, 64
        self.dark2 = nn.Sequential(
            Conv(base_channels, base_channels * 2, 3, 2),   # 完成卷积之后,320, 320, 64 -> 160, 160, 128
            C3(base_channels * 2, base_channels * 2, base_depth),  # 完成C3之后,160, 160, 128 -> 160, 160, 128
        )

        self.dark3 = nn.Sequential(
            Conv(base_channels * 2, base_channels * 4, 3, 2),  # 完成卷积之后,160, 160, 128 -> 80, 80, 256
            C3(base_channels * 4, base_channels * 4, base_depth * 3),   # 完成C3之后,80, 80, 256 -> 80, 80, 256
                                                                        #                   在这里引出有效特征层80, 80, 256
                                                                        #                   进行加强特征提取网络FPN的构建
        )

        self.dark4 = nn.Sequential(
            Conv(base_channels * 4, base_channels * 8, 3, 2),  # 完成卷积之后,80, 80, 256 -> 40, 40, 512
            C3(base_channels * 8, base_channels * 8, base_depth * 3),     # 完成C3之后,40, 40, 512 -> 40, 40, 512
                                                                          #                   在这里引出有效特征层40, 40, 512
                                                                          #                   进行加强特征提取网络FPN的构建
        )

        self.dark5 = nn.Sequential(
            Conv(base_channels * 8, base_channels * 16, 3, 2),  # 完成卷积之后,40, 40, 512 -> 20, 20, 1024
            SPP(base_channels * 16, base_channels * 16),      # 完成SPP之后,20, 20, 1024 -> 20, 20, 1024
            C3(base_channels * 16, base_channels * 16, base_depth, shortcut=False),    # 完成C3之后,20, 20, 1024 -> 20, 20, 1024
        )

    def forward(self, x):
        x = self.stem(x)  # Focus
        x = self.dark2(x)
        x = self.dark3(x)  # dark3的输出为80, 80, 256,是一个有效特征层
        feat_1 = x
        x = self.dark4(x)  # dark4的输出为40, 40, 512,是一个有效特征层
        feat_2 = x
        x = self.dark5(x)  # dark5的输出为20, 20, 1024,是一个有效特征层
        feat_3 = x
        return  feat_1, feat_2, feat_3

6)在特征利用部分,YoloV5提取多特征层进行目标检测,一共提取三个特征层。
三个特征层位于主干部分CSPdarknet的不同位置,分别位于中间层,中下层,底层,当输入为(640,640,3)的时候,三个特征层的shape分别为feat1=(80,80,256)、feat2=(40,40,512)、feat3=(20,20,1024)。

在获得三个有效特征层后,我们利用这三个有效特征层进行FPN层的构建,构建方式为:

a)   feat3=(20,20,1024)的特征层进行1次1X1卷积调整通道后获得P5,P5进行上采样UmSampling2d后与feat2=(40,40,512)特征层进行结合,然后使用CSPLayer进行特征提取获得P5_upsample,此时获得的特征层为(40,40,512)。


b)  P5_upsample=(40,40,512)的特征层进行1次1X1卷积调整通道后获得P4,P4进行上采样UmSampling2d后与feat1=(80,80,256)特征层进行结合,然后使用CSPLayer进行特征提取P3_out,此时获得的特征层为(80,80,256)。


c)  P3_out=(80,80,256)的特征层进行一次3x3卷积进行下采样,下采样后与P4堆叠,然后使用CSPLayer进行特征提取P4_out,此时获得的特征层为(40,40,512)。


d)  P4_out=(40,40,512)的特征层进行一次3x3卷积进行下采样,下采样后与P5堆叠,然后使用CSPLayer进行特征提取P5_out,此时获得的特征层为(20,20,1024)。

特征金字塔可以将不同shape的特征层进行特征融合,有利于提取出更好的特征。

import torch
import torch.nn as nn

from utils.CSPDarknet import *


# 只使用's'模型,其他的暂时不考虑
class YoloNet(nn.Module):

    def __init__(self, anchors_mask, num_classes, input_shape=[640, 640]):
        super(YoloNet, self).__init__()
        dep_mul, wid_mul = 0.33, 0.50
        base_channels = int(wid_mul * 64)  # 64
        base_depth = max(round(dep_mul * 3), 1)  # 3

        self.backbone = CSPDarknet(base_channels, base_depth)

        self.upsample = nn.Upsample(scale_factor=2, mode="nearest")

        self.conv_for_feat3 = Conv(base_channels * 16, base_channels * 8, 1, 1)
        self.conv3_for_upsample1 = C3(base_channels * 16, base_channels * 8, base_depth, shortcut=False)

        self.conv_for_feat2 = Conv(base_channels * 8, base_channels * 4, 1, 1)
        self.conv3_for_upsample2 = C3(base_channels * 8, base_channels * 4, base_depth, shortcut=False)

        self.down_sample1 = Conv(base_channels * 4, base_channels * 4, 3, 2)
        self.conv3_for_downsample1 = C3(base_channels * 8, base_channels * 8, base_depth, shortcut=False)

        self.down_sample2 = Conv(base_channels * 8, base_channels * 8, 3, 2)
        self.conv3_for_downsample2 = C3(base_channels * 16, base_channels * 16, base_depth, shortcut=False)

        # 80, 80, 256 => 80, 80, 3 * (5 + num_classes) => 80, 80, 3 * (4 + 1 + num_classes)
        self.yolo_head_P3 = nn.Conv2d(base_channels * 4, len(anchors_mask[2]) * (5 + num_classes), 1)
        # 40, 40, 512 => 40, 40, 3 * (5 + num_classes) => 40, 40, 3 * (4 + 1 + num_classes)
        self.yolo_head_P4 = nn.Conv2d(base_channels * 8, len(anchors_mask[1]) * (5 + num_classes), 1)
        # 20, 20, 1024 => 20, 20, 3 * (5 + num_classes) => 20, 20, 3 * (4 + 1 + num_classes)
        self.yolo_head_P5 = nn.Conv2d(base_channels * 16, len(anchors_mask[0]) * (5 + num_classes), 1)

    def forward(self, x):
        feat1, feat2, feat3 = self.backbone(x)

        # 20, 20, 1024 -> 20, 20, 512
        P5 = self.conv_for_feat3(feat3)
        # 20, 20, 512 -> 40, 40, 512
        P5_upsample = self.upsample(P5)

        # 40, 40, 512 -> 40, 40, 1024
        P4 = torch.cat([P5_upsample, feat2], 1)
        # 40, 40, 1024 -> 40, 40, 512
        P4 = self.conv3_for_upsample1(P4)

        # 40, 40, 512 -> 40, 40, 256
        P4 = self.conv_for_feat2(P4)
        # 40, 40, 256 -> 80, 80, 256
        P4_upsample = self.upsample(P4)

        # 80, 80, 256 cat 80, 80, 256 -> 80, 80, 512
        P3 = torch.cat([P4_upsample, feat1], 1)
        # 80, 80, 512 -> 80, 80, 256
        P3 = self.conv3_for_upsample2(P3)

        # 80, 80, 256 -> 40, 40, 256
        P3_downsample = self.down_sample1(P3)
        # 40, 40, 256 cat 40, 40, 256 -> 40, 40, 512
        P4 = torch.cat([P3_downsample, P4], 1)
        # 40, 40, 512 -> 40, 40, 512
        P4 = self.conv3_for_downsample1(P4)

        # 40, 40, 512 -> 20, 20, 512
        P4_downsample = self.down_sample2(P4)
        # 20, 20, 512 cat 20, 20, 512 -> 20, 20, 1024
        P5 = torch.cat([P4_downsample, P5], 1)
        # 20, 20, 1024 -> 20, 20, 1024
        P5 = self.conv3_for_downsample2(P5)

        # ---------------------------------------------------#
        #   第三个特征层
        #   y3=(batch_size,75,80,80)
        # ---------------------------------------------------#
        out2 = self.yolo_head_P3(P3)
        # ---------------------------------------------------#
        #   第二个特征层
        #   y2=(batch_size,75,40,40)
        # ---------------------------------------------------#
        out1 = self.yolo_head_P4(P4)
        # ---------------------------------------------------#
        #   第一个特征层
        #   y1=(batch_size,75,20,20)
        # ---------------------------------------------------#
        out0 = self.yolo_head_P5(P5)

        return out0, out1, out2
这里,得到整个Yolo V5的网络,输出三个特征层,下面开始训练Yolo V5算法模型。

2、训练

        训练使用的是COCO128数据集进行训练,电脑的显卡太弱了,只能选择这种少量的数据集作为验证。训练部分没啥好说的,慢慢等待就是了。需要注意的是,训练好的模型文件格式是yolov5.pt格式文件,需要使用torch.onnx.export函数将其转化为onnx格式的文件。之后,使用RKNN-TooKit2提供的API,将onnx格式的文件转化为rknn格式的文件。

三、部署

        部署到RK3588S平台之前,需要先将环境搭建好,连接好硬件。使用RKNN-TooKit2提供的API进行测试,测试代码比较简单,这里就不放了。常用的API可以参考博文:RKNN-Toolkit2相关API介绍(二)

看着API介绍就能写好代码。

测试结果如下:

Yolo v5训练并移植到RK3588S平台_第4张图片

自己训练的模型文件,跟使用别人的模型文件进行测试,差别还是有的,没有别人的准确率高。

四、总结

        RK3588S部署Yolo V5还是比较容易的,当然也需要进行一些必要的修改。主要的困难依然在Yolo V5算法本身身上,至少我部署的时候没发现特别大的问题。

你可能感兴趣的:(边缘计算)