YOLO V8网络（Pytorch）

整体网络结构图（来自网络）：

1. 导入

# ---------- YOLO V8 ----------
import torch
import torch.nn as nn

2. ConvBlock

三部分组成：Conv2d + BatchNorm2d + SiLU

class ConvBlock(nn.Module):
    def __init__(self, input_channel, output_channel, kernel_size, stride, padding):
        super(ConvBlock, self).__init__()
        self.conv = nn.Conv2d(input_channel, output_channel, kernel_size=kernel_size, stride=stride, padding=padding, bias=False)
        self.bn = nn.BatchNorm2d(output_channel)
        self.silu = nn.SiLU(inplace=True)

    def forward(self, x):
        out = self.silu(self.bn(self.conv(x)))
        return out

3. BottleNeck

CSPBlock中的一部分

class BottleNeck(nn.Module):
    def __init__(self, ConvBlock, input_channel, output_channel, add=True):
        super(BottleNeck, self).__init__()
        self.ConvBlock1 = ConvBlock(input_channel, output_channel//2, 3, 1, 1)
        self.ConvBlock2 = ConvBlock(output_channel//2, output_channel, 3, 1, 1)
        self.add = add

    def forward(self, x):
        residual = x
        x = self.ConvBlock1(x)
        out = self.ConvBlock2(x)
        if self.add:
            out = out + residual
        return out

4. CSPBlock

class CSPBlock(nn.Module):
    def __init__(self, ConvBlock, BottleNeck, input_channel, output_channel, block_num, add):
        super(CSPBlock, self).__init__()
        self.ConvBlock1 = ConvBlock(input_channel, output_channel, 1, 1, 0)
        self.BottleNeck = BottleNeck(ConvBlock, output_channel//2, output_channel//2, add)
        self.ConvBlock2 = ConvBlock((block_num + 2)*(output_channel//2), output_channel, 1, 1, 0)
        self.block_num = block_num

    def forward(self, x):
        x = self.ConvBlock1(x)
        y = list(x.chunk(2, 1))
        for _ in range(self.block_num):
            y.append(self.BottleNeck(y[-1]))
        x = torch.cat(y, 1)
        out = self.ConvBlock2(x)
        return out

5. SPPF

class SPPF(nn.Module):
    def __init__(self, ConvBlock, input_channel, output_channel, k=5):
        super(SPPF, self).__init__()
        self.ConvBlock1 = ConvBlock(input_channel, input_channel//2, 1, 1, 0)
        self.maxpooling = nn.MaxPool2d(kernel_size=k, stride=1, padding=k//2)
        self.ConvBlock2 = ConvBlock(input_channel//2*4, output_channel, 1, 1, 0)

    def forward(self, x):
        x = self.ConvBlock1(x)
        y1 = self.maxpooling(x)
        y2 = self.maxpooling(y1)
        y3 = self.maxpooling(y2)
        out = torch.cat((x, y1, y2, y3), 1)
        out = self.ConvBlock2(out)
        return out

6. BackBone

class Net_BackBone(nn.Module):
    def __init__(self, ConvBlock, BottleNeck, CSPBlock, SPPF):
        super(Net_BackBone, self).__init__()
        self.Layer1 = ConvBlock(3, 64, 3, 2, 1)   
        self.Layer2 = nn.Sequential(
            ConvBlock(64, 128, 3, 2, 1),
            CSPBlock(ConvBlock, BottleNeck, 128, 128, 3, True)
        )
        self.Layer3 = nn.Sequential(
            ConvBlock(128, 256, 3, 2, 1),
            CSPBlock(ConvBlock, BottleNeck, 256, 256, 6, True)
        )
        self.Layer4 = nn.Sequential(
            ConvBlock(256, 512, 3, 2, 1),
            CSPBlock(ConvBlock, BottleNeck, 512, 512, 6, True)
        )
        self.Layer5 = nn.Sequential(
            ConvBlock(512, 512, 3, 2, 1),
            CSPBlock(ConvBlock, BottleNeck, 512, 512, 3, True),
            SPPF(ConvBlock, 512, 512, 5)
        )
 
    def forward(self, x):  # 3*640*640
        x = self.Layer1(x)  # 64*320*320
        x = self.Layer2(x)  # 128*160*160
        x = self.Layer3(x)  # 256*80*80
        feature1 = x
        x = self.Layer4(x)  # 512*40*40
        feature2 = x
        x = self.Layer5(x)  # 512*20*20
        feature3 = x
        return feature1, feature2, feature3

7. Neck

Backbone和Head间的金字塔结构，用来加强特征提取

class Net_Neck(nn.Module):
    def __init__(self, Net_BackBone, ConvBlock, BottleNeck, CSPBlock, SPPF):
        super(Net_Neck, self).__init__()
        self.Backbone = Net_BackBone(ConvBlock, BottleNeck, CSPBlock, SPPF)
        self.CSPBlock_3_2 = CSPBlock(ConvBlock, BottleNeck, 1024, 512, 3, add=False)
        self.CSPBlock_2_1 = CSPBlock(ConvBlock, BottleNeck, 768, 256, 3, add=False)
        self.CSPBlock_1_2 = CSPBlock(ConvBlock, BottleNeck, 768, 512, 3, add=False)
        self.CSPBlock_2_3 = CSPBlock(ConvBlock, BottleNeck, 1024, 512, 3, add=False)
        self.Upsample = nn.Upsample(scale_factor=2, mode="nearest")
        self.Down_sample1 = ConvBlock(256, 256, 3, 2, 1)
        self.Down_sample2 = ConvBlock(512, 512, 3, 2, 1)

    def forward(self, x):
        feat1, feat2, feat3 = self.Backbone(x)  # 分别是256*80*80，512*40*40，512*20*20

        feat3_up = self.Upsample(feat3)
        feat3_up_feat2 = torch.cat([feat3_up, feat2], 1)
        feat3_up_feat2 = self.CSPBlock_3_2(feat3_up_feat2)

        feat2_up = self.Upsample(feat3_up_feat2)
        feat2_up_feat1 = torch.cat([feat2_up, feat1], 1)
        feat2_up_feat1 = self.CSPBlock_2_1(feat2_up_feat1)  # 256*80*80

        feat1_down = self.Down_sample1(feat2_up_feat1)
        feat1_down_feat2 = torch.cat([feat1_down, feat3_up_feat2], 1)
        feat1_down_feat2 = self.CSPBlock_1_2(feat1_down_feat2)  # 512*40*40

        feat2_down = self.Down_sample2(feat1_down_feat2)
        feat2_down_feat3 = torch.cat([feat2_down, feat3], 1)
        feat2_down_feat3 = self.CSPBlock_2_3(feat2_down_feat3)  # 512*20*20

        return feat2_up_feat1, feat1_down_feat2, feat2_down_feat3

8. Head

还没有写出来

9. 测试

model=Net_Neck(Net_BackBone, ConvBlock, BottleNeck, CSPBlock, SPPF)
x=torch.randn(16, 3, 640, 640)
fe1,fe2,fe3=model(x)
print(fe1.size())
print(fe2.size())
print(fe3.size())