【计算机视觉】:基于PyTorch的YoloV5目标检测平台

基于PyTorch的YoloV5目标检测平台

参考:
文章转载记录来自
Bubbliiiing大佬博客
Bubbliiiing大佬GIthub地址
Bubbliiiing大佬B站视频讲解
20系列及以下显卡环境配置pytorch代码对应的pytorch版本为1.2,配置方法地址
30系显卡由于框架更新不可使用上述环境配置教程。 配置如下: pytorch代码对应的pytorch版本为1.7.0,cuda为11.0,cudnn为8.0.5
江大白知乎 、博客等
pytorch深度学习框架
YoloV5目标检测算法
YoloV5l版本

YoloV5改进的部分改进

1、主干部分:使用了 Focus网络结构,具体操作是在一张图片中每隔一个像素拿到一个值,这个时候获得了四个独立的特征层,然后将四个独立的特征层进行堆叠,此时宽高信息就集中到了通道信息,输入通道扩充了四倍。该结构在YoloV5第5版之前有所应用,最新版本中未使用(后面改为了6*6的卷积,效果更好)。

2、数据增强:Mosaic数据增强、Mosaic利用了四张图片进行拼接实现数据中增强,根据论文所说其拥有一个巨大的优点是丰富检测物体的背景!且在BN计算的时候一下子会计算四张图片的数据!

3、多正样本匹配:在之前的Yolo系列里面,在训练时每一个真实框对应一个正样本,即在训练时,每一个真实框仅由一个先验框负责预测。YoloV5中为了加快模型的训练效率,增加了正样本的数量,在训练时,每一个真实框可以由多个先验框负责预测。

  • 输入端:在模型训练阶段,提出了一些改进思路,主要包括Mosaic数据增强、自适应锚框计算、自适应图片缩放;
  • 基准网络:融合其它检测算法中的一些新思路,主要包括:Focus结构与CSP结构;
  • Neck网络:目标检测网络在BackBone与最后的Head输出层之间往往会插入一些层,Yolov5中添加了FPN+PAN结构;
  • Head输出层:输出层的锚框机制与YOLOv4相同,主要改进的是训练时的损失函数GIOU_Loss,以及预测框筛选的DIOU_nms。

YoloV5思路

一、 整体结构

【计算机视觉】:基于PyTorch的YoloV5目标检测平台
和之前版本的Yolo类似,整个YoloV5可以依然可以分为三个部分,分别是 Backbone,FPN以及Yolo Head

Backbone可以被称作YoloV5的 主干特征提取网络,根据它的结构以及之前Yolo主干的叫法,我一般叫它 CSPDarknet(Cross Stage Partial),输入的图片首先会在CSPDarknet里面进行 特征提取,提取到的特征可以被称作 特征层是输入图片的特征集合。在主干部分,我们 获取了三个特征层进行下一步网络的构建,这三个特征层我称它为 有效特征层

FPN+PAN:FPN可以被称作YoloV5的 加强特征提取网络(特征金字塔),在主干部分获得的三个有效特征层会在这一部分进行 特征融合,特征融合的目的是 结合不同尺度的特征信息。在FPN部分,已经获得的有效特征层被用于继续提取特征。在YoloV5里 依然使用到了Panet (基于提议的实例分割框架下的路径聚合网络Path Aggregation Network)的结构,我们不仅会对特征进行上采样实现特征融合,还会对特征再次进行下采样实现特征融合。

PANet(Path Aggregation Network)最大的贡献是提出了一个 自顶向下和自底向上的双向融合骨干网络,同时在最底层和最高层之间添加了一条“short-cut”,用于缩短层之间的路径。
PANet还提出了自适应特征池化和全连接融合两个模块。
其中自适应特征池化可以用于聚合不同层之间的特征,保证特征的完整性和多样性,而通过全连接融合可以得到更加准确的预测mask。

Yolo Head是YoloV5的 分类器与回归器(分类:离散变量预测如天气晴或阴,回归:连续变量预测如气温几度),通过CSPDarknet和FPN,我们已经可以获得三个加强过的有效特征层。每一个特征层都有宽、高和通道数,此时我们可以 将特征图看作一个又一个特征点的集合,每一个特征点都有通道数个特征。Yolo Head实际上所做的工作就是 对特征点进行判断,判断特征点是否有物体与其对应。与以前版本的Yolo一样,YoloV5所用的解耦头是一起的,也就是分类和回归在一个1X1卷积里实现。

因此,整个YoloV5网络所作的工作就是 特征提取-特征加强-预测特征点对应的物体情况。

; 二、网络结构分析

(1)输入端:Mosaic数据增强、自适应锚框计算、自适应图片缩放
(2)Backbone:Focus结构,CSP结构
(3)Neck:FPN+PAN结构
(4)Prediction:GIOU_Loss

1、主干网络Backbone(即CSPDarknet)

YoloV5使用的Backbone(主干特征提取网络)为CSPDarknet,五个重要特点:
1、使用了 残差网络Residual,CSPDarknet中的残差卷积可以分为两个部分, 主干部分是一次 1X1的卷积和一次 3X3的卷积残差边部分不做任何处理,直接将主干的输入与输出结合。 整个YoloV5的主干部分都由残差卷积构成

【计算机视觉】:基于PyTorch的YoloV5目标检测平台

class Bottleneck(nn.Module):

    def __init__(self, c1, c2, shortcut=True, g=1, e=0.5):
        super(Bottleneck, self).__init__()
        c_ = int(c2 * e)
        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(Cross Stage Partial Network)网络结构(每层是CSPlayer),CSPnet结构并不算复杂,就是将原来的残差块的堆叠进行了一个拆分,拆成左右两部分: 主干部分继续进行原来的残差块的堆叠;另一部分则像一个残差边一样,经过少量处理直接连接到最后。因此可以认为CSP中存在一个大的残差边。

引出一个大的残差边,跳过特征提取过程,将大结构块的输入和输出直接相接
在 CSPResNe(X)t 结构中,输入特征图在通道维度被分为两个部分,第一部分被保留下来,第二部分则经过多个残差块向后传递,最后将两者在 CSPNet 结构的末端进行合并,这样跨阶段拆分与合并的网络构造有效降低了梯度信息重复的可能性,增加了梯度组合的多样性,有利于提高模型的学习能力,并且降低了网络中的数据传递量与计算量。

