models_rssrai.py

import json
from functools import partial
from pathlib import Path

import attr
import torch
import torch.cuda
import torch.nn as nn
import torch.nn.functional as F

from collections import OrderedDict


@attr.s(slots=True)
class HyperParams:
    classes = attr.ib(default=list(range(16)))
    net = attr.ib(default='UNet')
    n_channels = attr.ib(default=4)  # max 20
    total_classes = 16
    thresholds = attr.ib(default=[0.5])
    pre_buffer = attr.ib(default=0.0)

    patch_inner = attr.ib(default=128)
    patch_border = attr.ib(default=16)

    augment_rotations = attr.ib(default=10.0)  # degrees
    augment_flips = attr.ib(default=1)
    augment_channels = attr.ib(default=0.0)

    validation_square = attr.ib(default=400)

    dropout = attr.ib(default=0.0)
    bn = attr.ib(default=1)
    activation = attr.ib(default='relu')
    top_scale = attr.ib(default=2)
    log_loss = attr.ib(default=1.0)
    dice_loss = attr.ib(default=0.0)
    jaccard_loss = attr.ib(default=0.0)
    dist_loss = attr.ib(default=0.0)
    dist_dice_loss = attr.ib(default=0.0)
    dist_jaccard_loss = attr.ib(default=0.0)

    filters_base = attr.ib(default=32)

    n_epochs = attr.ib(default=150)
    oversample = attr.ib(default=0.0)
    lr = attr.ib(default=0.0001)
    lr_decay = attr.ib(default=0.0)
    weight_decay = attr.ib(default=0.0)
    batch_size = attr.ib(default=128)

    @property
    def n_classes(self):
        return len(self.classes)

    @property
    def has_all_classes(self):
        return self.n_classes == self.total_classes

    @property
    def needs_dist(self):
        return (self.dist_loss != 0 or self.dist_dice_loss != 0 or
                self.dist_jaccard_loss != 0)

    @classmethod
    def from_dir(cls, root: Path):
        params = json.loads(root.joinpath('hps.json').read_text())
        fields = {field.name for field in attr.fields(HyperParams)}
        return cls(**{k: v for k, v in params.items() if k in fields})

    def update(self, hps_string: str):
        if hps_string:
            values = dict(pair.split('=') for pair in hps_string.split(','))
            for field in attr.fields(HyperParams):
                v = values.pop(field.name, None)
                if v is not None:
                    default = field.default
                    assert not isinstance(default, bool)
                    if isinstance(default, (int, float, str)):
                        v = type(default)(v)
                    elif isinstance(default, list):
                        v = [type(default[0])(x) for x in v.split('-')]
                    setattr(self, field.name, v)
            if values:
                raise ValueError('Unknown hyperparams: {}'.format(values))


class BaseNet(nn.Module):
    def __init__(self, hps: HyperParams):
        super().__init__()
        self.hps = hps
        if hps.dropout:
            self.dropout2d = nn.Dropout2d(p=hps.dropout)
        else:
            self.dropout2d = lambda x: x
        self.register_buffer('global_step', torch.IntTensor(1).zero_())


def conv3x3(in_, out):
    return nn.Conv2d(in_, out, 3, padding=1)


def concat(xs):
    return torch.cat(xs, 1)


class MiniNet(BaseNet):
    def __init__(self, hps):
        super().__init__(hps)
        self.conv1 = nn.Conv2d(hps.n_channels, 4, 1)
        self.conv2 = nn.Conv2d(4, 8, 3, padding=1)
        self.conv3 = nn.Conv2d(8, hps.n_classes, 3, padding=1)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = self.conv3(x)
        b = self.hps.patch_border
        return F.sigmoid(x[:, :, b:-b, b:-b])


class OldNet(BaseNet):
    def __init__(self, hps):
        super().__init__(hps)
        self.conv1 = nn.Conv2d(hps.n_channels, 64, 5, padding=2)
        self.conv2 = nn.Conv2d(64, 64, 5, padding=2)
        self.conv3 = nn.Conv2d(64, 64, 5, padding=2)
        self.conv4 = nn.Conv2d(64, hps.n_classes, 7, padding=3)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = self.conv4(x)
        b = self.hps.patch_border
        return F.sigmoid(x[:, :, b:-b, b:-b])


