yolov5骨干代码

以下是YOLOv5的骨干网络代码示例：

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

class CSPDarknet53(nn.Module):
    def __init__(self, num_classes=1000):
        super(CSPDarknet53, self).__init__()
        
        # CSPDarknet53 architecture
        self.conv1 = Conv(3, 32, kernel_size=3, stride=1)
        self.conv2 = Conv(32, 64, kernel_size=3, stride=2)
        self.conv3 = Bottleneck(64, 64)
        self.conv4 = Conv(64, 128, kernel_size=3, stride=2)
        self.conv5 = nn.Sequential(
            BottleneckCSP(128, 128, n=2),
            Conv(128, 256, kernel_size=3, stride=2)
        )
        self.conv6 = nn.Sequential(
            BottleneckCSP(256, 256, n=8),
            Conv(256, 512, kernel_size=3, stride=2)
        )
        self.conv7 = nn.Sequential(
            BottleneckCSP(512, 512, n=8),
            Conv(512, 1024, kernel_size=3, stride=2)
        )
        self.conv8 = nn.Sequential(
            BottleneckCSP(1024, 1024, n=4),
            Conv(1024, 2048, kernel_size=1)
        )

        # Classification layer
        self.fc = nn.Linear(2048, num_classes)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)
        x = self.conv5(x)
        x = self.conv6(x)
        x = self.conv7(x)
        x = self.conv8(x)
        x = F.adaptive_avg_pool2d(x, (1, 1))
        x = torch.flatten(x, 1)
        x = self.fc(x)

        return x

class Conv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1):
        super(Conv, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding=kernel_size//2, bias=False)
        self.bn = nn.BatchNorm2d(out_channels)
        self.act = nn.LeakyReLU(0.1)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.act(x)

        return x

class Bottleneck(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(Bottleneck, self).__init__()
        hidden_channels = out_channels // 2
        self.conv1 = Conv(in_channels, hidden_channels, kernel_size=1)
        self.conv2 = Conv(hidden_channels, out_channels, kernel_size=3)

    def forward(self, x):
        x1 = self.conv1(x)
        x2 = self.conv2(x1)
        out = x + x2

        return out

class BottleneckCSP(nn.Module):
    def __init__(self, in_channels, out_channels, n=1):
        super(BottleneckCSP, self).__init__()
        hidden_channels = out_channels // 2
        self.conv1 = Conv(in_channels, hidden_channels, kernel_size=1)
        self.conv2 = Conv(in_channels, hidden_channels, kernel_size=1)
        self.conv3 = Conv(hidden_channels, hidden_channels, kernel_size=1)
        self.conv4 = Conv(hidden_channels, hidden_channels, kernel_size=1)
        self.conv5 = Conv(hidden_channels * 2, out_channels, kernel_size=1)
        self.convs = nn.ModuleList([self.conv2, self.conv3] * n)
        self.bn = nn.BatchNorm2d(out_channels)
        self.act = nn.LeakyReLU(0.1)

    def forward(self, x):
        x1 = self.conv1(x)

        x2 = self.conv2(x1)
        for conv in self.convs:
            x2 = conv(x2)
        x3 = self.conv3(x1)
        out = torch.cat((x2, x3), dim=1)
        out = self.conv4(out)
        out = self.conv5(out)

        out = self.bn(out)
        out = self.act(out)

        return out

希望对你有所帮助！