【计算机视觉】:基于PyTorch的YoloV5目标检测平台

class C3(nn.Module):

    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
        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)
        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)
            )
            , dim=1))

3、使用了Focus网络结构,这个网络结构是在YoloV5里面使用到比较有趣的网络结构,具体操作是在一张图片中 每隔一个像素拿到一个值,这个时候获得了 四个独立的特征层,然后将四个独立的特征层进行 堆叠此时宽高信息就集中到了通道信息,输入通道扩充了四倍。拼接起来的特征层相对于原先的 三通道变成了十二个通道,下图很好的展示了Focus结构,一看就能明白。( 高宽压缩通道扩张)(新的YoloV5用6*6的卷积替代了Focus,更快)

【计算机视觉】:基于PyTorch的YoloV5目标检测平台

class Focus(nn.Module):
    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):
        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具备 无上界有下界、平滑、非单调的特性。SiLU在深层模型上的效果优于 ReLU。可以看做是平滑的ReLU激活函数。

【计算机视觉】:基于PyTorch的YoloV5目标检测平台

class SiLU(nn.Module):
    @staticmethod
    def forward(x):
        return x * torch.sigmoid(x)

ReLU激活函数

【计算机视觉】:基于PyTorch的YoloV5目标检测平台
5、使用了SPP结构,通过不同池化核大小的最大池化进行特征提取,提高网络的感受野。在YoloV4中,SPP是用在FPN里面的,在YoloV5中,SPP模块被用在了主干特征提取网络中。
SPP
【计算机视觉】:基于PyTorch的YoloV5目标检测平台
卷积调整通道数之后利用不同池化核的最大池化核(5 9 13 )来进行特征提取 除了5 9 13 还有一个没有处理的
四个部分堆叠然后卷积标准化加激活函数进行通道数的调整

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))

主干代码CSPDarknet.py:

import torch
import torch.nn as nn

class SiLU(nn.Module):
    @staticmethod
    def forward(x):
        return x * torch.sigmoid(x)

def autopad(k, p=None):
    if p is None:
        p = k // 2 if isinstance(k, int) else [x // 2 for x in k]
    return p

class Focus(nn.Module):
    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):
        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 Conv(nn.Module):
    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):
        super(Conv, self).__init__()
        self.conv   = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
        self.bn     = nn.BatchNorm2d(c2, eps=0.001, momentum=0.03)
        self.act    = 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 fuseforward(self, x):
        return self.act(self.conv(x))

class Bottleneck(nn.Module):

    def __init__(self, c1, c2, shortcut=True, g=1, e=0.5):
        super(Bottleneck, self).__init__()
        c_ = int(c2 * e)
        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):
        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)
        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)
            )
            , dim=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 CSPDarknet(nn.Module):

    def __init__(self, base_channels, base_depth, phi, pretrained):
        super().__init__()

        self.stem       = Focus(3, base_channels, k=3)

        self.dark2 = nn.Sequential(

            Conv(base_channels, base_channels * 2, 3, 2),

            C3(base_channels * 2, base_channels * 2, base_depth),
        )

        self.dark3 = nn.Sequential(
            Conv(base_channels * 2, base_channels * 4, 3, 2),
            C3(base_channels * 4, base_channels * 4, base_depth * 3),
        )

        self.dark4 = nn.Sequential(
            Conv(base_channels * 4, base_channels * 8, 3, 2),
            C3(base_channels * 8, base_channels * 8, base_depth * 3),
        )

        self.dark5 = nn.Sequential(
            Conv(base_channels * 8, base_channels * 16, 3, 2),
            SPP(base_channels * 16, base_channels * 16),
            C3(base_channels * 16, base_channels * 16, base_depth, shortcut=False),
        )
        if pretrained:
            url = {
                's' : 'https://github.com/bubbliiiing/yolov5-pytorch/releases/download/v1.0/cspdarknet_s_backbone.pth',
                'm' : 'https://github.com/bubbliiiing/yolov5-pytorch/releases/download/v1.0/cspdarknet_m_backbone.pth',
                'l' : 'https://github.com/bubbliiiing/yolov5-pytorch/releases/download/v1.0/cspdarknet_l_backbone.pth',
                'x' : 'https://github.com/bubbliiiing/yolov5-pytorch/releases/download/v1.0/cspdarknet_x_backbone.pth',
            }[phi]
            checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", model_dir="./model_data")
            self.load_state_dict(checkpoint, strict=False)
            print("Load weights from ", url.split('/')[-1])

    def forward(self, x):
        x = self.stem(x)
        x = self.dark2(x)

        x = self.dark3(x)
        feat1 = x

        x = self.dark4(x)
        feat2 = x

        x = self.dark5(x)
        feat3 = x
        return feat1, feat2, feat3

2、构建FPN特征金字塔(加强特征提取网络)进行加强特征提取

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

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

  1. feat3=(20,20,1024)的特征层进行1次1X1卷积调整通道后获得P5, P5进行上采样UmSampling2d后与feat2=(40,40,512)特征层进行结合,然后使用CSPLayer进行特征提取获得P5_upsample,此时获得的特征层为(40,40,512)。
  2. P5_upsample=(40,40,512)的特征层进行1次1X1卷积调整通道后获得P4, P4进行上采样UmSampling2d后与feat1=(80,80,256)特征层进行结合,然后使用CSPLayer进行特征提取P3_out,此时获得的特征层为(80,80,256)。
  3. P3_out=(80,80,256)的特征层进行一次3×3卷积进行下采样, 下采样后与P4堆叠,然后使用CSPLayer进行特征提取P4_out,此时获得的特征层为(40,40,512)。
  4. P4_out=(40,40,512)的特征层进行一次3×3卷积进行下采样, 下采样后与P5堆叠,然后使用CSPLayer进行特征提取P5_out,此时获得的特征层为(20,20,1024)。

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

