Pytorch的学习——Yolov3的搭建

Yolov3的搭建

整个Yolov3一共经历三次下采样三次上采样，最后一共有三个输出，输出的维度分别为[通道数， 8， 8]， [通道数， 16， 16]， [通道数， 32， 32]

import torch.nn as nn


# yolo的darknet模块
class Darknet(nn.Module):
    def __init__(self, num_in, num_out, ksize, stride=1, padding=1):
        super(Darknet, self).__init__()
        self.conv1 = nn.Conv2d(num_in, num_out, ksize, stride, padding)
        self.bn1 = nn.BatchNorm2d(num_out)
        self.leakyReLU = nn.LeakyReLU(0.1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.leakyReLU(x)
        return x


# 残差模块 同resnet的残差块
class ResidualBlock(nn.Module):
    def __init__(self, num_in, num_out, num_block):
        super(ResidualBlock, self).__init__()
        self.num_block = num_block
        self.dark_conv1 = Darknet(num_in, num_out, ksize=3, stride=2, padding=1)
        self.dark_conv2 = []
        self.convblock = None
        for i in range(self.num_block):
            layers = list([])
            layers.append(Darknet(num_out, num_out//2, ksize=1, stride=1, padding=0))
            layers.append(Darknet(num_out//2, num_out, ksize=3, stride=1, padding=1))
            self.dark_conv2.append(nn.Sequential(*layers))
        self.dark_conv2 = nn.ModuleList(self.dark_conv2)

    def forward(self, x):
        x = self.dark_conv1(x)
        for convblock in self.dark_conv2:
            self.convblock = convblock
            residual = x
            x = self.convblock(x)
            x = x + residual
        return x


# darknet集成模块
class LastLayer(nn.Module):
    def __init__(self, num_in, num_out, num_out2):
        super(LastLayer, self).__init__()
        self.dark_conv1 = Darknet(num_in, num_out, ksize=1, stride=1, padding=0)
        self.dark_conv2 = Darknet(num_out, num_out * 2, ksize=3, stride=1, padding=1)
        self.dark_conv3 = Darknet(num_out * 2, num_out, ksize=1, stride=1, padding=0)
        self.dark_conv4 = Darknet(num_out, num_out * 2, ksize=3, stride=1, padding=1)
        self.dark_conv5 = Darknet(num_out * 2, num_out, ksize=1, stride=1, padding=0)

        self.dark_conv6 = Darknet(num_out, num_out * 2, ksize=3, stride