class SmallNet(BaseNet):
    def __init__(self, hps):
        super().__init__(hps)
        self.conv1 = nn.Conv2d(hps.n_channels, 64, 3, padding=1)
        self.conv2 = nn.Conv2d(64, 64, 3, padding=1)
        self.conv3 = nn.Conv2d(64, 64, 3, padding=1)
        self.conv4 = nn.Conv2d(64, 128, 3, padding=1)
        self.conv5 = nn.Conv2d(128, hps.n_classes, 3, padding=1)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = F.relu(self.conv4(x))
        x = self.conv5(x)
        b = self.hps.patch_border
        return F.sigmoid(x[:, :, b:-b, b:-b])


# UNet:
# http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net/u-net-architecture.png


class SmallUNet(BaseNet):
    def __init__(self, hps):
        super().__init__(hps)
        self.conv1 = nn.Conv2d(hps.n_channels, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 32, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.upsample = nn.UpsamplingNearest2d(scale_factor=2)
        self.conv3 = nn.Conv2d(32, 64, 3, padding=1)
        self.conv4 = nn.Conv2d(64, 64, 3, padding=1)
        self.conv5 = nn.Conv2d(64, 32, 3, padding=1)
        self.conv6 = nn.Conv2d(64, 32, 3, padding=1)
        self.conv7 = nn.Conv2d(32, hps.n_classes, 3, padding=1)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x1 = self.pool(x)
        x1 = F.relu(self.conv3(x1))
        x1 = F.relu(self.conv4(x1))
        x1 = F.relu(self.conv5(x1))
        x1 = self.upsample(x1)
        x = torch.cat([x, x1], 1)
        x = F.relu(self.conv6(x))
        x = self.conv7(x)
        b = self.hps.patch_border
        return F.sigmoid(x[:, :, b:-b, b:-b])


class UNetModule(nn.Module):
    def __init__(self, hps: HyperParams, in_: int, out: int):
        super().__init__()
        self.conv1 = conv3x3(in_, out)
        self.conv2 = conv3x3(out, out)
        self.bn = hps.bn
        self.activation = getattr(F, hps.activation)
        if self.bn:
            self.bn1 = nn.BatchNorm2d(out)
            self.bn2 = nn.BatchNorm2d(out)

    def forward(self, x):
        x = self.conv1(x)
        if self.bn:
            x = self.bn1(x)
        x = self.activation(x)
        x = self.conv2(x)
        if self.bn:
            x = self.bn2(x)
        x = self.activation(x)
        return x


class UNet(BaseNet):
    module = UNetModule
    filter_factors = [1, 2, 4, 8, 16]

    def __init__(self, hps: HyperParams):
        super().__init__(hps)
        self.pool = nn.MaxPool2d(2, 2)
        self.pool_top = nn.MaxPool2d(hps.top_scale, hps.top_scale)
        self.upsample = nn.UpsamplingNearest2d(scale_factor=2)
        self.upsample_top = nn.UpsamplingNearest2d(scale_factor=hps.top_scale)
        filter_sizes = [hps.filters_base * s for s in self.filter_factors]
        self.down, self.up = [], []
        for i, nf in enumerate(filter_sizes):
            low_nf = hps.n_channels if i == 0 else filter_sizes[i - 1]
            self.down.append(self.module(hps, low_nf, nf))
            setattr(self, 'down_{}'.format(i), self.down[-1])
            if i != 0:
                self.up.append(self.module(hps, low_nf + nf, low_nf))
                setattr(self, 'conv_up_{}'.format(i), self.up[-1])
        self.conv_final = nn.Conv2d(filter_sizes[0], hps.n_classes, 1)

    def forward(self, x):
        xs = []
        for i, down in enumerate(self.down):
            if i == 0:
                x_in = x
            elif i == 1:
                x_in = self.pool_top(xs[-1])
            else:
                x_in = self.pool(xs[-1])
            x_out = down(x_in)
            x_out = self.dropout2d(x_out)
            xs.append(x_out)

        x_out = xs[-1]
        for i, (x_skip, up) in reversed(list(enumerate(zip(xs[:-1], self.up)))):
            upsample = self.upsample_top if i == 0 else self.upsample
            x_out = up(torch.cat([upsample(x_out), x_skip], 1))
            x_out = self.dropout2d(x_out)

        x_out = self.conv_final(x_out)
        b = self.hps.patch_border
        return torch.sigmoid(x_out[:, :, b:-b, b:-b])


class Conv3BN(nn.Module):
    def __init__(self, hps: HyperParams, in_: int, out: int):
        super().__init__()
        self.conv = conv3x3(in_, out)
        self.bn = nn.BatchNorm2d(out) if hps.bn else None
        self.activation = getattr(F, hps.activation)

    def forward(self, x):
        x = self.conv(x)
        if self.bn is not None:
            x = self.bn(x)
        x = self.activation(x, inplace=True)
        return x


class UNet3lModule(nn.Module):
    def __init__(self, hps: HyperParams, in_: int, out: int):
        super().__init__()
        self.l1 = Conv3BN(hps, in_, out)
        self.l2 = Conv3BN(hps, out, out)
        self.l3 = Conv3BN(hps, out, out)

    def forward(self, x):
        x = self.l1(x)
        x = self.l2(x)
        x = self.l3(x)
        return x


class UNet3l(UNet):
    module = UNet3lModule


class UNet2Module(nn.Module):
    def __init__(self, hps: HyperParams, in_: int, out: int):
        super().__init__()
        self.l1 = Conv3BN(hps, in_, out)
        self.l2 = Conv3BN(hps, out, out)

    def forward(self, x):
        x = self.l1(x)
        x = self.l2(x)
        return x


class UNet2(BaseNet):
    def __init__(self, hps):
        super().__init__(hps)
        b = hps.filters_base
        self.filters = [b * 2, b * 2, b * 4, b * 8, b * 16]
        self.upsample = nn.UpsamplingNearest2d(scale_factor=2)
        self.down, self.down_pool, self.mid, self.up = [[] for _ in range(4)]
        for i, nf in enumerate(self.filters):
            low_nf = hps.n_channels if i == 0 else self.filters[i - 1]
            self.down_pool.append(
                nn.Conv2d(low_nf, low_nf, 3, padding=1, stride=2))
            setattr(self, 'down_pool_{}'.format(i), self.down_pool[-1])
            self.down.append(UNet2Module(hps, low_nf, nf))
            setattr(self, 'down_{}'.format(i), self.down[-1])
            if i != 0:
                self.mid.append(Conv3BN(hps, low_nf, low_nf))
                setattr(self, 'mid_{}'.format(i), self.mid[-1])
                self.up.append(UNet2Module(hps, low_nf + nf, low_nf))
                setattr(self, 'up_{}'.format(i), self.up[-1])
        self.conv_final = nn.Conv2d(self.filters[0], hps.n_classes, 1)

    def forward(self, x):
        xs = []
        for i, (down, down_pool) in enumerate(zip(self.down, self.down_pool)):
            x_out = down(down_pool(xs[-1]) if xs else x)
            xs.append(x_out)

        x_out = xs[-1]
        for x_skip, up, mid in reversed(list(zip(xs[:-1], self.up, self.mid))):
            x_out = up(torch.cat([self.upsample(x_out), mid(x_skip)], 1))

        x_out = self.conv_final(x_out)
        b = self.hps.patch_border
        return F.sigmoid(x_out[:, :, b:-b, b:-b])


class BasicConv2d(nn.Module):
    def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0):
        super().__init__()
        self.conv = nn.Conv2d(
            in_planes, out_planes,
            kernel_size=kernel_size, stride=stride, padding=padding, bias=False)
        self.bn = nn.BatchNorm2d(out_planes)
        self.relu = nn.ReLU(inplace=True)

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