nets/yolo.py

import torch
import torch.nn as nn

from nets.ConvNext import ConvNeXt_Small, ConvNeXt_Tiny
from nets.CSPdarknet import C3, Conv, CSPDarknet
from nets.Swin_transformer import Swin_transformer_Tiny

class YoloBody(nn.Module):
    def __init__(self, anchors_mask, num_classes, phi, backbone='cspdarknet', pretrained=False, input_shape=[640, 640]):
        super(YoloBody, self).__init__()
        depth_dict          = {'s' : 0.33, 'm' : 0.67, 'l' : 1.00, 'x' : 1.33,}
        width_dict          = {'s' : 0.50, 'm' : 0.75, 'l' : 1.00, 'x' : 1.25,}
        dep_mul, wid_mul    = depth_dict[phi], width_dict[phi]

        base_channels       = int(wid_mul * 64)
        base_depth          = max(round(dep_mul * 3), 1)

        self.backbone_name  = backbone
        if backbone == "cspdarknet":

            self.backbone   = CSPDarknet(base_channels, base_depth, phi, pretrained)
        else:

            self.backbone       = {
                'convnext_tiny'         : ConvNeXt_Tiny,
                'convnext_small'        : ConvNeXt_Small,
                'swin_transfomer_tiny'  : Swin_transformer_Tiny,
            }[backbone](pretrained=pretrained, input_shape=input_shape)
            in_channels         = {
                'convnext_tiny'         : [192, 384, 768],
                'convnext_small'        : [192, 384, 768],
                'swin_transfomer_tiny'  : [192, 384, 768],
            }[backbone]
            feat1_c, feat2_c, feat3_c = in_channels
            self.conv_1x1_feat1 = Conv(feat1_c, base_channels * 4, 1, 1)
            self.conv_1x1_feat2 = Conv(feat2_c, base_channels * 8, 1, 1)
            self.conv_1x1_feat3 = Conv(feat3_c, base_channels * 16, 1, 1)

        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)

        self.yolo_head_P3 = nn.Conv2d(base_channels * 4, len(anchors_mask[2]) * (5 + num_classes), 1)

        self.yolo_head_P4 = nn.Conv2d(base_channels * 8, len(anchors_mask[1]) * (5 + num_classes), 1)

        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)
        if self.backbone_name != "cspdarknet":
            feat1 = self.conv_1x1_feat1(feat1)
            feat2 = self.conv_1x1_feat2(feat2)
            feat3 = self.conv_1x1_feat3(feat3)

        P5          = self.conv_for_feat3(feat3)

        P5_upsample = self.upsample(P5)

        P4          = torch.cat([P5_upsample, feat2], 1)

        P4          = self.conv3_for_upsample1(P4)

        P4          = self.conv_for_feat2(P4)

        P4_upsample = self.upsample(P4)

        P3          = torch.cat([P4_upsample, feat1], 1)

        P3          = self.conv3_for_upsample2(P3)

        P3_downsample = self.down_sample1(P3)

        P4 = torch.cat([P3_downsample, P4], 1)

        P4 = self.conv3_for_downsample1(P4)

        P4_downsample = self.down_sample2(P4)

        P5 = torch.cat([P4_downsample, P5], 1)

        P5 = self.conv3_for_downsample2(P5)

        out2 = self.yolo_head_P3(P3)

        out1 = self.yolo_head_P4(P4)

        out0 = self.yolo_head_P5(P5)
        return out0, out1, out2

3、利用Yolo Head获得预测结果

【计算机视觉】:基于PyTorch的YoloV5目标检测平台
利用FPN特征金字塔,我们可以 获得三个加强特征,这三个加强特征的shape分别为(20,20,1024)、(40,40,512)、(80,80,256),然后我们利用这三个shape的特征层传入Yolo Head获得预测结果

对于每一个特征层,我们可以获得利用一个卷积调整通道数,最终的通道数和需要区分的种类个数相关,在YoloV5里,每一个特征层上每一个特征点存在3个先验框。

如果使用的是voc训练集,类则为20种,最后的维度应该为75 = 3×25,三个特征层的shape为(20,20,75),(40,40,75),(80,80,75)。本人使用的也是VOC数据集
最后的75可以拆分成3个25,对应3个先验框的25个参数,25可以拆分成4+1+20。
前4个参数用于判断每一个特征点的回归参数,回归参数调整后可以获得预测框;
第5个参数用于判断每一个特征点是否包含物体;
最后20个参数用于判断每一个特征点所包含的物体种类。

如果使用的是coco训练集,类则为80种,最后的维度应该为255 = 3×85,三个特征层的shape为(20,20,255),(40,40,255),(80,80,255)
最后的255可以拆分成3个85,对应3个先验框的85个参数,85可以拆分成4+1+80。
前4个参数用于判断每一个特征点的回归参数,回归参数调整后可以获得预测框;
第5个参数用于判断每一个特征点是否包含物体;
最后80个参数用于判断每一个特征点所包含的物体种类。

nets/yolo.py

import torch
import torch.nn as nn

from nets.ConvNext import ConvNeXt_Small, ConvNeXt_Tiny
from nets.CSPdarknet import C3, Conv, CSPDarknet
from nets.Swin_transformer import Swin_transformer_Tiny

