您的位置:

详解Yolov5s网络结构图

一、网络结构概述

1、Yolov5s是一种基于Deep CNN的目标检测算法,网络结构采用轻量化实现,旨在提高检测速度和准确率。

2、其网络结构可以分为主干网络和检测头两部分,主干网络利用CSP Darknet53,检测头部分则包括多个卷积层和检测层,其中检测层主要实现目标检测的过程。

3、相较于其他目标检测算法,Yolov5s在推理阶段不需要借助Anchor,从而降低了复杂度,提高了速度。

二、主干网络

1、CSP Darknet53作为Yolov5s的主干网络,由极深的CNN网络和空间扭曲层(CSP:Cross Stage Partial connection)组成,设计的主要目的是有效地减少CNN层的计算复杂度和过拟合的问题。

2、空间扭曲层采用预测残差的方式,把输入特征图拆分成两个部分,其中一部分再通过一系列卷积、BN、激活等层的处理后作为输出,另一部分通过卷积层处理后再与输出特征图融合。这种处理方式既可以提升特征的抽象能力,也可以减少训练参数和计算复杂度。

3、整个主干网络采用预训练的方式进行优化,使用了ImageNet数据集,Master Training Set(MTS)和Target Training Set(TTS)。

class CSPDarknet(nn.Module):
    # CSPDarknet结构定义
    def __init__(self, depths, wid_mult=1.0):
        super(CSPDarknet, self).__init__()
        # 定义三个不同深度的卷积层
        depths = [int(x * wid_mult) for x in depths]
        self.base = nn.ModuleList([
            # 初始化一个ConvBNLeakyReLU类,输入通道数为3,输出通道数为32,卷积核大小为3,步长为1
            ConvBNLeakyReLU(3, depths[0], 3, 1),
            # 注意CSPBlock里面也有两个卷积层,不过这里第一个卷积层不计在depths里
            CSPBlock(depths[0], depths[1], n=1),
            CSPBlock(depths[1], depths[2], n=2),
            CSPBlock(depths[2], depths[2], n=8),
            CSPBlock(depths[2], depths[1], n=2, shortcut=False),
        ])
        # 最后再接一个卷积层,通道数为depths[1],输出通道数为depths[2],卷积核大小为1,步长为1
        # 这里由于不需要经过BN和激活函数,所以可以用nn.Conv2d
        self.tip = nn.Conv2d(depths[1], depths[2], 1, 1)

三、检测头

1、检测头部分包括若干个卷积层和检测层,其中检测层主要实现目标检测的过程。检测头中的卷积层主要是进行特征图的尺度调整和特征图的融合,而检测层则定义用于预测类别和边界框的模型。

2、Yolov5的检测层可以进行三个不同尺度的预测,其输出由5个信息组成,分别是中心坐标和长宽,以及类别置信度,其中长和宽采取的是先验框的形式,而检测结果的置信度则是经过softmax处理后的结果。

3、为了减少FPN结构带来的计算延迟,Yolov5采用了SPPnet结构,可以通过不同尺度的池化操作得到不同大小的感受野从而提高检测准确率。

