Self-Implemented ResNet Structure with pytorch. ( Structure configuration adaptable )

ResNetBlock.py

import torch
import torch.nn as nn 
import torch.nn.functional as F 

# apply 函数针对net.children()， apply的参数为一个函数，该函数的输入是 Module
def weights_init(m):
    class_name = m.__class__.__name__
    if class_name.find("conv") != -1:
        m.weight.data.normal_(0,0.02)
    elif class_name.find("norm") !=  -1:
        m.weight.data.normal_(0,0.02)
        m.bias.data.fill_(0)

def ConvLayer(in_channels, out_channels, kernel_size, stride, padding = 0, bias = False, 
        norm_layer = None, no_activation = False, *args, **kws):
    if norm_layer is None:
        norm_layer = nn.BatchNorm2d
    blocks = [
        nn.Conv2d(in_channels, out_channels, kernel_size, 
            stride=stride, padding=padding, bias=bias),
        norm_layer(out_channels),
    ]
    if not no_activation:
        blocks.append(nn.ReLU(True))
    return blocks


def PreConvLayer(in_channels, out_channels, kernel_size, stride, padding = 0, bias = False, 
        norm_layer = None, *args, **kws):
    """ Conv Blocks of pre-activation structure
    """
    if norm_layer is None:
        norm_layer = nn.BatchNorm2d
    blocks = [
        norm_layer(in_channels),
        nn.ReLU(True),
        nn.Conv2d(in_channels, out_channels, kernel_size, 
            stride=stride, padding=padding, bias=bias),
    ]
    return blocks


class BasicBlock(nn.Module):
    """ Basic Block structure in ResNet, constructed with two 3x3 conv layer.
    The first conv layer change the size as well as the number of channel, the second 
    keep them the same.
    """
    def __init__(self, inc, outc, stride = 1, use_pre_activation = False):
        super(BasicBlock, self).__init__()

        self.use_pre_activation = use_pre_activation
        if use_pre_activation:
            Conv = PreConvLayer
        else:
            Conv = ConvLayer

        self.unit = nn.Sequential(
            *Conv(inc, outc, 3, stride, 1), # 默认采用 3,2,1 的结构，大小减半。
            *Conv(outc, outc, 3, 1, 1, no_activation=True),
        )
        self.projector = None
        if stride != 1 or inc != outc: # 如果大小改变或者通道数改变，需要1x1卷积变换x大小或者维度，1,2,0 大小减半， 1,1,0 不变。
            self.projector = nn.Sequential( 
                *ConvLayer(inc, outc, 1, stride, no_activation=True) 
            )

    def forward(self, x):
        y = self.unit(x) 
        z = x
        if self.projector is not None:
            z = self.projector(x)
        if self.use_pre_activation:
            return y + z
        else:
            return F.relu(z + y, True)


class BottleNeck(nn.Module):
    """ BottleNeck structure in ResNet, constructed with 1x1, 3x3, 1x1 conv layer.
    """
    def __init__(self, inc, outc, stride = 1, channel_down_scale = 4, use_pre_activation = False):
        super(BottleNeck, self).__init__()

        self.use_pre_activation = use_pre_activation
        if use_pre_activation:
            Conv = PreConvLayer
        else:
            Conv = ConvLayer

        self.unit = nn.Sequential(
            *Conv(inc, outc // channel_down_scale, 1, 1), # 1x1 做通道的 down sample
            *Conv(outc // channel_down_scale, outc // channel_down_scale, 3, stride, 1), # 中间的ConvLayer做大小变换。
            *Conv(outc // channel_down_scale, outc, 1, 1, no_activation=True) # 1x1 恢复通道数
        )
        self.projector = None
        if stride != 1 or inc != outc:
            self.projector = nn.Sequential( # 如果大小要改变，需要使用1x1的卷积层变换 x 的维度和大小，默认采用  1,2,0 的结构，大小减半。
                *ConvLayer(inc, outc, 1, stride, no_activation=True) 
            )

    def forward(self, x):
        y = self.unit(x) 
        z = x
        if self.projector is not None:
            z = self.projector(x)
        if self.use_pre_activation:
            return y + z
        else:
            return F.relu(z + y, True)


class ResBlocks(nn.Module):
    def __init__(self, num_blocks, last_channel, channel, down_scale = True, 
                use_bottle_block = False, pre_act = False):
        super(ResBlocks, self).__init__()
        self.blocks = []
        if use_bottle_block:
            BuildBlock = BottleNeck
        else:
            BuildBlock = BasicBlock

        for _ in range(num_blocks - 1): # 先添加输入输出不变的Block
            self.blocks.append(BuildBlock(channel, channel, 1, use_pre_activation = pre_act))
        self.blocks.insert(0, 
            BuildBlock(last_channel, channel, 2 if down_scale else 1, use_pre_activation= pre_act)
        ) # 第一层的Block可能改变通道数和大小。
        self.blocks = nn.Sequential(*self.blocks)

    def forward(self, x):
        return self.blocks(x)


class ResNet(nn.Module):
    def __init__(self, input_channel = 3, num_classes = 1000, configure = 18, pre_act = False):
        super(ResNet, self).__init__()

        # add configuration here for more ResNet likely structure.
        resnet18 = [[2,2,2,2], [64,128,256,512], [False, True, True, True], False]
        resnet34 = [[3,4,6,3], [64,128,256,512], [False, True, True, True], False]
        resnet50 = [[3,4,6,3], [256,512,1024,2048], [False, True, True, True], True]
        resnet101 = [[3,4,23,3], [256,512,1024,2048], [False, True, True, True], True]
        resnet152 = [[3,4,36,3], [256,512,1024,2048], [False, True, True, True], True]

        if configure == 18:
            config = resnet18
        elif configure == 34:
            config = resnet34
        elif configure == 50:
            config = resnet50
        elif configure == 101:
            config = resnet101
        elif configure == 152:
            config = resnet152
        else:
            raise Exception("No such pre-set configuration for ResNet-{}".format(configure))

        num_blocks_config = config[0]
        channel_config = config[1] # output channel of one blocks
        downscale_config = config[2] # change the size of image or not
        use_bottle_block = config[3] # use bottlenet structure for resnet50 and over

        self.feature = [
            *ConvLayer(input_channel, 64, 7, 2, 3),
            nn.MaxPool2d(3, 2, 1)
        ]

        last_channel = 64
        for num_blocks, channel, downscale in zip(num_blocks_config, channel_config, downscale_config):
            self.feature.append(ResBlocks(num_blocks, last_channel, channel, downscale, use_bottle_block, pre_act))
            last_channel = channel

        self.feature = nn.Sequential(*self.feature)
        self.avg_pooling = nn.AdaptiveAvgPool2d(1)
        self.classifier =  nn.Linear(last_channel, num_classes)


    def forward(self, x):
        x = self.feature(x)
        x = self.avg_pooling(x).view(x.size(0), -1)
        return self.classifier(x)


if __name__ == "__main__":
    
    image = torch.randn(4,3,224,224)
    net = ResNet(num_classes = 10, configure=152, pre_act=True)
    net.apply(weights_init)
    print(net(image).size())
    print(net)