class YoloBody(nn.Module):
    def __init__(self, anchors_mask, num_classes, phi, backbone='cspdarknet', pretrained=False, input_shape=[640, 640]):
        super(YoloBody, self).__init__()
        depth_dict          = {'s' : 0.33, 'm' : 0.67, 'l' : 1.00, 'x' : 1.33,}
        width_dict          = {'s' : 0.50, 'm' : 0.75, 'l' : 1.00, 'x' : 1.25,}
        dep_mul, wid_mul    = depth_dict[phi], width_dict[phi]

        base_channels       = int(wid_mul * 64)
        base_depth          = max(round(dep_mul * 3), 1)

        self.backbone_name  = backbone
        if backbone == "cspdarknet":

            self.backbone   = CSPDarknet(base_channels, base_depth, phi, pretrained)
        else:

            self.backbone       = {
                'convnext_tiny'         : ConvNeXt_Tiny,
                'convnext_small'        : ConvNeXt_Small,
                'swin_transfomer_tiny'  : Swin_transformer_Tiny,
            }[backbone](pretrained=pretrained, input_shape=input_shape)
            in_channels         = {
                'convnext_tiny'         : [192, 384, 768],
                'convnext_small'        : [192, 384, 768],
                'swin_transfomer_tiny'  : [192, 384, 768],
            }[backbone]
            feat1_c, feat2_c, feat3_c = in_channels
            self.conv_1x1_feat1 = Conv(feat1_c, base_channels * 4, 1, 1)
            self.conv_1x1_feat2 = Conv(feat2_c, base_channels * 8, 1, 1)
            self.conv_1x1_feat3 = Conv(feat3_c, base_channels * 16, 1, 1)

        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)

        self.yolo_head_P3 = nn.Conv2d(base_channels * 4, len(anchors_mask[2]) * (5 + num_classes), 1)

        self.yolo_head_P4 = nn.Conv2d(base_channels * 8, len(anchors_mask[1]) * (5 + num_classes), 1)

        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)
        if self.backbone_name != "cspdarknet":
            feat1 = self.conv_1x1_feat1(feat1)
            feat2 = self.conv_1x1_feat2(feat2)
            feat3 = self.conv_1x1_feat3(feat3)

        P5          = self.conv_for_feat3(feat3)

        P5_upsample = self.upsample(P5)

        P4          = torch.cat([P5_upsample, feat2], 1)

        P4          = self.conv3_for_upsample1(P4)

        P4          = self.conv_for_feat2(P4)

        P4_upsample = self.upsample(P4)

        P3          = torch.cat([P4_upsample, feat1], 1)

        P3          = self.conv3_for_upsample2(P3)

        P3_downsample = self.down_sample1(P3)

        P4 = torch.cat([P3_downsample, P4], 1)

        P4 = self.conv3_for_downsample1(P4)

        P4_downsample = self.down_sample2(P4)

        P5 = torch.cat([P4_downsample, P5], 1)

        P5 = self.conv3_for_downsample2(P5)

        out2 = self.yolo_head_P3(P3)

        out1 = self.yolo_head_P4(P4)

        out0 = self.yolo_head_P5(P5)
        return out0, out1, out2

三、预测结果的解码

1、获得预测框与得分

由第二步我们可以获得三个特征层的预测结果,shape分别为 (N,20,20,75),(N,40,40,75),(N,80,80,75) 的数据。

但是这个预测结果并不对应着最终的预测框在图片上的位置,还需要解码才可以完成。在YoloV5里,每一个特征层上每一个特征点存在3个先验框。

每个特征层最后的255可以拆分成3个75,对应3个先验框的75个参数,我们先将其reshape一下,其结果为(N,20,20,3,25),(N,40.40,3,25),(N,80,80,3,25)。

如果使用的是voc训练集,类则为20种,最后的维度应该为75 = 3×25,三个特征层的shape为(20,20,25),(40,40,25),(80,80,25)。本人使用的也是VOC数据集
最后的75可以拆分成3个25,对应3个先验框的25个参数,25可以拆分成4+1+20。
前4个参数用于判断每一个特征点的回归参数,回归参数调整后可以获得预测框;
第5个参数用于判断每一个特征点是否包含物体;
最后20个参数用于判断每一个特征点所包含的物体种类。

(N,20,20,3,25) 这个特征层为例, 该特征层相当于将图像划分成20×20个特征点,如果某个特征点落在物体的对应框内,就用于预测该物体。

如图所示,蓝色的点为20×20的特征点,此时我们对左图 黑色点的三个先验框进行解码操作演示:
1、进行中心预测点的计算, 利用Regression预测结果前两个序号的内容对特征点的三个先验框中心坐标进行偏移,偏移后是右图红色的三个点;
2、进行预测框宽高的计算, 利用Regression预测结果后两个序号的内容求指数后获得预测框的宽高;
3、此时获得的预测框就可以绘制在图片上了。

【计算机视觉】:基于PyTorch的YoloV5目标检测平台
除去这样的解码操作,还有 非极大抑制的操作需要进行, 防止同一种类的框的堆积

utils/utils_bbox.py

    def decode_box(self, inputs):
        outputs = []
        for i, input in enumerate(inputs):

            batch_size      = input.size(0)
            input_height    = input.size(2)
            input_width     = input.size(3)

            stride_h = self.input_shape[0] / input_height
            stride_w = self.input_shape[1] / input_width

            scaled_anchors = [(anchor_width / stride_w, anchor_height / stride_h) for anchor_width, anchor_height in self.anchors[self.anchors_mask[i]]]

            prediction = input.view(batch_size, len(self.anchors_mask[i]),
                                    self.bbox_attrs, input_height, input_width).permute(0, 1, 3, 4, 2).contiguous()

            x = torch.sigmoid(prediction[..., 0])
            y = torch.sigmoid(prediction[..., 1])

            w = torch.sigmoid(prediction[..., 2])
            h = torch.sigmoid(prediction[..., 3])

            conf        = torch.sigmoid(prediction[..., 4])

            pred_cls    = torch.sigmoid(prediction[..., 5:])

            FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor
            LongTensor  = torch.cuda.LongTensor if x.is_cuda else torch.LongTensor

            grid_x = torch.linspace(0, input_width - 1, input_width).repeat(input_height, 1).repeat(
                batch_size * len(self.anchors_mask[i]), 1, 1).view(x.shape).type(FloatTensor)
            grid_y = torch.linspace(0, input_height - 1, input_height).repeat(input_width, 1).t().repeat(
                batch_size * len(self.anchors_mask[i]), 1, 1).view(y.shape).type(FloatTensor)

            anchor_w = FloatTensor(scaled_anchors).index_select(1, LongTensor([0]))
            anchor_h = FloatTensor(scaled_anchors).index_select(1, LongTensor([1]))
            anchor_w = anchor_w.repeat(batch_size, 1).repeat(1, 1, input_height * input_width).view(w.shape)
            anchor_h = anchor_h.repeat(batch_size, 1).repeat(1, 1, input_height * input_width).view(h.shape)

            pred_boxes          = FloatTensor(prediction[..., :4].shape)
            pred_boxes[..., 0]  = x.data * 2. - 0.5 + grid_x
            pred_boxes[..., 1]  = y.data * 2. - 0.5 + grid_y
            pred_boxes[..., 2]  = (w.data * 2) ** 2 * anchor_w
            pred_boxes[..., 3]  = (h.data * 2) ** 2 * anchor_h

            _scale = torch.Tensor([input_width, input_height, input_width, input_height]).type(FloatTensor)
            output = torch.cat((pred_boxes.view(batch_size, -1, 4) / _scale,
                                conf.view(batch_size, -1, 1), pred_cls.view(batch_size, -1, self.num_classes)), -1)
            outputs.append(output.data)
        return outputs