class InceptionModule(nn.Module):
    def __init__(self, hps: HyperParams, in_: int, out: int):
        super().__init__()
        out_1 = out * 3 // 8
        out_2 = out * 2 // 8
        self.conv1x1 = BasicConv2d(in_, out_1, kernel_size=1)
        self.conv3x3_pre = BasicConv2d(in_, in_ // 2, kernel_size=1)
        self.conv3x3 = BasicConv2d(in_ // 2, out_1, kernel_size=3, padding=1)
        self.conv5x5_pre = BasicConv2d(in_, in_ // 4, kernel_size=1)
        self.conv5x5 = BasicConv2d(in_ // 4, out_2, kernel_size=5, padding=2)
        assert hps.bn
        assert hps.activation == 'relu'

    def forward(self, x):
        return torch.cat([
            self.conv1x1(x),
            self.conv3x3(self.conv3x3_pre(x)),
            self.conv5x5(self.conv5x5_pre(x)),
        ], 1)


class Inception2Module(nn.Module):
    def __init__(self, hps: HyperParams, in_: int, out: int):
        super().__init__()
        self.l1 = InceptionModule(hps, in_, out)
        self.l2 = InceptionModule(hps, out, out)

    def forward(self, x):
        x = self.l1(x)
        x = self.l2(x)
        return x


class InceptionUNet(UNet):
    module = InceptionModule


class Inception2UNet(UNet):
    module = Inception2Module


class SimpleSegNet(BaseNet):
    def __init__(self, hps):
        super().__init__(hps)
        s = hps.filters_base
        self.pool = nn.MaxPool2d(2, 2)
        self.upsample = nn.UpsamplingNearest2d(scale_factor=2)
        self.input_conv = BasicConv2d(hps.n_channels, s, 1)
        self.enc_1 = BasicConv2d(s * 1, s * 2, 3, padding=1)
        self.enc_2 = BasicConv2d(s * 2, s * 4, 3, padding=1)
        self.enc_3 = BasicConv2d(s * 4, s * 8, 3, padding=1)
        self.enc_4 = BasicConv2d(s * 8, s * 8, 3, padding=1)
        # https://github.com/pradyu1993/segnet - decoder lacks relu (???)
        self.dec_4 = BasicConv2d(s * 8, s * 8, 3, padding=1)
        self.dec_3 = BasicConv2d(s * 8, s * 4, 3, padding=1)
        self.dec_2 = BasicConv2d(s * 4, s * 2, 3, padding=1)
        self.dec_1 = BasicConv2d(s * 2, s * 1, 3, padding=1)
        self.conv_final = nn.Conv2d(s, hps.n_classes, 1)

    def forward(self, x):
        # Input
        x = self.input_conv(x)
        # Encoder
        x = self.enc_1(x)
        x = self.pool(x)
        x = self.enc_2(x)
        x = self.pool(x)
        x = self.enc_3(x)
        x = self.pool(x)
        x = self.enc_4(x)
        # Decoder
        x = self.dec_4(x)
        x = self.upsample(x)
        x = self.dec_3(x)
        x = self.upsample(x)
        x = self.dec_2(x)
        x = self.upsample(x)
        x = self.dec_1(x)
        # Output
        x = self.conv_final(x)
        b = self.hps.patch_border
        return F.sigmoid(x[:, :, b:-b, b:-b])


class DenseLayer(nn.Module):
    def __init__(self, in_, out, *, dropout, bn):
        super().__init__()
        self.bn = nn.BatchNorm2d(in_) if bn else None
        self.activation = nn.ReLU(inplace=True)
        self.conv = conv3x3(in_, out)
        self.dropout = nn.Dropout2d(p=dropout) if dropout else None

    def forward(self, x):
        x = self.activation(x)
        if self.bn is not None:
            x = self.bn(x)
        x = self.conv(x)
        if self.dropout is not None:
            x = self.dropout(x)
        return x


class DenseBlock(nn.Module):
    def __init__(self, in_, k, n_layers, dropout, bn):
        super().__init__()
        self.out = k * n_layers
        layer_in = in_
        self.layers = []
        for i in range(n_layers):
            layer = DenseLayer(layer_in, k, dropout=dropout, bn=bn)
            self.layers.append(layer)
            setattr(self, 'layer_{}'.format(i), layer)
            layer_in += k

    def forward(self, x):
        inputs = [x]
        outputs = []
        for i, layer in enumerate(self.layers[:-1]):
            outputs.append(layer(inputs[i]))
            inputs.append(concat([outputs[i], inputs[i]]))
        return torch.cat([self.layers[-1](inputs[-1])] + outputs, 1)


class DenseUNetModule(DenseBlock):
    def __init__(self, hps: HyperParams, in_: int, out: int):
        n_layers = 4
        super().__init__(in_, out // n_layers, n_layers,
                         dropout=hps.dropout, bn=hps.bn)


class DenseUNet(UNet):
    module = DenseUNetModule


class DownBlock(nn.Module):
    def __init__(self, in_, out, scale, *, dropout, bn):
        super().__init__()
        self.in_ = in_
        self.bn = nn.BatchNorm2d(in_) if bn else None
        self.activation = nn.ReLU(inplace=True)
        self.conv = nn.Conv2d(in_, out, 1)
        self.dropout = nn.Dropout2d(p=dropout) if dropout else None
        self.pool = nn.MaxPool2d(scale, scale)

    def forward(self, x):
        if self.bn is not None:
            x = self.bn(x)
        x = self.activation(x)
        x = self.conv(x)
        if self.dropout is not None:
            x = self.dropout(x)
        x = self.pool(x)
        return x


class UpBlock(nn.Module):
    def __init__(self, in_, out, scale):
        super().__init__()
        self.up_conv = nn.Conv2d(in_, out, 1)
        self.upsample = nn.UpsamplingNearest2d(scale_factor=scale)

    def forward(self, x):
        return self.upsample(self.up_conv(x))


class DenseNet(BaseNet):
    """ https://arxiv.org/pdf/1611.09326v2.pdf
    """
    def __init__(self, hps):
        super().__init__(hps)
        k = hps.filters_base
        block_layers = [3, 5, 7, 5, 3]
        block_in = [n * k for n in [3, 8, 16, 8, 4]]
        scale_factors = [4, 2]
        dense = partial(DenseBlock, dropout=hps.dropout, bn=hps.bn)
        self.input_conv = nn.Conv2d(hps.n_channels, block_in[0], 3, padding=1)
        self.blocks = []
        self.scales = []
        self.n_layers = len(block_layers) // 2
        for i, (in_, l) in enumerate(zip(block_in, block_layers)):
            if i < self.n_layers:
                block = dense(in_, k, l)
                scale = DownBlock(block.out + in_, block_in[i + 1],
                                  scale_factors[i],
                                  dropout=hps.dropout, bn=hps.bn)
            elif i == self.n_layers:
                block = dense(in_, k, l)
                scale = None
            else:
                block = dense(in_ + self.scales[2 * self.n_layers - i].in_,
                              k, l)
                scale = UpBlock(self.blocks[-1].out, in_,
                                scale_factors[2 * self.n_layers - i])
            setattr(self, 'block_{}'.format(i), block)
            setattr(self, 'scale_{}'.format(i), scale)
            self.blocks.append(block)
            self.scales.append(scale)
        self.output_conv = nn.Conv2d(self.blocks[-1].out, hps.n_classes, 1)

    def forward(self, x):
        # Input
        x = self.input_conv(x)
        # Network
        skips = []
        for i, (block, scale) in enumerate(zip(self.blocks, self.scales)):
            if i < self.n_layers:
                x = concat([block(x), x])
                skips.append(x)
                x = scale(x)
            elif i == self.n_layers:
                x = block(x)
            else:
                x = block(concat([scale(x), skips[2 * self.n_layers - i]]))
        # Output
        x = self.output_conv(x)
        b = self.hps.patch_border
        return F.sigmoid(x[:, :, b:-b, b:-b])


class _ImagePool(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.conv = _ConvBnReLU(in_ch, out_ch, 1, 1, 0, 1)

    def forward(self, x):
        _, _, H, W = x.shape
        h = self.pool(x)
        h = self.conv(h)
        h = F.interpolate(h, size=(H, W), mode="bilinear", align_corners=False)
        return h

class _ASPP(nn.Module):
    """
    Atrous spatial pyramid pooling with image-level feature
    """

    def __init__(self, in_ch, out_ch, rates):
        super(_ASPP, self).__init__()
        self.stages = nn.Module()
        self.stages.add_module("c0", _ConvBnReLU(in_ch, out_ch, 1, 1, 0, 1))
        for i, rate in enumerate(rates):
            self.stages.add_module(
                "c{}".format(i + 1),
                _ConvBnReLU(in_ch, out_ch, 3, 1, padding=rate, dilation=rate),
            )
        self.stages.add_module("imagepool", _ImagePool(in_ch, out_ch))

    def forward(self, x):
        return torch.cat([stage(x) for stage in self.stages.children()], dim=1)


class DeepLabV3Plus(nn.Module):
    """
    DeepLab v3+: Dilated ResNet with multi-grid + improved ASPP + decoder
    """

    def __init__(self, hps: HyperParams):
        super(DeepLabV3Plus, self).__init__()

        n_classes = 16
        n_blocks = [3, 4, 6, 3]
        atrous_rates = [6, 12, 18]
        multi_grids = [1, 2, 4]
        output_stride = 16

        self.hps = hps

        # Stride and dilation
        if output_stride == 8:
            s = [1, 2, 1, 1]
            d = [1, 1, 2, 4]
        elif output_stride == 16:
            s = [1, 2, 2, 1]
            d = [1, 1, 1, 2]

        # Encoder
        ch = [64 * 2 ** p for p in range(6)]
        self.layer1 = _Stem(ch[0])
        self.layer2 = _ResLayer(n_blocks[0], ch[0], ch[2], s[0], d[0])
        self.layer3 = _ResLayer(n_blocks[1], ch[2], ch[3], s[1], d[1])
        self.layer4 = _ResLayer(n_blocks[2], ch[3], ch[4], s[2], d[2])
        self.layer5 = _ResLayer(n_blocks[3], ch[4], ch[5], s[3], d[3], multi_grids)
        self.aspp = _ASPP(ch[5], 256, atrous_rates)
        concat_ch = 256 * (len(atrous_rates) + 2)
        self.add_module("fc1", _ConvBnReLU(concat_ch, 256, 1, 1, 0, 1))

        # Decoder
        self.reduce = _ConvBnReLU(256, 48, 1, 1, 0, 1)
        self.fc2 = nn.Sequential(
            OrderedDict(
                [
                    ("conv1", _ConvBnReLU(304, 256, 3, 1, 1, 1)),
                    ("conv2", _ConvBnReLU(256, 256, 3, 1, 1, 1)),
                    ("conv3", nn.Conv2d(256, n_classes, kernel_size=1)),
                ]
            )
        )

        self.register_buffer('global_step', torch.IntTensor(1).zero_())

    def forward(self, x):
        h = self.layer1(x)
        h = self.layer2(h)
        h_ = self.reduce(h)
        h = self.layer3(h)
        h = self.layer4(h)
        h = self.layer5(h)
        h = self.aspp(h)
        h = self.fc1(h)
        h = F.interpolate(h, size=h_.shape[2:], mode="bilinear", align_corners=False)
        h = torch.cat((h, h_), dim=1)
        h = self.fc2(h)
        h = F.interpolate(h, size=x.shape[2:], mode="bilinear", align_corners=False)
        b = self.hps.patch_border

        return torch.sigmoid(h[:, :, b:-b, b:-b])



try:
    from encoding.nn import SyncBatchNorm

    _BATCH_NORM = SyncBatchNorm
except:
    _BATCH_NORM = nn.BatchNorm2d

_BOTTLENECK_EXPANSION = 4


class _ConvBnReLU(nn.Sequential):
    """
    Cascade of 2D convolution, batch norm, and ReLU.
    """

    BATCH_NORM = _BATCH_NORM

    def __init__(
            self, in_ch, out_ch, kernel_size, stride, padding, dilation, relu=True
    ):
        super(_ConvBnReLU, self).__init__()
        self.add_module(
            "conv",
            nn.Conv2d(
                in_ch, out_ch, kernel_size, stride, padding, dilation, bias=False
            ),
        )
        self.add_module("bn", _BATCH_NORM(out_ch, eps=1e-5, momentum=0.999))

        if relu:
            self.add_module("relu", nn.ReLU())


class _Bottleneck(nn.Module):
    """
    Bottleneck block of MSRA ResNet.
    """

    def __init__(self, in_ch, out_ch, stride, dilation, downsample):
        super(_Bottleneck, self).__init__()
        mid_ch = out_ch // _BOTTLENECK_EXPANSION
        self.reduce = _ConvBnReLU(in_ch, mid_ch, 1, stride, 0, 1, True)
        self.conv3x3 = _ConvBnReLU(mid_ch, mid_ch, 3, 1, dilation, dilation, True)
        self.increase = _ConvBnReLU(mid_ch, out_ch, 1, 1, 0, 1, False)
        self.shortcut = (
            _ConvBnReLU(in_ch, out_ch, 1, stride, 0, 1, False)
            if downsample
            else lambda x: x  # identity
        )

    def forward(self, x):
        h = self.reduce(x)
        h = self.conv3x3(h)
        h = self.increase(h)
        h += self.shortcut(x)
        return F.relu(h)


class _ResLayer(nn.Sequential):
    """
    Residual layer with multi grids
    """

    def __init__(self, n_layers, in_ch, out_ch, stride, dilation, multi_grids=None):
        super(_ResLayer, self).__init__()

        if multi_grids is None:
            multi_grids = [1 for _ in range(n_layers)]
        else:
            assert n_layers == len(multi_grids)

        # Downsampling is only in the first block
        for i in range(n_layers):
            self.add_module(
                "block{}".format(i + 1),
                _Bottleneck(
                    in_ch=(in_ch if i == 0 else out_ch),
                    out_ch=out_ch,
                    stride=(stride if i == 0 else 1),
                    dilation=dilation * multi_grids[i],
                    downsample=(True if i == 0 else False),
                ),
            )


class _Stem(nn.Sequential):
    """
    The 1st conv layer.
    Note that the max pooling is different from both MSRA and FAIR ResNet.
    """

    def __init__(self, out_ch):
        super(_Stem, self).__init__()
        self.add_module("conv1", _ConvBnReLU(4, out_ch, 7, 2, 3, 1))
        self.add_module("pool", nn.MaxPool2d(3, 2, 1, ceil_mode=True))