CSPDarknet53代码

CSPDarknet53代码

CSPDarknet53网络

import math

#列表保持哈希性质（无序），进行有序排列
from collections import OrderedDict

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.modules.conv import Conv2d

class Mish(nn.Module):
    def __init__(self):
        super(Mish,self).__init__()
        
    def forward(self, x):
        #F.softplus ---> (ln(1 + e^x))
        return x * torch.tanh(F.softplus(x))

#标准卷积块
class BasicConv(nn.Module):
    #构建完一个类后，先写好类中主函数（继承）
    def __init__(self, input,output,kernel_size,stride = 1):
        super(BasicConv,self).__init__()
        #输入通道数；输出通道数；卷积核大小；步长；偏置
        self.conv = nn.Conv2d(input,output,kernel_size,stride,kernel_size//2,bias=False)
        self.bn = nn.BatchNorm2d(output)
        self.activate = Mish()

    def forward(self, x): 
        x = self.conv(x)
        x = self.bn(x)
        x = self.activate(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.block = nn.Sequential(
            BasicConv(channels,hidden_channels,1),
            BasicConv(hidden_channels,channels,3)
        )


    def forward(self,x):
        return x + self.block(x)


class Resblock_body(nn.Module):
    def __init__(self, in_channels, out_channels, num_blocks, first):
        super(Resblock_body, self).__init__()

        self.downsample_conv = BasicConv(in_channels, out_channels, 3, stride=2)

        if first:
            self.split_conv0 = BasicConv(out_channels, out_channels, 1)

            self.split_conv1 = BasicConv(out_channels, out_channels, 1)  
            self.blocks_conv = nn.Sequential(
                Resblock(channels=out_channels, hidden_channels=out_channels//2),
                BasicConv(out_channels, out_channels, 1)
            )

            self.concat_conv = BasicConv(out_channels*2, out_channels, 1)
        else:
            self.split_conv0 = BasicConv(out_channels, out_channels//2, 1)

            self.split_conv1 = BasicConv(out_channels, out_channels//2, 1)
            self.blocks_conv = nn.Sequential(
                *[Resblock(out_channels//2) for _ in range(num_blocks)],
                BasicConv(out_channels//2, out_channels//2, 1)
            )

            self.concat_conv = BasicConv(out_channels, out_channels, 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([x1, x0], 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 = BasicConv(3, self.inplanes, kernel_size=3,stride=1)
        self.feature_channels = [64,128,256,512,1024]

        '''
        self.stages = nn.ModuleList(
            [
                # channels = 32  ===>  channels = 64
                Resblock_body(self.inplanes, self.feature_channels[0], layers[0],first=True),
                # channels = 64  ===>  channels = 128
                Resblock_body(self.feature_channels[0], self.feature_channels[1], layers[1], first=False),
                # channels = 128  ===>  channels = 256
                Resblock_body(self.feature_channels[1], self.feature_channels[2], layers[2], first=False),
                # channels = 256  ===>  channels = 512
                Resblock_body(self.feature_channels[2], self.feature_channels[3], layers[3], first=False),
                # channels = 512  ===>  channels = 1024
                Resblock_body(self.feature_channels[3], self.feature_channels[4], layers[4], first=False)
            ]
        )
        '''
        self.conv2 = Resblock_body(self.inplanes, self.feature_channels[0], layers[0],first=True)
        self.conv3 = Resblock_body(self.feature_channels[0], self.feature_channels[1], layers[1], first=False)
        self.conv4 = Resblock_body(self.feature_channels[1], self.feature_channels[2], layers[2], first=False)
        self.conv5 = Resblock_body(self.feature_channels[2], self.feature_channels[3], layers[3], first=False)
        self.conv6 = Resblock_body(self.feature_channels[3], self.feature_channels[4], layers[4], first=False)

        self.num_feature = 1
        #isinstance用来判断一个对象是否是已知类型，例如a = 3，isinstance(a,int) = true
        #进行处理       
        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)
        out1 = self.stages[2](x)
        out2 = self.stages[3](x)
        out3 = self.stages[4](x)
        '''
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        out1 = self.conv4(x)
        out2 = self.conv5(x)
        out3 = self.conv6(x)
        return out1,out2,out3

#加载预训练模型
def darknet53(pretrained, **kwargs):
    model = CSPDarknet([1, 2, 8, 8, 4])
    if pretrained:
        if isinstance(pretrained, str):
            model.load_state_dict(torch.load(pretrained))
        else:
            raise Exception("darknet request a pretrained path. got [{}]".format(pretrained))
    return model

小知识点：args和kwargs

args的类型是元组turple，kwargs类型是可变关键数的字典

def test_args_kwargs(*args,**kwargs):
    print(type(args))
    print(type(kwargs))
    for v in args:
        print('Optional argument: ', v)
    for k,v in kwargs.items():
        print('Optional argument %s (kwargs): %s' %(k,v))

if __name__ == "__main__":
    args = (1,2,3)
    kwargs =  {"args1":3,"args2":2,"args3":1}
    test_args_kwargs(*args,**kwargs)
    test_args_kwargs(1,2,3,hhh=0)