2、得分筛选与非极大值抑制

得到最终的预测结果后还要进行 得分排序与非极大抑制筛选

得分筛选就是 筛选出得分满足confidence置信度的预测框
非极大抑制就是 筛选出一定区域内属于同一种类得分最大的框

得分筛选与非极大抑制的过程可以概括如下:
1、找出该图片中 得分大于门限函数的框。在进行 重合框筛选前就进行得分的筛选可以大幅度减少框的数量
2、对 种类进行循环,非极大抑制的作用是 筛选出一定区域内属于同一种类得分最大的框,对种类进行循环可以 帮助我们对每一个类分别进行非极大抑制。
3、根据得分对该种类进行 从大到小排序
4、每次取出得分最大的框,计算 其与其它所有预测框的重合程度,重合程度过大的则剔除。

得分筛选与非极大抑制后的结果就可以用于绘制预测框了。

nms_iou=0.3(越小越严格即距离相近的两个类别会选取得分高的。即框少了):

【计算机视觉】:基于PyTorch的YoloV5目标检测平台

nms_iou=0.01:

【计算机视觉】:基于PyTorch的YoloV5目标检测平台
    def non_max_suppression(self, prediction, num_classes, input_shape, image_shape, letterbox_image, conf_thres=0.5, nms_thres=0.4):

        box_corner          = prediction.new(prediction.shape)
        box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2
        box_corner[:, :, 1] = prediction[:, :, 1] - prediction[:, :, 3] / 2
        box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2
        box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2
        prediction[:, :, :4] = box_corner[:, :, :4]

        output = [None for _ in range(len(prediction))]

        for i, image_pred in enumerate(prediction):

            class_conf, class_pred = torch.max(image_pred[:, 5:5 + num_classes], 1, keepdim=True)

            conf_mask = (image_pred[:, 4] * class_conf[:, 0] >= conf_thres).squeeze()

            image_pred = image_pred[conf_mask]
            class_conf = class_conf[conf_mask]
            class_pred = class_pred[conf_mask]
            if not image_pred.size(0):
                continue

            detections = torch.cat((image_pred[:, :5], class_conf.float(), class_pred.float()), 1)

            unique_labels = detections[:, -1].cpu().unique()

            if prediction.is_cuda:
                unique_labels = unique_labels.cuda()
                detections = detections.cuda()

            for c in unique_labels:

                detections_class = detections[detections[:, -1] == c]

                keep = nms(
                    detections_class[:, :4],
                    detections_class[:, 4] * detections_class[:, 5],
                    nms_thres
                )

                max_detections = detections_class[keep]

                output[i] = max_detections if output[i] is None else torch.cat((output[i], max_detections))

            if output[i] is not None:
                output[i]           = output[i].cpu().numpy()
                box_xy, box_wh      = (output[i][:, 0:2] + output[i][:, 2:4])/2, output[i][:, 2:4] - output[i][:, 0:2]
                output[i][:, :4]    = self.yolo_correct_boxes(box_xy, box_wh, input_shape, image_shape, letterbox_image)

        return output

四、训练部分

1、计算loss所需内容

计算loss实际上是网络的预测结果和网络的真实结果的对比。
和网络的预测结果一样,网络的损失也由三个部分组成,分别是Reg部分、Obj部分、Cls部分。Reg部分是特征点的回归参数判断、Obj部分是特征点是否包含物体判断、Cls部分是特征点包含的物体的种类。

2、正样本的匹配过程

在YoloV5中,训练时正样本的匹配过程可以分为两部分。
a、匹配先验框。
b、匹配特征点。

所谓 正样本匹配,就是 寻找哪些先验框被认为有对应的真实框,并且负责这个真实框的预测

a、匹配先验框

在YoloV5网络中,一共设计了9个不同大小的先验框。每个输出的特征层对应3个先验框。

对于任何一个真实框gt,YoloV5不再使用iou进行正样本的匹配,而是直接采用高宽比进行匹配,即使用真实框和9个不同大小的先验框计算宽高比。

如果真实框与某个先验框的宽高比例大于设定阈值,则说明该真实框和该先验框匹配度不够,将该先验框认为是负样本。

比如 此时有一个真实框,它的宽高为[200, 200],是一个正方形。YoloV5默认设置的 9个先验框为[10,13], [16,30], [33,23], [30,61], [62,45], [59,119], [116,90], [156,198], [373,326]。设定 阈值门限为4。

此时我们需要计算该真实框和9个先验框的宽高比例。 比较宽高时存在两个情况, 一个是真实框的宽高比先验框大,一个是先验框的宽高比真实框大。 因此我们需要同时计算:真实框的宽高/先验框的宽高;先验框的宽高/真实框的宽高。然后在这其中选取最大值。

下个列表就是比较结果,这是一个shape为[9, 4]的矩阵,9代表9个先验框,4代表真实框的宽高/先验框的宽高;先验框的宽高/真实框的宽高。

[[20.         15.38461538  0.05        0.065     ]
 [12.5         6.66666667  0.08        0.15      ]
 [ 6.06060606  8.69565217  0.165       0.115     ]
 [ 6.66666667  3.27868852  0.15        0.305     ]
 [ 3.22580645  4.44444444  0.31        0.225     ]
 [ 3.38983051  1.68067227  0.295       0.595     ]
 [ 1.72413793  2.22222222  0.58        0.45      ]
 [ 1.28205128  1.01010101  0.78        0.99      ]
 [ 0.53619303  0.61349693  1.865       1.63      ]]

