YoloV3网络模型搭建

结构图是复制别人的

 

import torch
import torch.nn as nn
from collections import OrderedDict


class CBL(nn.Module):
    def __init__(self, channel_in, channel_out, ks, p=1, strides=(1, 1)):
        super(CBL, self).__init__()
        self.block = nn.Sequential(
            nn.Conv2d(channel_in, channel_out, ks, padding=p),
            nn.BatchNorm2d(channel_out),
            nn.LeakyReLU(0.1)
        )
    def forward(self, x):
        return self.block(x)


class ResidualBlock(nn.Module):
    # 利用一个1x1卷积下降通道数===>利用一个3x3卷积提取特征===>利用一个1x1卷积上升通道数
    def __init__(self, inp, planes):
        """这个类需要传进来两个参数，一个是【input_channel,output_channel】"""
        super(ResidualBlock, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(inp, planes[0], kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(planes[0]),
            nn.LeakyReLU(0.1))
        self.conv2 = nn.Sequential(
            nn.Conv2d(planes[0], planes[1], kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(planes[1]),
            nn.LeakyReLU(0.1))
        self.conv3 = nn.Sequential(
            nn.Conv2d(planes[1], planes[1], kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(planes[1]),
            nn.LeakyReLU(0.1))

    def forward(self, inputs):
        x = self.conv1(inputs)
        x = self.conv2(x)
        x = self.conv3(x)
        return (x + inputs)


class DarkNet(nn.Module):
    def __init__(self, blocks):
        super(DarkNet, self).__init__()
        self.inp = 32
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=self.inp, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(self.inp),
            nn.LeakyReLU(0.1))
        self.block1 = nn.Sequential(self._make_residual_layer(planes=[32, 64], block=blocks[0]),
                                    self._make_residual_layer(planes=[64, 128], block=blocks[1]),
                                    self._make_residual_layer(planes=[128, 256], block=blocks[2]))
        self.block2 = self._make_residual_layer(planes=[256, 512], block=blocks[3])
        self.block3 = self._make_residual_layer(planes=[512, 1024], block=blocks[4])

    def forward(self, inputs):
        #########################
        # 1.Darknet部分
        #########################
        x = self.conv1(inputs)
        x = self.block1(x)  # 包括前三个res
        feat1 = x
        x = self.block2(x)
        feat2 = x
        x = self.block3(x)
        feat3 = x
        return feat1, feat2, feat3

    def _make_residual_layer(self, planes, block):
        # planes是【input_channel,output_channel】，blocks是重复执行残差的次数
        layers = []

        for i in range(block):
            if i < 1:       # 第一次由于channel不一样，所以需要对通道数进行调整
                layers.append(
                    ("ds_conv", nn.Conv2d(self.inp, planes[1], kernel_size=3, stride=2, padding=1, bias=False)))
                layers.append(("ds_bn", nn.BatchNorm2d(planes[1])))
                layers.append(("ds_relu", nn.LeakyReLU(0.1)))
                self.inp = planes[1]
            else:
                layers.append(("residual_{}".format(i), ResidualBlock(self.inp, planes)))
        return nn.Sequential(OrderedDict(layers))


class YoloBody(nn.Module):
    def __init__(self, num_classes):
        super(YoloBody, self).__init__()
        self.darknet = DarkNet([1, 2, 8, 8, 4])
        self.num_classes = num_classes
        self.up1 = nn.Sequential(
            CBL(512, 256, 1, 0),
            nn.UpsamplingBilinear2d(scale_factor=2)
        )
        self.up2 = nn.Sequential(
            CBL(256, 128, 1, 0),
            nn.UpsamplingBilinear2d(scale_factor=2)
        )

    def forward(self, inputs):
        feat1, feat2, feat3 = self.darknet(inputs)
        # print(feat1.shape)
        #########################
        # 2.加强特征提取
        #########################
        big_map, x = self.make_five_conv(feat3, 1024, 512)
        big_output = self.output_conv(big_map, 512, (self.num_classes+5)*3)

        x = self.up1(x)
        x = torch.cat([feat2, x],dim=1)
        middel_map, x = self.make_five_conv(x, 768, 256)
        middel_output = self.output_conv(middel_map, 256, (self.num_classes+5)*3)

        x = self.up2(x)
        x = torch.cat([feat1, x], dim=1)
        small_map, _ = self.make_five_conv(x, 384, 128)
        small_output = self.output_conv(small_map, 128, (self.num_classes+5)*3)
        print(big_output.shape, middel_output.shape, small_output.shape)

        return (big_output, middel_output, small_output)

    def make_five_conv(self, x, channel_in, channel_out):
        x = CBL(channel_in, channel_out, ks=1, p=0)(x)
        x = CBL(channel_out, channel_out, ks=3, p=1)(x)
        x = CBL(channel_out, channel_out, ks=1, p=0)(x)
        x = CBL(channel_out, channel_out, ks=3, p=1)(x)
        x = CBL(channel_out, channel_out, ks=1, p=0)(x)
        return x,x

    def output_conv(self, x, channel_in, channel_out):
        x = CBL(channel_in, channel_out, ks=3)(x)
        y = nn.Conv2d(channel_out, channel_out, kernel_size=1, padding=0)(x)
        return y





if __name__ == '__main__':
    inputs = torch.zeros(size=(1,3, 416, 416))
    print(inputs.shape)
    # model = DarkNet(blocks=[1, 2, 8, 8, 4])
    model = YoloBody(num_classes=20)
    out = model.forward(inputs)