yolov3 neck部分搭建，以及完整的网络搭建代码，可直接运行

news2024/9/24 17:17:52

1. 用来实现共同的一些操作 common_module.py

2.1 使用到的通道配置 net_config.py

2.2 主干网络 darknet53.py

3. predict部分 neck.py

4. 拼凑成完整的yolov3网络 Yolov3.py

5. 网络结构如下

上次提到了搭建了主干网络，这个加上neck就可以构建完整的yolov3网络

红色的字体是通道数。当然，你也可以最定义，最后的输出通道数可以根据你的分类数定夺。

下面我们一起来实现一下整个代码。

1. 用来实现共同的一些操作 common_module.py


import torch.nn.functional as F
from torch import nn


class ConvBnLeakRelu(nn.Module):
    """
    基本卷积部分,两个功能
    改变形状：k:3  s:2   p：1
    升降维： k:1    s:1 p:0
    """

    def __init__(self, in_c, out_c, k, s):
        super(ConvBnLeakRelu, self).__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(in_c, out_c, k, s, padding=k // 2, bias=False),
            nn.BatchNorm2d(out_c),
            nn.LeakyReLU(0.1, inplace=True)
        )

    def forward(self, x):
        return self.layers(x)


class ConvBn(nn.Module):
    """
    基本卷积部分,两个功能,不跟激活，用于残差单元
    改变形状：k:3  s:2   p：1
    升降维： k:1    s:1 p:0
    """

    def __init__(self, in_c, out_c, k, s):
        super(ConvBn, self).__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(in_c, out_c, k, s, padding=k // 2, bias=False),
            nn.BatchNorm2d(out_c),
        )

    def forward(self, x):
        return self.layers(x)


class DownSample(nn.Module):
    """
    wh减半，通道增加1倍
    """

    def __init__(self, in_c):
        super(DownSample, self).__init__()
        self.downsample = ConvBnLeakRelu(in_c, in_c * 2, 3, 2)

    def forward(self, x):
        return self.downsample(x)


class ResBlock(nn.Module):
    def __init__(self, in_c):
        super(ResBlock, self).__init__()
        self.layers = nn.Sequential(
            ConvBnLeakRelu(in_c, in_c // 2, 1, 1),
            ConvBn(in_c // 2, in_c, 3, 1)
        )

    def forward(self, x):
        return F.leaky_relu(x + self.layers(x))


class ConvSet(nn.Module):
    # 通道减半，形状不变
    def __init__(self, in_c, out_c):
        super(ConvSet, self).__init__()
        self.layers = nn.Sequential(
            ConvBnLeakRelu(in_c, in_c * 2, 1, 1),
            ConvBnLeakRelu(in_c * 2, in_c * 2, 3, 1),
            ConvBnLeakRelu(in_c * 2, in_c, 1, 1),
            ConvBnLeakRelu(in_c, in_c, 3, 1),
            ConvBnLeakRelu(in_c, out_c, 1, 1),
        )

    def forward(self, x):
        return self.layers(x)


class UpSample(nn.Module):
    # 通道减半，形状扩大一倍
    def __init__(self, in_c):
        super(UpSample, self).__init__()
        self.layers = nn.Sequential(
            ConvBnLeakRelu(in_c, in_c // 2, 1, 1),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        )

    def forward(self, x):
        return self.layers(x)


class PredictHead(nn.Module):
    def __init__(self, in_c, out_c=8*3):
        super(PredictHead, self).__init__()
        self.layers = nn.Sequential(
            ConvBnLeakRelu(in_c, in_c, 3, 1),
            ConvBnLeakRelu(in_c, out_c, 1, 1)
        )
    def forward(self, x):
        return self.layers(x)

if __name__ == '__main__':
    from torchinfo import summary

    # net = DownSample(64)
    # net = ResBlock(64)
    # net = ConvSet(64)
    net = PredictHead(512, 24)
    summary(net, (3, 512, 13, 13))

2.1 使用到的通道配置 net_config.py


config = [
    [32, 1],
    [64, 2],
    [128, 8],
    [256, 8],
    [512, 4],
]

2.2 主干网络 darknet53.py


import torch
from net import net_config
from net.common_module import *


class DarkNet53(nn.Module):
    def __init__(self):
        super(DarkNet53, self).__init__()
        self.input_layer = nn.Sequential(ConvBnLeakRelu(3, 32, 3, 1))

        layers = []
        for in_c, num_block in net_config.config:
            layers += [self.make_layer(in_c, num_block)]
        self.layers = nn.Sequential(*layers)

    def make_layer(self, in_c, block_num):
        out_c = in_c * 2
        layer = [DownSample(in_c)]
        for i in range(block_num):
            layer.append(ResBlock(out_c))
        return nn.Sequential(*layer)

    def forward(self, x):
        # 32 416*416
        x = self.input_layer(x)
        out_52 = self.layers[:3](x)
        out_26 = self.layers[3](out_52)
        out_13 = self.layers[4](out_26)
        return out_52, out_26, out_13


if __name__ == '__main__':
    net = DarkNet53()
    x = torch.rand(1, 3, 416, 416)
    out_52, out_26, out_13 = net(x)
    print(out_52.shape, out_26.shape, out_13.shape, )
    # torch.Size([1, 256, 52, 52]) torch.Size([1, 512, 26, 26]) torch.Size([1, 1024, 13, 13])

    summary(net, (10, 3, 416, 416))

输出的就是这个框里面输出的3条线

3. predict部分 neck.py



import torch
from net.common_module import ConvSet, UpSample, PredictHead
from torch import nn


class Neck(nn.Module):
    def __init__(self):
        super(Neck, self).__init__()
        # 1024→512
        self.conv2dset_13 = ConvSet(1024, 512)
        self.predict_13 = PredictHead(512)

        # 512→256
        self.up26 = UpSample(512)
        self.conv2dset_26 = ConvSet(512 + 256, 256)
        self.predict_26 = PredictHead(256)
        # 256→128
        self.up52 = UpSample(256)
        self.conv2dset_52 = ConvSet(256 + 128, 128)
        self.predict_52 = PredictHead(128)

    def forward(self, out_52, out_26, out_13):
        out_13 = self.conv2dset_13(out_13)
        out_predict_13 = self.predict_13(out_13)

        out_up26 = self.up26(out_13)
        out_26 = self.conv2dset_26(torch.cat((out_up26, out_26), dim=1))
        out_predict_26 = self.predict_26(out_26)

        out_up52 = self.up52(out_26)
        out_52 = self.conv2dset_52(torch.cat((out_up52, out_52), dim=1))
        out_predict_52 = self.predict_52(out_52)

        return out_predict_52, out_predict_26, out_predict_13

4. 拼凑成完整的yolov3网络 Yolov3.py


import torch
from net.darknet53 import DarkNet53
from net.neck import Neck
from torch import nn


class YoloV3(nn.Module):
    def __init__(self):
        super(YoloV3, self).__init__()
        self.backone = DarkNet53()
        self.neck = Neck()

    def forward(self, x):
        out_52, out_26, out_13 = self.backone(x)
        out_predict_52, out_predict_26, out_predict_13 = self.neck(out_52, out_26, out_13)
        return out_predict_13, out_predict_26, out_predict_52


if __name__ == '__main__':
    net = YoloV3()
    x = torch.rand(1, 3, 416, 416)
    out_13, out_26, out_52 = net(x)

    print(out_52.shape, out_26.shape, out_13.shape)

    # summary(net, (10, 3, 416, 416))
    # torch.Size([1, 128, 52, 52]) torch.Size([1, 256, 26, 26]) torch.Size([1, 512, 13, 13])