然后对每个先验框的比较结果取最大值。获得下述矩阵:

[20.         12.5         8.69565217  6.66666667  4.44444444  3.38983051
   2.22222222  1.28205128  1.865     ]

之后我们判断,哪些先验框的比较结果的值小于门限。可以知道[59,119], [116,90], [156,198], [373,326]四个先验框均满足需求。

[116,90], [156,198], [373,326]属于20,20的特征层。
[59,119]属于40,40的特征层。

此时我们已经可以判断哪些大小的先验框可用于该真实框的预测。

b、匹配特征点

在过去的Yolo系列中,每个真实框由其中心点所在的网格内的左上角特征点来负责预测。

对于被选中的特征层,首先计算真实框落在哪个网格内, 此时该网格左上角特征点便是一个负责预测的特征点。

同时利用四舍五入规则,找出最近的两个网格,将这三个网格都认为是负责预测该真实框的。

【计算机视觉】:基于PyTorch的YoloV5目标检测平台

红色点表示该真实框的中心,除了当前所处的网格外,其2个最近的邻域网格也被选中。从这里就可以发现预测框的XY轴偏移部分的取值范围不再是0-1,而是0.5-1.5。

找到对应特征点后,对应特征点在a中被选中的先验框负责该真实框的预测

; 3、计算Loss

由第一部分可知,YoloV5的损失由三个部分组成:
1、Reg部分,由第2部分可知道每个真实框对应的先验框,获取到每个框对应的先验框后,取出该先验框对应的预测框,利用真实框和预测框计算CIOU损失,作为Reg部分的Loss组成。
2、Obj部分,由第2部分可知道每个真实框对应的先验框,所有真实框对应的先验框都是正样本,剩余的先验框均为负样本,根据正负样本和特征点的是否包含物体的预测结果计算交叉熵损失,作为Obj部分的Loss组成。
3、Cls部分,由第三部分可知道每个真实框对应的先验框,获取到每个框对应的先验框后,取出该先验框的种类预测结果,根据真实框的种类和先验框的种类预测结果计算交叉熵损失,作为Cls部分的Loss组成。

import math
from copy import deepcopy
from functools import partial

import numpy as np
import torch
import torch.nn as nn

class YOLOLoss(nn.Module):
    def __init__(self, anchors, num_classes, input_shape, cuda, anchors_mask = [[6,7,8], [3,4,5], [0,1,2]], label_smoothing = 0):
        super(YOLOLoss, self).__init__()

        self.anchors        = anchors
        self.num_classes    = num_classes
        self.bbox_attrs     = 5 + num_classes
        self.input_shape    = input_shape
        self.anchors_mask   = anchors_mask
        self.label_smoothing = label_smoothing

        self.threshold      = 4

        self.balance        = [0.4, 1.0, 4]
        self.box_ratio      = 0.05
        self.obj_ratio      = 1 * (input_shape[0] * input_shape[1]) / (640 ** 2)
        self.cls_ratio      = 0.5 * (num_classes / 80)
        self.cuda = cuda

    def clip_by_tensor(self, t, t_min, t_max):
        t = t.float()
        result = (t >= t_min).float() * t + (t < t_min).float() * t_min
        result = (result  t_max).float() * result + (result > t_max).float() * t_max
        return result

    def MSELoss(self, pred, target):
        return torch.pow(pred - target, 2)

    def BCELoss(self, pred, target):
        epsilon = 1e-7
        pred    = self.clip_by_tensor(pred, epsilon, 1.0 - epsilon)
        output  = - target * torch.log(pred) - (1.0 - target) * torch.log(1.0 - pred)
        return output

    def box_giou(self, b1, b2):
