Yolo v5训练并移植到RK3588S平台 |
您所在的位置：网站首页 › yolo预训练模型需要多少算力 › Yolo v5训练并移植到RK3588S平台 |
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算法大致可以分为三个部分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网络结构，具体则是在一个图像中，每隔一个像素拿到一个值，这个时候获得四个独立的特征层，然后将四个独立的特征层进行堆叠。此时，宽高信息集中到了通道信息，输入通道扩大了四倍，拼接起来的特征层相对于原先的三通道变成了十二个通道。如下图所示。
4）使用SiLU激活函数，SiLU激活是由Sigmoid和ReLU的改进版本。SiLU具有无上界、有下界、平滑、非单调的特点，公式为： F（x） = x × sigmoid(x)
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_36）在特征利用部分，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介绍就能写好代码。 测试结果如下：
自己训练的模型文件，跟使用别人的模型文件进行测试，差别还是有的，没有别人的准确率高。 四、总结 RK3588S部署Yolo V5还是比较容易的，当然也需要进行一些必要的修改。主要的困难依然在Yolo V5算法本身身上，至少我部署的时候没发现特别大的问题。 |
图1 Yolo V5结构图
【本文地址】
今日新闻 |
推荐新闻 |