xautodl/lib/models/CifarResNet.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from .initialization import initialize_resnet
from .SharedUtils import additive_func


class Downsample(nn.Module):
    def __init__(self, nIn, nOut, stride):
        super(Downsample, self).__init__()
        assert stride == 2 and nOut == 2 * nIn, "stride:{} IO:{},{}".format(
            stride, nIn, nOut
        )
        self.in_dim = nIn
        self.out_dim = nOut
        self.avg = nn.AvgPool2d(kernel_size=2, stride=2, padding=0)
        self.conv = nn.Conv2d(nIn, nOut, kernel_size=1, stride=1, padding=0, bias=False)

    def forward(self, x):
        x = self.avg(x)
        out = self.conv(x)
        return out


class ConvBNReLU(nn.Module):
    def __init__(self, nIn, nOut, kernel, stride, padding, bias, relu):
        super(ConvBNReLU, self).__init__()
        self.conv = nn.Conv2d(
            nIn, nOut, kernel_size=kernel, stride=stride, padding=padding, bias=bias
        )
        self.bn = nn.BatchNorm2d(nOut)
        if relu:
            self.relu = nn.ReLU(inplace=True)
        else:
            self.relu = None
        self.out_dim = nOut
        self.num_conv = 1

    def forward(self, x):
        conv = self.conv(x)
        bn = self.bn(conv)
        if self.relu:
            return self.relu(bn)
        else:
            return bn


class ResNetBasicblock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride):
        super(ResNetBasicblock, self).__init__()
        assert stride == 1 or stride == 2, "invalid stride {:}".format(stride)
        self.conv_a = ConvBNReLU(inplanes, planes, 3, stride, 1, False, True)
        self.conv_b = ConvBNReLU(planes, planes, 3, 1, 1, False, False)
        if stride == 2:
            self.downsample = Downsample(inplanes, planes, stride)
        elif inplanes != planes:
            self.downsample = ConvBNReLU(inplanes, planes, 1, 1, 0, False, False)
        else:
            self.downsample = None
        self.out_dim = planes
        self.num_conv = 2

    def forward(self, inputs):

        basicblock = self.conv_a(inputs)
        basicblock = self.conv_b(basicblock)

        if self.downsample is not None:
            residual = self.downsample(inputs)
        else:
            residual = inputs
        out = additive_func(residual, basicblock)
        return F.relu(out, inplace=True)


class ResNetBottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, stride):
        super(ResNetBottleneck, self).__init__()
        assert stride == 1 or stride == 2, "invalid stride {:}".format(stride)
        self.conv_1x1 = ConvBNReLU(inplanes, planes, 1, 1, 0, False, True)
        self.conv_3x3 = ConvBNReLU(planes, planes, 3, stride, 1, False, True)
        self.conv_1x4 = ConvBNReLU(
            planes, planes * self.expansion, 1, 1, 0, False, False
        )
        if stride == 2:
            self.downsample = Downsample(inplanes, planes * self.expansion, stride)
        elif inplanes != planes * self.expansion:
            self.downsample = ConvBNReLU(
                inplanes, planes * self.expansion, 1, 1, 0, False, False
            )
        else:
            self.downsample = None
        self.out_dim = planes * self.expansion
        self.num_conv = 3

    def forward(self, inputs):

        bottleneck = self.conv_1x1(inputs)
        bottleneck = self.conv_3x3(bottleneck)
        bottleneck = self.conv_1x4(bottleneck)

        if self.downsample is not None:
            residual = self.downsample(inputs)
        else:
            residual = inputs
        out = additive_func(residual, bottleneck)
        return F.relu(out, inplace=True)


class CifarResNet(nn.Module):
    def __init__(self, block_name, depth, num_classes, zero_init_residual):
        super(CifarResNet, self).__init__()

        # Model type specifies number of layers for CIFAR-10 and CIFAR-100 model
        if block_name == "ResNetBasicblock":
            block = ResNetBasicblock
            assert (depth - 2) % 6 == 0, "depth should be one of 20, 32, 44, 56, 110"
            layer_blocks = (depth - 2) // 6
        elif block_name == "ResNetBottleneck":
            block = ResNetBottleneck
            assert (depth - 2) % 9 == 0, "depth should be one of 164"
            layer_blocks = (depth - 2) // 9
        else:
            raise ValueError("invalid block : {:}".format(block_name))

        self.message = "CifarResNet : Block : {:}, Depth : {:}, Layers for each block : {:}".format(
            block_name, depth, layer_blocks
        )
        self.num_classes = num_classes
        self.channels = [16]
        self.layers = nn.ModuleList([ConvBNReLU(3, 16, 3, 1, 1, False, True)])
        for stage in range(3):
            for iL in range(layer_blocks):
                iC = self.channels[-1]
                planes = 16 * (2 ** stage)
                stride = 2 if stage > 0 and iL == 0 else 1
                module = block(iC, planes, stride)
                self.channels.append(module.out_dim)
                self.layers.append(module)
                self.message += "\nstage={:}, ilayer={:02d}/{:02d}, block={:03d}, iC={:3d}, oC={:3d}, stride={:}".format(
                    stage,
                    iL,
                    layer_blocks,
                    len(self.layers) - 1,
                    iC,
                    module.out_dim,
                    stride,
                )

        self.avgpool = nn.AvgPool2d(8)
        self.classifier = nn.Linear(module.out_dim, num_classes)
        assert (
            sum(x.num_conv for x in self.layers) + 1 == depth
        ), "invalid depth check {:} vs {:}".format(
            sum(x.num_conv for x in self.layers) + 1, depth
        )

        self.apply(initialize_resnet)
        if zero_init_residual:
            for m in self.modules():
                if isinstance(m, ResNetBasicblock):
                    nn.init.constant_(m.conv_b.bn.weight, 0)
                elif isinstance(m, ResNetBottleneck):
                    nn.init.constant_(m.conv_1x4.bn.weight, 0)

    def get_message(self):
        return self.message

    def forward(self, inputs):
        x = inputs
        for i, layer in enumerate(self.layers):
            x = layer(x)
        features = self.avgpool(x)
        features = features.view(features.size(0), -1)
        logits = self.classifier(features)
        return features, logits