Darknet网络代码

Darknet网络代码

 import math
 from collections import OrderedDict
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 ​
 class Mish(nn.Module):
     def __init__(self):
         super(Mish, self).__init__()
     def forward(self, x):
         x = x * torch.tanh(F.softmax(x))
     # F.softmax(x) = ln(1+e^x)
     # tanh(x) = (e^x-e^(-x))/(e^x+e^(-x))
         return x
 ​
 # 卷积+归一化+激活
 class CB(nn.Module):
     def __init__(self, in_channels, out_channels, kernel_size, stride=1):
         super(CB, self).__init__()
         self.conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=kernel_size // 2, bias=False)
         self.bn = nn.BatchNorm2d(out_channels)
         self.activation = Mish()
     def forward(self, x):
         x = self.conv(x)
         x = self.bn(x)
         x = self.activation(x)
         return x
 ​
 class Resblock(nn.Module):
     def __init__(self, channels, hidden_channels=None):
         super(Resblock, self).__init__()
         if hidden_channels is None:
             hidden_channels = channels
         self.resblock = nn.Sequential(
             nn.Conv2d(in_channels = channels, out_channels= hidden_channels, kernel_size=1),
             nn.Conv2d(in_channels=hidden_channels, out_channels=channels, kernel_size=3, padding=1)
         )
     def forward(self, x):
 ​
         x = x + self.resblock(x)
         return x
 ​
 class Resblock_body(nn.Module):
     def __init__(self, in_channels, out_channels, num_blocks, first):
         super(Resblock_body, self).__init__()
         self.downsample_conv = CB(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=2)
         if first:
             self.split_conv0 = CB(in_channels=out_channels, out_channels=out_channels, kernel_size=1)
             self.split_conv1 = CB(in_channels=out_channels, out_channels=out_channels, kernel_size=1)
             self.blocks_conv = nn.Sequential(
                 Resblock(channels=out_channels, hidden_channels=out_channels//2),
                 CB(in_channels=out_channels, out_channels=out_channels, kernel_size=1)
             )
             self.concat_conv = CB(in_channels=out_channels * 2, out_channels=out_channels, kernel_size=1)
         else:
             self.split_conv0 = CB(in_channels=out_channels, out_channels=out_channels // 2, kernel_size=1)
             self.split_conv1 = CB(in_channels=out_channels, out_channels=out_channels // 2, kernel_size=1)
             self.blocks_conv = nn.Sequential(
                 *[Resblock(out_channels // 2) for _ in range(num_blocks)],
                 CB(in_channels=out_channels // 2, out_channels=out_channels // 2, kernel_size=1)
             )
             self.concat_conv = CB(in_channels=out_channels, out_channels=out_channels, kernel_size=1)
     def forward(self, x):
         x = self.downsample_conv(x)
         x0 = self.split_conv0(x)
         x1 = self.split_conv1(x)
         x1 = self.blocks_conv(x1)
         x = torch.cat([x0, x1], dim=1)
         x = self.concat_conv(x)
         return x
 class CspDarknet(nn.Module):
     def __init__(self, layers):
         super(CspDarknet, self).__init__()
         self.inplanes = 32
         self.conv1 = CB(in_channels = 3, out_channels=self.inplanes, kernel_size=3, stride=1) 
        self.feature_channels = [64, 128, 256, 512, 1024] 
        self.stages =nn.ModuleList 
        self.stages = nn.ModuleList([ 
            # 416,416,32 -> 208,208,64 
            Resblock_body(in_channels=self.inplanes, out_channels=self.feature_channels[0], num_blocks=layers[0], first=True), 
            # 208,208,64 -> 104,104,128 
            Resblock_body(in_channels=self.feature_channels[0], out_channels=self.feature_channels[1], num_blocks=layers[1], first=False), 
            # 104,104,128 -> 52,52,256 
            Resblock_body(in_channels=self.feature_channels[1], out_channels=self.feature_channels[2], num_blocks=layers[2], first=False), 
            # 52,52,256 -> 26,26,512 
            Resblock_body(in_channels=self.feature_channels[2], out_channels=self.feature_channels[3], num_blocks=layers[3], first=False), 
            # 26,26,512 -> 13,13,1024 
            Resblock_body(in_channels=self.feature_channels[3], out_channels=self.feature_channels[4], num_blocks=layers[4], first=False) 
        ]) 
        self.num_features=1 
        for m in self.modules(): 
            if isinstance(m, nn.Conv2d): 
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 
                m.weight.data.normal_(0, math.sqrt(2. / n)) 
            elif isinstance(m, nn.BatchNorm2d): 
                m.weight.data.fill_(1) 
                m.bias.data.zero_() 
    def forward(self, x): 
        x = self.conv1(x) 
        x = self.stages[0](x) 
        x = self.stages[1](x) 
        out3 = self.stages[2](x) 
        out4 = self.stages[3](out3) 
        out5 = self.stages[4](out4) 
        return out3, out4, out5 
def darknet53(pretrained): 
    model = CspDarknet([1, 2, 8, 8, 4]) 
    if pretrained: 
        model.load_state_dict(torch.load("D:/finishi_project/model_data/CSPdarknet53_backbone_weights.pth")) 
    return model 
​ 
# from torchsummary import summary 
# dar = darknet53(False) 
# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 
# summary(dar, input_size=(3, 416, 416)) 
​ 
​

 

代码没有注释，欢迎留言共同讨论，顺便给个关注，感谢。