Darknet53 pytorch版

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

class ConvolutionalLayer(nn.Module):
    def __init__(self,in_channels,out_channels,kernal_size,stride,padding):
        super(ConvolutionalLayer, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels,out_channels,kernal_size,stride,padding),
            nn.BatchNorm2d(out_channels),
            nn.LeakyReLU(0.1)
        )
    def forward(self,x):
        return self.conv(x)

class ResidualLayer(nn.Module):
    def __init__(self,in_channels):
        super(ResidualLayer, self).__init__()
        self.reseblock = nn.Sequential(
            ConvolutionalLayer(in_channels,in_channels // 2,kernal_size=1,stride=1,padding=0),
            ConvolutionalLayer(in_channels // 2,in_channels,kernal_size=3,stride=1,padding=1)
        )

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

class DownSampleLayer(nn.Module):
    def __init__(self,in_channels,out_channels):
        super(DownSampleLayer, self).__init__()
        self.conv = nn.Sequential(
            ConvolutionalLayer(in_channels,out_channels,kernal_size=3,stride=2,padding=1)
        )

    def forward(self,x):
        return self.conv(x)

class UpSampleLayer(nn.Module):
    def __init__(self):
        super(UpSampleLayer, self).__init__()
    def forward(self,x):
        return F.interpolate(x,scale_factor=2,mode='nearest')

class ConvolutionalSetLayer(nn.Module):
    def __init__(self,in_channel,out_channel):
        super(ConvolutionalSetLayer, self).__init__()
        self.conv = nn.Sequential(
            ConvolutionalLayer(in_channel,out_channel,kernal_size=1,stride=1,padding=0),
            ConvolutionalLayer(out_channel,in_channel,kernal_size=3,stride=1,padding=1),
            ConvolutionalLayer(in_channel,out_channel,kernal_size=1,stride=1,padding=0),
            ConvolutionalLayer(out_channel,in_channel,kernal_size=3,stride=1,padding=1),
            ConvolutionalLayer(in_channel,out_channel,kernal_size=1,stride=1,padding=0)
        )

    def forward(self,x):
        return self.conv(x)

class DarkNet53(nn.Module):
    def __init__(self):
        super(DarkNet53, self).__init__()

        self.feature_52 = nn.Sequential(
            ConvolutionalLayer(3,32,3,1,1),
            DownSampleLayer(32,64),

            ResidualLayer(64),

            DownSampleLayer(64,128),

            ResidualLayer(128),
            ResidualLayer(128),

            DownSampleLayer(128,256),

            ResidualLayer(256),
            ResidualLayer(256),
            ResidualLayer(256),
            ResidualLayer(256),
            ResidualLayer(256),
            ResidualLayer(256),
            ResidualLayer(256),
            ResidualLayer(256)
        )

        self.feature_26 = nn.Sequential(
            DownSampleLayer(256,512),

            ResidualLayer(512),
            ResidualLayer(512),
            ResidualLayer(512),
            ResidualLayer(512),
            ResidualLayer(512),
            ResidualLayer(512),
            ResidualLayer(512),
            ResidualLayer(512),
        )

        self.feature_13 = nn.Sequential(
            DownSampleLayer(512,1024),

            ResidualLayer(1024),
            ResidualLayer(1024),
            ResidualLayer(1024),
            ResidualLayer(1024)
        )

        self.convolset_13 = nn.Sequential(
            ConvolutionalSetLayer(1024,512)
        )

        self.convolset_26 = nn.Sequential(
            ConvolutionalSetLayer(768,256)
        )

        self.convolset_52 = nn.Sequential(
            ConvolutionalSetLayer(384,128)
        )

        self.detection_13 = nn.Sequential(
            ConvolutionalLayer(512,1024,3,1,1),
            nn.Conv2d(1024,15,1,1,0)
        )

        self.detection_26 = nn.Sequential(
            ConvolutionalLayer(256,512,3,1,1),
            nn.Conv2d(512,15,1,1,0)
        )

        self.detection_52 = nn.Sequential(
            ConvolutionalLayer(128,256,3,1,1),
            nn.Conv2d(256,15,1,1,0)
        )

        self.up_26 = nn.Sequential(
            ConvolutionalLayer(512,256,1,1,0),
            UpSampleLayer()
        )

        self.up_52 = nn.Sequential(
            ConvolutionalLayer(256,128,1,1,0),
            UpSampleLayer()
        )

    def forward(self,x):
        h_52 = self.feature_52(x)
        h_26 = self.feature_26(h_52)
        h_13 = self.feature_13(h_26)
        conval_13 = self.convolset_13(h_13)
        detection_13 = self.detection_13(conval_13)
        up_26 = self.up_26(conval_13)
        route_26 = torch.cat((up_26,h_26),dim=1)
        conval_26 = self.convolset_26(route_26)
        detection_26 = self.detection_26(conval_26)
        up_52 = self.up_52(conval_26)
        route_52 = torch.cat((up_52,h_52),dim=1)
        conval_52 = self.convolset_52(route_52)
        detection_52 = self.detection_52(conval_52)

        return detection_13,detection_26,detection_52