1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47
| import torch import torch.nn as nn
class Conv(nn.Module): def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0): super().__init__() self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=False) self.bn = nn.BatchNorm2d(out_channels) self.act = nn.LeakyReLU(0.1) def forward(self, x): return self.act(self.bn(self.conv(x)))
class CSPBlock(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() hidden_channels = out_channels // 2 self.conv1 = Conv(in_channels, hidden_channels, 1) self.conv2 = Conv(in_channels, hidden_channels, 1) self.conv3 = Conv(2 * hidden_channels, out_channels, 1) def forward(self, x): x1 = self.conv1(x) x2 = self.conv2(x) x = torch.cat([x1, x2], dim=1) return self.conv3(x)
class YOLOv4(nn.Module): def __init__(self, num_classes=80): super().__init__() self.backbone = nn.Sequential( Conv(3, 32, 3, 1, 1), Conv(32, 64, 3, 2, 1), CSPBlock(64, 64), Conv(64, 128, 3, 2, 1), CSPBlock(128, 128), Conv(128, 256, 3, 2, 1), CSPBlock(256, 256), Conv(256, 512, 3, 2, 1), CSPBlock(512, 512), Conv(512, 1024, 3, 2, 1), CSPBlock(1024, 1024), ) def forward(self, x): return self.backbone(x)
|