"""
        输入为:
        ----------
        b1: tensor, shape=(batch, feat_w, feat_h, anchor_num, 4), xywh
        b2: tensor, shape=(batch, feat_w, feat_h, anchor_num, 4), xywh

        返回为:
        -------
        giou: tensor, shape=(batch, feat_w, feat_h, anchor_num, 1)
"""

        b1_xy       = b1[..., :2]
        b1_wh       = b1[..., 2:4]
        b1_wh_half  = b1_wh/2.

        b1_mins     = b1_xy - b1_wh_half
        b1_maxes    = b1_xy + b1_wh_half

        b2_xy       = b2[..., :2]
        b2_wh       = b2[..., 2:4]
        b2_wh_half  = b2_wh/2.

        b2_mins     = b2_xy - b2_wh_half
        b2_maxes    = b2_xy + b2_wh_half

        intersect_mins  = torch.max(b1_mins, b2_mins)
        intersect_maxes = torch.min(b1_maxes, b2_maxes)
        intersect_wh    = torch.max(intersect_maxes - intersect_mins, torch.zeros_like(intersect_maxes))
        intersect_area  = intersect_wh[..., 0] * intersect_wh[..., 1]
        b1_area         = b1_wh[..., 0] * b1_wh[..., 1]
        b2_area         = b2_wh[..., 0] * b2_wh[..., 1]
        union_area      = b1_area + b2_area - intersect_area
        iou             = intersect_area / union_area

        enclose_mins    = torch.min(b1_mins, b2_mins)
        enclose_maxes   = torch.max(b1_maxes, b2_maxes)
        enclose_wh      = torch.max(enclose_maxes - enclose_mins, torch.zeros_like(intersect_maxes))

        enclose_area    = enclose_wh[..., 0] * enclose_wh[..., 1]
        giou            = iou - (enclose_area - union_area) / enclose_area

        return giou

    def smooth_labels(self, y_true, label_smoothing, num_classes):
        return y_true * (1.0 - label_smoothing) + label_smoothing / num_classes

    def forward(self, l, input, targets=None, y_true=None):

        bs      = input.size(0)
        in_h    = input.size(2)
        in_w    = input.size(3)

        stride_h = self.input_shape[0] / in_h
        stride_w = self.input_shape[1] / in_w

        scaled_anchors  = [(a_w / stride_w, a_h / stride_h) for a_w, a_h in self.anchors]

        prediction = input.view(bs, len(self.anchors_mask[l]), self.bbox_attrs, in_h, in_w).permute(0, 1, 3, 4, 2).contiguous()

        x = torch.sigmoid(prediction[..., 0])
        y = torch.sigmoid(prediction[..., 1])

        w = torch.sigmoid(prediction[..., 2])
        h = torch.sigmoid(prediction[..., 3])

        conf = torch.sigmoid(prediction[..., 4])

        pred_cls = torch.sigmoid(prediction[..., 5:])

        pred_boxes = self.get_pred_boxes(l, x, y, h, w, targets, scaled_anchors, in_h, in_w)

        if self.cuda:
            y_true          = y_true.type_as(x)

        loss    = 0
        n       = torch.sum(y_true[..., 4] == 1)
        if n != 0:

            giou        = self.box_giou(pred_boxes, y_true[..., :4]).type_as(x)
            loss_loc    = torch.mean((1 - giou)[y_true[..., 4] == 1])
            loss_cls    = torch.mean(self.BCELoss(pred_cls[y_true[..., 4] == 1], self.smooth_labels(y_true[..., 5:][y_true[..., 4] == 1], self.label_smoothing, self.num_classes)))
            loss        += loss_loc * self.box_ratio + loss_cls * self.cls_ratio

            tobj        = torch.where(y_true[..., 4] == 1, giou.detach().clamp(0), torch.zeros_like(y_true[..., 4]))
        else:
            tobj        = torch.zeros_like(y_true[..., 4])
        loss_conf   = torch.mean(self.BCELoss(conf, tobj))

        loss        += loss_conf * self.balance[l] * self.obj_ratio

        return loss

    def get_near_points(self, x, y, i, j):
        sub_x = x - i
        sub_y = y - j
        if sub_x > 0.5 and sub_y > 0.5:
            return [[0, 0], [1, 0], [0, 1]]
        elif sub_x < 0.5 and sub_y > 0.5:
            return [[0, 0], [-1, 0], [0, 1]]
        elif sub_x < 0.5 and sub_y < 0.5:
            return [[0, 0], [-1, 0], [0, -1]]
        else:
            return [[0, 0], [1, 0], [0, -1]]

    def get_target(self, l, targets, anchors, in_h, in_w):

        bs              = len(targets)

        noobj_mask      = torch.ones(bs, len(self.anchors_mask[l]), in_h, in_w, requires_grad = False)

        box_best_ratio = torch.zeros(bs, len(self.anchors_mask[l]), in_h, in_w, requires_grad = False)

        y_true          = torch.zeros(bs, len(self.anchors_mask[l]), in_h, in_w, self.bbox_attrs, requires_grad = False)
        for b in range(bs):
            if len(targets[b])==0:
                continue

            batch_target = torch.zeros_like(targets[b])

            batch_target[:, [0,2]] = targets[b][:, [0,2]] * in_w
            batch_target[:, [1,3]] = targets[b][:, [1,3]] * in_h
            batch_target[:, 4] = targets[b][:, 4]
            batch_target = batch_target.cpu()

            ratios_of_gt_anchors = torch.unsqueeze(batch_target[:, 2:4], 1) / torch.unsqueeze(torch.FloatTensor(anchors), 0)
            ratios_of_anchors_gt = torch.unsqueeze(torch.FloatTensor(anchors), 0) /  torch.unsqueeze(batch_target[:, 2:4], 1)
            ratios               = torch.cat([ratios_of_gt_anchors, ratios_of_anchors_gt], dim = -1)
            max_ratios, _        = torch.max(ratios, dim = -1)

            for t, ratio in enumerate(max_ratios):

                over_threshold = ratio < self.threshold
                over_threshold[torch.argmin(ratio)] = True

                for k, mask in enumerate(self.anchors_mask[l]):

                    if not over_threshold[mask]:
                        continue

                    i = torch.floor(batch_target[t, 0]).long()
                    j = torch.floor(batch_target[t, 1]).long()

                    offsets = self.get_near_points(batch_target[t, 0], batch_target[t, 1], i, j)
                    for offset in offsets:
                        local_i = i + offset[0]
                        local_j = j + offset[1]

                        if local_i >= in_w or local_i < 0 or local_j >= in_h or local_j < 0:
                            continue

                        if box_best_ratio[b, k, local_j, local_i] != 0:
                            if box_best_ratio[b, k, local_j, local_i] > ratio[mask]:
                                y_true[b, k, local_j, local_i, :] = 0
                            else:
                                continue

                        c = batch_target[t, 4].long()

                        noobj_mask[b, k, local_j, local_i] = 0

                        y_true[b, k, local_j, local_i, 0] = batch_target[t, 0]
                        y_true[b, k, local_j, local_i, 1] = batch_target[t, 1]
                        y_true[b, k, local_j, local_i, 2] = batch_target[t, 2]
                        y_true[b, k, local_j, local_i, 3] = batch_target[t, 3]
                        y_true[b, k, local_j, local_i, 4] = 1
                        y_true[b, k, local_j, local_i, c + 5] = 1

                        box_best_ratio[b, k, local_j, local_i] = ratio[mask]

        return y_true, noobj_mask

    def get_pred_boxes(self, l, x, y, h, w, targets, scaled_anchors, in_h, in_w):

        bs = len(targets)

        grid_x = torch.linspace(0, in_w - 1, in_w).repeat(in_h, 1).repeat(
            int(bs * len(self.anchors_mask[l])), 1, 1).view(x.shape).type_as(x)
        grid_y = torch.linspace(0, in_h - 1, in_h).repeat(in_w, 1).t().repeat(
            int(bs * len(self.anchors_mask[l])), 1, 1).view(y.shape).type_as(x)

        scaled_anchors_l = np.array(scaled_anchors)[self.anchors_mask[l]]
        anchor_w = torch.Tensor(scaled_anchors_l).index_select(1, torch.LongTensor([0])).type_as(x)
        anchor_h = torch.Tensor(scaled_anchors_l).index_select(1, torch.LongTensor([1])).type_as(x)

        anchor_w = anchor_w.repeat(bs, 1).repeat(1, 1, in_h * in_w).view(w.shape)
        anchor_h = anchor_h.repeat(bs, 1).repeat(1, 1, in_h * in_w).view(h.shape)

        pred_boxes_x    = torch.unsqueeze(x * 2. - 0.5 + grid_x, -1)
        pred_boxes_y    = torch.unsqueeze(y * 2. - 0.5 + grid_y, -1)
        pred_boxes_w    = torch.unsqueeze((w * 2) ** 2 * anchor_w, -1)
        pred_boxes_h    = torch.unsqueeze((h * 2) ** 2 * anchor_h, -1)
        pred_boxes      = torch.cat([pred_boxes_x, pred_boxes_y, pred_boxes_w, pred_boxes_h], dim = -1)
        return pred_boxes

