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