您的位置:

YoloV5源码解析

一、模型框架概述

YoloV5模型框架采用了改进版本的SPP(Spatial Pyramid Pooling)结构和PANet(Path Aggregation Network)模块,同时采用了Swish激活函数,Padding=Same卷积核和Focus分离卷积等技术。整个模型包括Backbone、Neck和Head三个部分,其中Backbone部分为CSPdarknet53卷积神经网络,Neck部分为PANet模块,Head部分为YOLOv5模型的检测头部分。

二、数据预处理

YoloV5模型的输入数据格式为416x416的RGB图像,输入图像会经过归一化而统一成0~1之间的数值,之后进行图像增强操作,主要包括随机缩放、随机裁剪、随机反转、色彩调整等操作,这些操作能够增加模型的鲁棒性,提高模型的泛化能力。


def pre_process(img, img_size):
    img = letterbox(img, new_shape=img_size)[0]
    img = img[:, :, ::-1].transpose(2, 0, 1).astype(np.float32)
    img /= 255.0
    return img

三、Backbone

YoloV5的Backbone部分采用了CSPdarknet53卷积神经网络结构。CSPdarknet53中的CSP(Block)模块是一种卷积神经网络结构,它采用了残差连接、通道分离和聚合(Bottleneck)的方式,融合了两条不同的特征传递路径,用于深度特征的提取和运算。使用这种方式可以避免特征信息损失,提高模型的拟合能力。


class CSP(nn.Module):
    def __init__(self, in_channels, out_channels, n=1, shortcut=True, activation=True):
        super(CSP, self).__init__()
        
        self.shortcut = shortcut
        self.activation = activation
        self.layers = nn.Sequential()
        
        self.layers.add_module("conv1", nn.Conv2d(in_channels, out_channels // 2, 1, 1, 0, bias=False))
        self.layers.add_module("bn1", nn.BatchNorm2d(out_channels // 2))
        self.layers.add_module("act1", nn.LeakyReLU(0.1))
        for i in range(n):
            self.layers.add_module("conv2_%d" % i, nn.Conv2d(out_channels // 2, out_channels // 2, 3, 1, 1, bias=False))
            self.layers.add_module("bn2_%d" % i, nn.BatchNorm2d(out_channels // 2))
            self.layers.add_module("act2_%d" % i, nn.LeakyReLU(0.1))
        self.layers.add_module("conv3", nn.Conv2d(out_channels // 2, out_channels, 1, 1, 0, bias=False))
        self.layers.add_module("bn3", nn.BatchNorm2d(out_channels))
        if self.activation:
            self.layers.add_module("act3", nn.LeakyReLU(0.1))

    def forward(self, x):
        if self.shortcut:
            x1, x2 = x.chunk(2, dim=1)
            out = self.layers(x2)
            out = torch.cat([out, x1], dim=1)
        else:
            out = self.layers(x)
        return out

四、Neck

YoloV5的Neck部分采用了PANet模块。PANet是一种基于特征金字塔(FPN)和聚合网络的模块,可用于网络特征表达的增强,实现多尺度特征的聚合与利用。PANet模块中采用了一个自上而下及自下而上的聚合过程,该过程能够解决不同尺度特征信息的传递问题,提高模型对小目标的检测率和准确率。


class PANet(nn.Module):
    def __init__(self, channels):
        super(PANet, self).__init__()
        
        self.layers1 = nn.Sequential(
            nn.Conv2d(channels[1], channels[1], kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(channels[1]),
            nn.LeakyReLU(0.1),
            nn.Conv2d(channels[1], channels[0], kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(channels[0]),
            nn.LeakyReLU(0.1)
        )
        self.layers2 = nn.Sequential(
            nn.Conv2d(channels[0], channels[0], kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(channels[0]),
            nn.LeakyReLU(0.1),
            nn.Conv2d(channels[0], channels[0], kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(channels[0]),
            nn.LeakyReLU(0.1)
        )
        self.layers3 = nn.Sequential(
            nn.Conv2d(channels[0], channels[0], kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(channels[0]),
            nn.LeakyReLU(0.1),
            nn.Conv2d(channels[0], channels[0], kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(channels[0]),
            nn.LeakyReLU(0.1)
        )
        self.layers4 = nn.Sequential(
            nn.Conv2d(channels[0], channels[0], kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(channels[0]),
            nn.LeakyReLU(0.1),
            nn.Conv2d(channels[0], channels[0], kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(channels[0]),
            nn.LeakyReLU(0.1)
        )
        self.fuse = nn.Sequential(
            nn.Conv2d(channels[0], channels[0], kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(channels[0]),
            nn.LeakyReLU(0.1)
        )
        self.downsample = nn.MaxPool2d(kernel_size=2, stride=2)

    def forward(self, x):
        out = self.layers1(x[0])
        x2 = self.fuse(torch.cat((out, F.upsample(x[1], scale_factor=2)), dim=1))
        x3 = self.layers2(x2)
        x4 = self.layers3(x3)
        x5 = self.layers4(x4)
        x5 = self.downsample(x5)
        return x2, x3, x4, x5

五、Head

YoloV5的Head部分采用了YOLOv5模型的检测头部分。在YoloV5中,检测头主要包括Anchor定义、预测调整、处理NMS、处理特征图等操作。采用这种方式可以有效地提高模型的准确率和召回率,同时还能够保证目标检测的稳定性和一致性。


class Detection(nn.Module):
    def __init__(self, in_channels, num_classes, anchors):
        super(Detection, self).__init__()
        self.num_anchors = len(anchors)
        self.num_classes = num_classes
        self.in_channels = in_channels

        self.conv = nn.Conv2d(self.in_channels, self.num_anchors * (self.num_classes + 5), kernel_size=1, stride=1, padding=0)

        self.init_conv2d()

    def forward(self, x):
        out = self.conv(x)
        out = out.permute(0, 2, 3, 1)
        return out.reshape(out.shape[0], -1, self.num_classes + 5)

    def init_conv2d(self):
        bias_value = -4.0
        nn.init.normal_(self.conv.weight, std=0.01)
        nn.init.constant_(self.conv.bias, bias_value)

六、总结

综上所述,YoloV5模型采用了CSPdarknet53卷积神经网络、PANet模块等技术,通过Backbone、Neck和Head三部分实现了高效、准确的目标检测。除此之外,模型还采用了图像增强、Swish激活函数、Padding=Same卷积核和Focus分离卷积等技术来提高模型的鲁棒性和泛化能力。