训练自己的YoloV5模型

注意文件存放的目录

【计算机视觉】:基于PyTorch的YoloV5目标检测平台

; 一、数据集的准备

大佬为大家准备好的数据集链接
VOC数据集下载地址如下,里面已经包括了训练集、测试集、验证集(与测试集一样),无需再次划分:
链接: https://pan.baidu.com/s/19Mw2u_df_nBzsC2lg20fQA
提取码: j5ge

训练前将标签文件放在VOCdevkit文件夹下的VOC2007文件夹下的Annotation中。

【计算机视觉】:基于PyTorch的YoloV5目标检测平台

训练前将图片文件放在VOCdevkit文件夹下的VOC2007文件夹下的JPEGImages中。

【计算机视觉】:基于PyTorch的YoloV5目标检测平台

划分测试集test.txt(用来进行MAP计算,MAP是指标)、训练集train.txt、验证集val,txt、trainval.txt没用到

【计算机视觉】:基于PyTorch的YoloV5目标检测平台

; 二、数据集的处理

在完成数据集的摆放之后,我们需要对数据集进行下一步的处理,目的是获得训练用的2007_train.txt以及2007_val.txt,需要用到根目录下的voc_annotation.py。

voc_annotation.py里面有一些参数需要设置。
分别是annotation_mode、classes_path、trainval_percent、train_percent、VOCdevkit_path,第一次训练可以仅修改classes_path


annotation_mode     = 0

classes_path        = 'model_data/voc_classes.txt'

trainval_percent    = 0.9
train_percent       = 0.9

VOCdevkit_path  = 'VOCdevkit'

classes_path用于指向检测类别所对应的txt,以voc数据集为例,我们用的txt为:

【计算机视觉】:基于PyTorch的YoloV5目标检测平台

三、开始训练网络

通过voc_annotation.py我们已经生成了2007_train.txt以及2007_val.txt,此时我们可以开始训练了。
训练的参数较多,大家可以在下载库后仔细看注释,其中最重要的部分依然是train.py里的classes_path。
classes_path用于指向检测类别所对应的txt,这个txt和voc_annotation.py里面的txt一样!训练自己的数据集必须要修改!

【计算机视觉】:基于PyTorch的YoloV5目标检测平台
修改完classes_path后就可以运行train.py开始训练了,在训练多个epoch后,权值会生成在logs文件夹中。
其它参数的作用如下:

    Cuda            = True

    distributed     = False

    sync_bn         = False

    fp16            = False

    classes_path    = 'model_data/voc_classes.txt'

    anchors_path    = 'model_data/yolo_anchors.txt'
    anchors_mask    = [[6, 7, 8], [3, 4, 5], [0, 1, 2]]

    model_path      = 'model_data/yolov5_s.pth'

    input_shape     = [640, 640]

    backbone        = 'cspdarknet'

    pretrained      = False

    phi             = 's'

    mosaic              = True
    mosaic_prob         = 0.5
    mixup               = True
    mixup_prob          = 0.5
    special_aug_ratio   = 0.7

    label_smoothing     = 0

    Init_Epoch          = 0
    Freeze_Epoch        = 50
    Freeze_batch_size   = 16

    UnFreeze_Epoch      = 200
    Unfreeze_batch_size = 8

    Freeze_Train        = True

    Init_lr             = 1e-2
    Min_lr              = Init_lr * 0.01

    optimizer_type      = "sgd"
    momentum            = 0.937
    weight_decay        = 5e-4

    lr_decay_type       = "cos"

    save_period         = 10

    save_dir            = 'logs'

    eval_flag           = True
    eval_period         = 10

    num_workers         = 4

    train_annotation_path   = '2007_train.txt'
    val_annotation_path     = '2007_val.txt'

四、训练结果预测

训练结果预测需要用到两个文件,分别是yolo.py和predict.py。
我们首先需要去yolo.py里面修改model_path以及classes_path,这两个参数必须要修改。

model_path指向训练好的权值文件,在logs文件夹里。
classes_path指向检测类别所对应的txt。

class YOLO(object):
    _defaults = {

        "model_path"        : 'logs/ep200-loss0.048-val_loss0.048.pth',

        "classes_path"      : 'model_data/voc_classes.txt',

        "anchors_path"      : 'model_data/yolo_anchors.txt',
        "anchors_mask"      : [[6, 7, 8], [3, 4, 5], [0, 1, 2]],

        "input_shape"       : [640, 640],

        "backbone"          : 'cspdarknet',

        "phi"               : 's',

        "confidence"        : 0.5,

        "nms_iou"           : 0.3,

        "letterbox_image"   : True,

        "cuda"              : True,
    }

完成修改后就可以运行predict.py进行检测了。

Original: https://blog.csdn.net/qq_45802659/article/details/125381945
Author: JRWu
Title: 【计算机视觉】:基于PyTorch的YoloV5目标检测平台

原创文章受到原创版权保护。转载请注明出处:https://www.johngo689.com/708765/

转载文章受原作者版权保护。转载请注明原作者出处!

(0)

大家都在看

亲爱的 Coder【最近整理,可免费获取】👉 最新必读书单  | 👏 面试题下载  | 🌎 免费的AI知识星球