# Codeblock 4a
class Backbone(nn.Module):
    def __init__(self):
        super().__init__()
        # in_channels, out_channels, kernel_size, stride, padding, maxpool_flag
        self.stage0 = ConvBlock(3, 64, 7, 2, 3, maxpool_flag=True)      #(1)
        self.stage1 = ConvBlock(64, 192, 3, 1, 1, maxpool_flag=True)    #(2)
        
        self.stage2 = nn.ModuleList([
            ConvBlock(192, 128, 1, 1, 0), 
            ConvBlock(128, 256, 3, 1, 1), 
            ConvBlock(256, 256, 1, 1, 0),
            ConvBlock(256, 512, 3, 1, 1, maxpool_flag=True)      #(3)
        ])
        
        
        self.stage3 = nn.ModuleList([])
        for _ in range(4):
            self.stage3.append(ConvBlock(512, 256, 1, 1, 0))
            self.stage3.append(ConvBlock(256, 512, 3, 1, 1))
            
        self.stage3.append(ConvBlock(512, 512, 1, 1, 0))
        self.stage3.append(ConvBlock(512, 1024, 3, 1, 1, maxpool_flag=True))  #(4)
        
        
        self.stage4 = nn.ModuleList([])
        for _ in range(2):
            self.stage4.append(ConvBlock(1024, 512, 1, 1, 0))
            self.stage4.append(ConvBlock(512, 1024, 3, 1, 1))
        
        self.stage4.append(ConvBlock(1024, 1024, 3, 1, 1))
        self.stage4.append(ConvBlock(1024, 1024, 3, 2, 1))    #(5)
        
        
        self.stage5 = nn.ModuleList([])
        self.stage5.append(ConvBlock(1024, 1024, 3, 1, 1))
        self.stage5.append(ConvBlock(1024, 1024, 3, 1, 1))