Yolov4网络代码

from collections import OrderedDict

import torch

import torch.nn as nn

from Darknet_53 import darknet53

def conv(in_channels, out_channels, kernel_size, stride=1):

    pad = (kernel_size-1)//2 if kernel_size else 0

    return nn.Sequential(OrderedDict(

        [

            ("conv", nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=pad)),

            ("bn", nn.BatchNorm2d(out_channels)),

            ("relu", nn.LeakyReLU(0.1))

        ]

    ))

class SPP(nn.Module):

    def __init__(self, pool_sizes=[5, 9, 13]):

        super(SPP, self).__init__()

        self.maxpools = nn.ModuleList([nn.MaxPool2d(pool_size, 1, pool_size//2) for pool_size in pool_sizes])

    def forward(self, x):

        features = [maxpool(x) for maxpool in self.maxpools[::-1]]

        features = torch.cat(features + [x], dim=1)

        return features

class Upsample(nn.Module):

    def __init__(self, in_channels, out_channels):

        super(Upsample, self).__init__()

        self.upsample = nn.Sequential(

            conv(in_channels=in_channels, out_channels=out_channels,kernel_size=1),

            nn.Upsample(scale_factor=2, mode="nearest")

        )

    def forward(self, x):

        x = self.upsample(x)

        return x

def conv_three(channels_list, in_channels):

    m = nn.Sequential(

        conv(in_channels=in_channels, out_channels=channels_list[0], kernel_size=1),

        conv(in_channels=channels_list[0], out_channels=channels_list[1], kernel_size=3),

        conv(in_channels=channels_list[1], out_channels=channels_list[0], kernel_size=1)

    )

    return m

def conv_five(channels_list, in_channels):

    m = nn.Sequential(

        conv(in_channels=in_channels, out_channels=channels_list[0], kernel_size=1),

        conv(in_channels=channels_list[0], out_channels=channels_list[1], kernel_size=3),

        conv(in_channels=channels_list[1], out_channels=channels_list[0], kernel_size=1),

        conv(in_channels=channels_list[0], out_channels=channels_list[1], kernel_size=3),

        conv(in_channels=channels_list[1], out_channels=channels_list[0], kernel_size=1)

    )

    return m

def Yolov4_head(channels_list, in_channels):

    m = nn.Sequential(

        conv(in_channels=in_channels, out_channels=channels_list[0], kernel_size=3),

        conv(in_channels=channels_list[0], out_channels=channels_list[1], kernel_size=1)

    )

    return m

class YoloBody(nn.Module):

    def __init__(self, anchors_mask, num_classes, pretrained = False):

        super(YoloBody, self).__init__()

        self.backbone = darknet53(pretrained)

        self.conv1=conv_three(channels_list=[512, 1024], in_channels=1024)

        self.spp = SPP()

        self.conv2=conv_three(channels_list=[512, 1024], in_channels=2048)

        self.upsample1 = Upsample(512, 256)

        self.conv_for_p4 = conv(in_channels=512, out_channels=256, kernel_size=1)

        self.make_five_conv1=conv_five(channels_list=[256, 512], in_channels=512)

        self.upsample2 = Upsample(in_channels=256, out_channels=128)

        self.conv_for_p3=conv(in_channels=256, out_channels=128, kernel_size=1)

        self.make_five_conv2=conv_five(channels_list=[128, 256], in_channels=256)

        # 3*(5+num_classes) = 3*(5+20) = 3*(4+1+20)=75

        self.yolo_head3=Yolov4_head(channels_list= [256, len(anchors_mask[0]) * (5 + num_classes)], in_channels=128)

        self.down_sample1 = conv(in_channels=128, out_channels=256, kernel_size=3, stride=2)

        self.make_five_conv3 = conv_five(channels_list=[256, 512], in_channels=512)

        # 3*(5+num_classes) = 3*(5+20) = 3*(4+1+20)=75

        self.yolo_head2 = Yolov4_head(channels_list=[512, len(anchors_mask[1]) * (5 + num_classes)], in_channels=256)

        self.down_sample2 = conv(in_channels=256, out_channels=512, kernel_size=3, stride=2)

        self.make_five_conv4 = conv_five(channels_list=[512, 1024], in_channels=1024)

        # 3*(5+num_classes)=3*(5+20)=3*(4+1+20)=75

        self.yolo_head1 = Yolov4_head(channels_list=[1024, len(anchors_mask[2]) * (5 + num_classes)], in_channels=512)

    def forward(self, x):

        x2, x1, x0 = self.backbone(x)

        # 13,13,1024 -> 13,13,512 -> 13,13,1024 -> 13,13,512 -> 13,13,2048

        p5 = self.conv1(x0)

        p5 = self.spp(p5)

        # 13,13,2048 -> 13,13,512 -> 13,13,1024 -> 13,13,512

        p5 = self.conv2(p5)

        # 13,13,512 -> 13,13,256 -> 26,26,256

        p5_upsample = self.upsample1(p5)

        # 26,26,512 -> 26,26,256

        p4 = self.conv_for_p4(x1)

        # 26,26,256 + 26,26,256 -> 26,26,512

        p4 = torch.cat([p4, p5_upsample], axis=1)

        # 26,26,512 -> 26,26,256 -> 26,26,512 -> 26,26,256 -> 26,26,512 -> 26,26,256

        p4 = self.make_five_conv1(p4)

        # 26,26,256 -> 26,26,128 -> 52,52,128

        p4_upsample = self.upsample2(p4)

        # 52,52,256 -> 52,52,128

        p3 = self.conv_for_p3(x2)

        p3=torch.cat([p3, p4_upsample], axis=1)

        p3=self.make_five_conv2(p3)

        p3_downsample=self.down_sample1(p3)

        p4=torch.cat([p3_downsample, p4], axis=1)

        p4=self.make_five_conv3(p4)

        p4_downsample=self.down_sample2(p4)

        p5=torch.cat([p4_downsample, p5], axis=1)

        p5=self.make_five_conv4(p5)

        out2=self.yolo_head3(p3)

        out1=self.yolo_head2(p4)

        out0=self.yolo_head1(p5)

        return out0, out1, out2

# from torchsummary import summary

# yoloyolo=YoloBody(anchors_mask=["0","0","0"], num_classes=5, pretrained = False)

# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# summary(yoloyolo, input_size=(3, 416, 416))

# print(yoloyolo)
代码没有注释，欢迎留言共同讨论，顺便给个关注，感谢。
巴特西

YOLOV4网络

Yolov4网络代码

最新文章

热门文章
