def __init__(self, in_channels, out_channels, repeat=0):
super().__init__()
self._out_channels = out_channels
self.conv0 = ConvBlock(in_channels, out_channels, size=3, stride=2)
repeat += 1
self.tail = []
self.adds = []
for i in range(repeat):
name = f'seq{i}'
name_add = f'add{i}'
setattr(
self, name,
nn.Sequential(
ConvBlock(out_channels,
out_channels // 2,
size=1,
stride=1),
ConvBlock(out_channels // 2,
out_channels,
size=3,
stride=1)))
setattr(self, name_add, Add())
self.tail.append(getattr(self, name))
self.adds.append(getattr(self, name_add))
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# Darknet Encoder
# ======================================================================
self.features = EfficientEncoder(in_channels=self.in_channels)
f_out1, f_out2 = self.features.out_channels
self.sequence_branch1_1 = nn.Sequential()
self.sequence_branch1_1.add_module(
'branch1_conv1',
ConvBlock(f_out2, max(8, 128 // self.kernels_divider), size=1))
self.sequence_branch1_1.add_module('branch1_upsample', Upsample(2))
self.sequence_branch1_2 = nn.Sequential()
self.sequence_branch1_2.add_module('branch1_concat', Concat(1))
self.sequence_branch1_2.add_module(
'branch1_conv2',
ConvBlock(f_out1 + self.sequence_branch1_1[0].out_channels,
max(8, 128 // self.kernels_divider)))
self.sequence_branch1_2.add_module(
'branch1_conv3',
nn.Conv2d(self.sequence_branch1_2[-1].out_channels,
self.yolo_layer_input_size,
kernel_size=1))
self.sequence_branch2 = nn.Sequential()
self.sequence_branch2.add_module(
'branch2_conv1',
ConvBlock(f_out2, max(8, 128 // self.kernels_divider)))
self.sequence_branch2.add_module(
'branch2_conv2',
nn.Conv2d(self.sequence_branch2[-1].out_channels,
self.yolo_layer_input_size,
kernel_size=1))
# ======================================================================
# YOLO Layers
# ======================================================================
self.yolo1, self.yolo2 = self._create_yolo_layers()
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
kd = self.kernels_divider
self.down1 = ConvPoolBlock(self.in_channels, 16 // kd)
self.down2 = ConvPoolBlock(self.down1.out_channels, 32 // kd)
self.down3 = ConvPoolBlock(self.down2.out_channels, 64 // kd)
self.down4 = Down(self.down3.out_channels, 128 // kd)
self.down5 = Down(self.down4.out_channels, 256 // kd)
self.down = [
self.down1, self.down2, self.down3, self.down4, self.down5
]
self.seq = nn.Sequential()
self.seq.add_module(
'conv1', ConvBlock(self.down[-1].out_channels, 512 // kd, 1))
self.seq_y1 = nn.Sequential()
self.seq_y1.add_module(
'conv1', ConvBlock(self.seq[-1].out_channels, 1024 // kd, 3))
self.seq_y1.add_module(
'conv2',
nn.Conv2d(self.seq_y1[-1].out_channels, self.yolo_layer_input_size,
1, 1))
self.seqy2_1 = nn.Sequential()
self.seqy2_1.add_module(
'conv', ConvBlock(self.seq[-1].out_channels, 256 // kd, 1))
self.seqy2_1.add_module('up', Upsample(2))
self.seqy2_2 = nn.Sequential()
self.seqy2_2.add_module('concat', Concat(1))
self.seqy2_2.add_module(
'conv1',
Conv(self.seqy2_1[0].out_channels + self.down4.out_channels,
512 // kd))
self.seqy2_2.add_module(
'conv2', ConvBlock(self.seqy2_2[-1].out_channels, 256 // kd, 1))
self.seqy2_3 = nn.Sequential()
self.seqy2_3.add_module(
'conv6', ConvBlock(self.seqy2_2[-1].out_channels, 512 // kd))
self.seqy2_3.add_module(
'conv7',
nn.Conv2d(self.seqy2_3[-1].out_channels,
self.yolo_layer_input_size, 1, 1))
self.seqy3_1 = nn.Sequential()
self.seqy3_1.add_module(
'conv', ConvBlock(self.seqy2_2[-1].out_channels, 128 // kd, 1))
self.seqy3_1.add_module('up', Upsample(2))
self.seqy3_2 = nn.Sequential()
self.seqy3_2.add_module('concat', Concat(1))
self.seqy3_2.add_module(
'conv1',
Conv(self.seqy3_1[0].out_channels + self.down3.out_channels,
256 // kd))
self.seqy3_2.add_module(
'conv2',
ConvBlock(self.seqy3_2[-1].out_channels,
self.yolo_layer_input_size))
self.yolo1, self.yolo2, self.yolo3 = self._create_yolo_layers()
def __init__(self, in_channels, out_channels):
super().__init__()
self.conv = ConvBlock(in_channels, out_channels // 2, 1)
self.conv_pool = ConvPoolBlock(out_channels // 2, out_channels)
self.out_channels = out_channels
def __init__(self, in_shape, out_channels):
super().__init__()
self.conv1 = ConvBlock(in_shape, out_channels // 2, 1)
self.conv2 = ConvBlock(out_channels // 2, out_channels, 3)
self.out_channels = out_channels
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# Darknet Encoder
# ======================================================================
self.sequence_1 = nn.Sequential()
self.sequence_1.add_module(
'conv1',
ConvPoolBlock(self.in_channels, max(8,
16 // self.kernels_divider)))
self.sequence_1.add_module(
'conv2',
ConvPoolBlock(self.sequence_1[-1].out_channels,
max(8, 32 // self.kernels_divider)))
self.sequence_1.add_module(
'conv3',
ConvPoolBlock(self.sequence_1[-1].out_channels,
max(8, 64 // self.kernels_divider)))
self.sequence_1.add_module(
'conv4',
ConvPoolBlock(self.sequence_1[-1].out_channels,
max(8, 128 // self.kernels_divider)))
self.sequence_1.add_module(
'conv5',
ConvBlock(self.sequence_1[-1].out_channels,
max(8, 256 // self.kernels_divider)))
self.sequence_2 = nn.Sequential()
self.sequence_2.add_module('max_pool5', nn.MaxPool2d(2, 2))
self.sequence_2.add_module(
'conv6',
ConvPoolBlock(self.sequence_1[-1].out_channels,
max(8, 512 // self.kernels_divider),
pool_stride=1))
self.sequence_2.add_module(
'conv7',
ConvBlock(self.sequence_2[-1].out_channels,
max(8, 1024 // self.kernels_divider)))
self.sequence_2.add_module(
'conv8',
ConvBlock(self.sequence_2[-1].out_channels,
max(8, 256 // self.kernels_divider),
size=1))
self.sequence_branch1_1 = nn.Sequential()
self.sequence_branch1_1.add_module(
'branch1_conv1',
ConvBlock(self.sequence_2[-1].out_channels,
max(8, 128 // self.kernels_divider),
size=1))
self.sequence_branch1_1.add_module('branch1_upsample', Upsample(2))
self.sequence_branch1_2 = nn.Sequential()
self.sequence_branch1_2.add_module('branch1_concat', Concat(1))
self.sequence_branch1_2.add_module(
'branch1_conv2',
ConvBlock(
self.sequence_1[-1].out_channels +
self.sequence_branch1_1[0].out_channels,
max(8, 256 // self.kernels_divider)))
self.sequence_branch1_2.add_module(
'branch1_conv3',
nn.Conv2d(self.sequence_branch1_2[-1].out_channels,
self.yolo_layer_input_size,
kernel_size=1))
self.sequence_branch2 = nn.Sequential()
self.sequence_branch2.add_module(
'branch2_conv1',
ConvBlock(self.sequence_2[-1].out_channels,
max(8, 512 // self.kernels_divider)))
self.sequence_branch2.add_module(
'branch2_conv2',
nn.Conv2d(self.sequence_branch2[-1].out_channels,
self.yolo_layer_input_size,
kernel_size=1))
# ======================================================================
# YOLO Layers
# ======================================================================
self.yolo1, self.yolo2 = self._create_yolo_layers()
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
kd = self.kernels_divider
# Darknet Encoder
# ======================================================================
self.conv1 = ConvBlock(self.in_channels, 32 // kd, stride=1, size=3)
self.down1 = DownSample(self.conv1.out_channels, 64 // kd, repeat=0)
self.down2 = DownSample(self.down1.out_channels, 128 // kd, repeat=1)
self.down3 = DownSample(self.down2.out_channels, 256 // kd, repeat=7)
self.down4 = DownSample(self.down3.out_channels, 512 // kd, repeat=7)
self.down5 = DownSample(self.down4.out_channels, 1024 // kd, repeat=3)
self.down = [
self.down1, self.down2, self.down3, self.down4, self.down5
]
self.sequence_spp = nn.Sequential()
self.sequence_spp.add_module(
'conv1',
ConvBlock(self.down[-1].out_channels, 512 // kd, size=1, stride=1))
self.sequence_spp.add_module(
'conv2',
ConvBlock(self.sequence_spp[-1].out_channels,
1024 // kd,
size=3,
stride=1))
self.sequence_spp.add_module(
'conv3',
ConvBlock(self.sequence_spp[-1].out_channels,
512 // kd,
size=1,
stride=1))
self.spp1 = MaxPool(size=5, stride=1)
self.spp2 = MaxPool(size=9, stride=1)
self.spp3 = MaxPool(size=13, stride=1)
self.branch1_1 = nn.Sequential()
self.branch1_1.add_module('concat', Concat(1))
self.branch1_1.add_module(
'conv1',
ConvBlock(self.sequence_spp[-1].out_channels * 4,
512 // kd,
size=1,
stride=1))
self.branch1_1.add_module(
'conv2',
ConvBlock(self.branch1_1[-1].out_channels,
1024 // kd,
size=3,
stride=1))
self.branch1_1.add_module(
'conv3',
ConvBlock(self.branch1_1[-1].out_channels,
512 // kd,
size=1,
stride=1))
self.branch1_2 = nn.Sequential()
self.branch1_2.add_module(
'conv1',
ConvBlock(self.branch1_1[-1].out_channels,
1024 // kd,
size=3,
stride=1))
self.branch1_2.add_module(
'conv2',
ConvBlock(self.branch1_2[-1].out_channels,
self.yolo_layer_input_size,
size=1,
stride=1))
self.branch2_1 = nn.Sequential(
ConvBlock(self.branch1_1[-1].out_channels,
256 // kd,
size=1,
stride=1), Upsample(2))
self.branch2_2 = nn.Sequential()
self.branch2_2.add_module('concat', Concat(1))
self.branch2_2.add_module(
'conv1',
ConvBlock(self.branch2_1[0].out_channels +
self.down[3].out_channels,
256 // kd,
size=1,
stride=1))
self.branch2_2.add_module(
'conv2',
ConvBlock(self.branch2_2[-1].out_channels,
512 // kd,
size=3,
stride=1))
self.branch2_2.add_module(
'conv3',
ConvBlock(self.branch2_2[-1].out_channels,
256 // kd,
size=1,
stride=1))
self.branch2_2.add_module(
'conv4',
ConvBlock(self.branch2_2[-1].out_channels,
512 // kd,
size=3,
stride=1))
self.branch2_2.add_module(
'conv5',
ConvBlock(self.branch2_2[-1].out_channels,
256 // kd,
size=1,
stride=1))
self.branch2_3 = nn.Sequential()
self.branch2_3.add_module(
'conv6',
ConvBlock(self.branch2_2[-1].out_channels,
512 // kd,
size=3,
stride=1))
self.branch2_3.add_module(
'conv7',
ConvBlock(self.branch2_3[-1].out_channels,
self.yolo_layer_input_size,
size=1,
stride=1))
self.branch3_1 = nn.Sequential(
ConvBlock(self.branch2_2[-1].out_channels,
128 // kd,
size=1,
stride=1), Upsample(2))
self.branch3_2 = nn.Sequential()
self.branch3_2.add_module('concat', Concat(1))
self.branch3_2.add_module(
'conv1',
ConvBlock(self.branch3_1[0].out_channels +
self.down[2].out_channels,
128 // kd,
size=1,
stride=1))
self.branch3_2.add_module(
'conv2',
ConvBlock(self.branch3_2[-1].out_channels,
256 // kd,
size=3,
stride=1))
self.branch3_2.add_module(
'conv3',
ConvBlock(self.branch3_2[-1].out_channels,
128 // kd,
size=1,
stride=1))
self.branch3_2.add_module(
'conv4',
ConvBlock(self.branch3_2[-1].out_channels,
256 // kd,
size=3,
stride=1))
self.branch3_2.add_module(
'conv5',
ConvBlock(self.branch3_2[-1].out_channels,
128 // kd,
size=1,
stride=1))
self.branch3_2.add_module(
'conv6',
ConvBlock(self.branch3_2[-1].out_channels,
256 // kd,
size=3,
stride=1))
self.branch3_2.add_module(
'conv7',
ConvBlock(self.branch3_2[-1].out_channels,
self.yolo_layer_input_size,
size=1,
stride=1))
# ======================================================================
# YOLO Layers
# ======================================================================
self.yolo1, self.yolo2, self.yolo3 = self._create_yolo_layers()