class Detect(nn.Module):
    """
    检测头,负责将特征图传输到预测层,实现目标的检测
    """
    def __init__(self, nc, anchors):
        super(Detect, self).__init__()
        self.anchors = torch.Tensor(anchors)
        # 类别的个数,包含背景类别
        self.nc = nc
        # 利用nn.ModuleList定义多个卷积层,依次是1个卷积层和3个卷积层,其中第2个卷积层采用SPP结构
        self.m = nn.ModuleList(nn.Conv2d(x, (self.nc + 5) * len(self.anchors), 1) for x in [512, 1024, 512])
        self.export = True
        self.half = False

    def forward(self, x):
        z = []
        for i in range(3):
            # 通过nn.Conv2d之后,矩阵的形状为[batch_size, anchor_num*(5+nc), grid_xy, grid_xy]
            # 后续还需要在grid_xy这个维度拆分出anchor_num个通道,在不同尺寸的预测结果中调用
            # hidden.shape => [batch_size, anchor_num*(5+nc), grid_xy, grid_xy]
            # grid => [grid_xy, grid_xy]
            # stride => [img_size/grid_xy, img_size/grid_xy]
            hidden = self.m[i](x)
            grid = hidden.shape[-2:]
            stride = self.img_size // grid[-1]
            # view(-1, anchor_num, 5+nc, grid_xy, grid_xy)
            # permute(0, 1, 3, 4, 2)
            # 这里是我们在上文提到的拆开hidden中的anchor_num个通道,并调整形状和次序
            hidden = hidden.view(hidden.shape[0], len(self.anchors), 5 + self.nc, grid[0], grid[1]).permute(0, 1, 3, 4, 2).contiguous()
            # 对于前两个数,也就是预测框的中心坐标,我们希望从相对网格坐标中心偏移量预测绝对坐标
            # 在 along_width 和 along_height 两个维度,生成?×?个尺度为 [1, 1, ?, ?]的网格,
            # 偏移量先乘26倍作为目标框中心坐标的初始化,与anchor配对后减去自己的偏移量
            aa = self.anchors.clone().view(len(self.anchors), 1, 1, 2).repeat(1, grid[0], grid[1], 1).cuda()
            # 已经sigmoid求过,所以只需要用exp()还原即可
            hidden[..., 0:2] = (hidden[..., 0:2].sigmoid() * 2.0 - 0.5 + aa) * stride
            # 对于宽和高,先用exp()还原,再与ancor相乘得到相对于当前网格左上角的绝对距离
            hidden[..., 2:4] = (hidden[..., 2:4].sigmoid() * 2) ** 2 * aa * stride
            z.append(hidden.view(hidden.shape[0], -1, 5 + self.nc))
        return torch.cat(z, 1).detach()

四、模型训练

1、Yolov5s模型在训练时,采用GIOU损失函数来衡量预测框和真实框之间的差异,同时采用了焦点损失函数来平衡正负样本的数量,从而提高模型的泛化能力。

2、模型的训练过程采用了分步训练的方法,首先只训练主干网络,再训练检测头,最后将整个模型联合训练,从而提高模型的收敛速度和准确率。

model = Model(cfg).to(device)

# Step 1: 只训练主干网络
optimizer = optim.SGD(model.backbone.parameters(), lr=lr0, momentum=momentum, nesterov=True)

# Step 2: 只训练检测头
optimizer = optim.SGD(model.detect.parameters(), lr=lr0, momentum=momentum, nesterov=True)

# Step 3: 模型联合训练
optimizer = optim.SGD(
    [{'params': model.detect.parameters()}, {'params': model.backbone.parameters(), 'lr': lr0 * 0.1}],
    lr=lr0, momentum=momentum, nesterov=True)

五、模型优化

1、为了提高模型的性能和效果,可以考虑采用数据增强的方式来增加训练集的大小,从而提高模型的泛化能力。

2、可以考虑采用小批量随机梯度下降算法(Min-Batch SGD)来训练模型,从而加快收敛速度。

3、可以考虑使用反向梯度裁剪(Gradient Clipping)来避免梯度爆炸问题。

def train(data_loader, model, loss_func, optimizer, epoch):
    # 模型转换为train状态
    model.train()
    for i, (img, target) in enumerate(data_loader):
        # 将数据推送到GPU
        img, target = img.to(device), target.to(device)
        # 进行预测
        loss, _, _ = model(img, target)
        # 反向传播
        loss.backward()
        # 反向梯度裁剪
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0, norm_type=2)
        # 更新模型
        optimizer.step()
        optimizer.zero_grad()
        # 打印日志
        if i % 10 == 0:
            print(f'Epoch[{epoch}], Step[{i}/{len(data_loader)}], Loss: {loss.item()}')