tresnet_v2.py

import math
from collections import OrderedDict
from functools import partial

import torch.nn as nn
import torch
from torch.nn import Module as Module

from src.models.tresnet.layers.anti_aliasing import AntiAliasDownsampleLayer
from .layers.adaptive_avgmax_pool import SelectAdaptivePool2d
from .layers.general_layers import SEModule, Flatten, SpaceToDepthModule
from .layers.bottleneck_head import bottleneck_head

from inplace_abn import InPlaceABN
from inplace_abn import ABN

model_urls = []  # TBD
attn_layer = SEModule
stem_type = SpaceToDepthModule


class Linear_normalize(torch.nn.Module):
    def __init__(self, in_features, out_features, bias=True, normalize_embedding=False,
                 normalize_weights=False):
        super(Linear_normalize, self).__init__()
        self.normalize_embedding = normalize_embedding
        self.normalize_weights = normalize_weights
        self.in_features = in_features
        self.out_features = out_features
        self.weight = torch.nn.Parameter(torch.Tensor(out_features, in_features))
        if bias:
            self.bias = torch.nn.Parameter(torch.Tensor(out_features))
        else:
            self.register_parameter('bias', None)
        self.scale_embeddings = 32
        self.reset_parameters()

    def reset_parameters(self):
        torch.nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        if self.bias is not None:
            fan_in, _ = torch.nn.init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / math.sqrt(fan_in)
            torch.nn.init.uniform_(self.bias, -bound, bound)

    def forward(self, input):
        if self.normalize_embedding:
            input_t = torch.nn.functional.normalize(input) * self.scale_embeddings
        else:
            input_t = input
        if self.normalize_weights:
            weights_t = torch.nn.functional.normalize(self.weight)
        else:
            weights_t = self.weight

        return torch.nn.functional.linear(input_t, weights_t, self.bias)


def InplacABN_to_ABN(module: nn.Module) -> nn.Module:
    # convert all InplaceABN layer to bit-accurate ABN layers.
    if isinstance(module, InPlaceABN):
        module_new = ABN(module.num_features, activation=module.activation,
                         activation_param=module.activation_param)
        for key in module.state_dict():
            module_new.state_dict()[key].copy_(module.state_dict()[key])
        module_new.training = module.training
        module_new.weight.data = module_new.weight.abs() + module_new.eps
        return module_new
    for name, child in reversed(module._modules.items()):
        new_child = InplacABN_to_ABN(child)
        if new_child != child:
            module._modules[name] = new_child
    return module


def conv3x3(in_planes, out_planes, stride=1):
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)


def conv3x3_depth(planes, stride=1):
    return nn.Conv2d(planes, planes, groups=planes, kernel_size=3, stride=stride, padding=1, bias=False)


def conv2d(ni, nf, stride):
    return nn.Sequential(
        nn.Conv2d(ni, nf, kernel_size=3, stride=stride, padding=1, bias=False),
        nn.BatchNorm2d(nf),
        nn.ReLU(inplace=True)
    )


def conv2d_ABN(ni, nf, stride, activation="leaky_relu", kernel_size=3, activation_param=1e-2, groups=1):
    activation_param = 1e-6
    return nn.Sequential(
        nn.Conv2d(ni, nf, kernel_size=kernel_size, stride=stride, padding=kernel_size // 2, groups=groups,
                  bias=False),
        InPlaceABN(num_features=nf, activation=activation, activation_param=activation_param)
    )


class BasicBlock(Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None, use_se=True, anti_alias_layer=None):
        super(BasicBlock, self).__init__()
        if stride == 1:
            self.conv1 = conv2d_ABN(inplanes, planes, stride=1, activation_param=1e-3)
        else:
            if anti_alias_layer is None:
                self.conv1 = conv2d_ABN(inplanes, planes, stride=2, activation_param=1e-3)
            else:
                self.conv1 = nn.Sequential(conv2d_ABN(inplanes, planes, stride=1, activation_param=1e-3),
                                           anti_alias_layer(channels=planes, filt_size=3, stride=2))

        self.conv2 = conv2d_ABN(planes, planes, stride=1, activation="identity")
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride
        reduce_layer_planes = max(planes * self.expansion // 4, 64)
        self.se = attn_layer(channels=planes * self.expansion, reduction_channels=reduce_layer_planes) if \
            use_se else None

    def forward(self, x):
        if self.downsample is not None:
            residual = self.downsample(x)
        else:
            residual = x

        out = self.conv1(x)
        out = self.conv2(out)

        if self.se is not None: out = self.se(out)

        out += residual

        out = self.relu(out)

        return out


class Bottleneck(Module):
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None, use_se=True, anti_alias_layer=None):
        super(Bottleneck, self).__init__()
        self.conv1 = conv2d_ABN(inplanes, planes, kernel_size=1, stride=1, activation="leaky_relu",
                                activation_param=1e-3)
        if stride == 1:
            self.conv2 = conv2d_ABN(planes, planes, kernel_size=3, stride=1, activation="leaky_relu",
                                    activation_param=1e-3)
        else:
            if anti_alias_layer is None:
                self.conv2 = conv2d_ABN(planes, planes, kernel_size=3, stride=2, activation="leaky_relu",
                                        activation_param=1e-3)
            else:
                self.conv2 = nn.Sequential(conv2d_ABN(planes, planes, kernel_size=3, stride=1,
                                                      activation="leaky_relu", activation_param=1e-3),
                                           anti_alias_layer(channels=planes, filt_size=3, stride=2))

        self.conv3 = conv2d_ABN(planes, planes * self.expansion, kernel_size=1, stride=1,
                                activation="identity")

        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

        reduce_layer_planes = max(planes * self.expansion // 8, 64)
        self.se = attn_layer(planes, reduce_layer_planes) if use_se else None

    def forward(self, x):
        if self.downsample is not None:
            residual = self.downsample(x)
        else:
            residual = x

        out = self.conv1(x)
        out = self.conv2(out)
        if self.se is not None: out = self.se(out)

        out = self.conv3(out)
        out = out + residual  # no inplace
        out = self.relu(out)

        return out


layer_1_2_type = BasicBlock


class TResNetV2(Module):

    def __init__(self, layers, in_chans=3, num_classes=1000, width_factor=1.0,remove_model_jit=False):
        super(TResNetV2, self).__init__()
        ## body
        self.inplanes = int(int(64 * width_factor + 4) / 8) * 8
        self.planes = int(int(64 * width_factor + 4) / 8) * 8
        # print("self.inplanes ={}".format(self.inplanes))
        if stem_type == SpaceToDepthModule:
            # JIT layers
            SpaceToDepth = stem_type(remove_model_jit=remove_model_jit)
            conv1 = conv2d_ABN(in_chans * 16, self.planes, stride=1, kernel_size=3)
        else:
            stem_chs_1 = stem_chs_2 = 32
            SpaceToDepth = nn.Sequential(*[
                conv2d_ABN(in_chans, stem_chs_1, kernel_size=3, stride=2,
                           activation="leaky_relu"),
                conv2d_ABN(stem_chs_1, stem_chs_2, kernel_size=3, stride=1,
                           activation="leaky_relu"),
                conv2d_ABN(stem_chs_2, self.inplanes, kernel_size=3, stride=1,
                           activation="leaky_relu"),
                nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
            ])
            conv1 = torch.nn.Identity()

        anti_alias_layer = partial(AntiAliasDownsampleLayer, remove_model_jit=remove_model_jit)
        global_pool_layer = SelectAdaptivePool2d(pool_type='avg', flatten=True)
        layer1 = self._make_layer(layer_1_2_type, self.planes, layers[0], stride=1, use_se=True,
                                  anti_alias_layer=anti_alias_layer)  # 56x56
        layer2 = self._make_layer(layer_1_2_type, self.planes * 2, layers[1], stride=2, use_se=True,
                                  anti_alias_layer=anti_alias_layer)  # 28x28
        layer3 = self._make_layer(Bottleneck, self.planes * 4, layers[2], stride=2, use_se=True,
                                  anti_alias_layer=anti_alias_layer)  # 14x14
        layer4 = self._make_layer(Bottleneck, self.planes * 8, layers[3], stride=2, use_se=False,
                                  anti_alias_layer=anti_alias_layer)  # 7x7

        # layer4[-1].relu = nn.Identity()
        # layer4[-1].conv3[1].bias.requires_grad = False

        self.body = nn.Sequential(OrderedDict([
            ('SpaceToDepth', SpaceToDepth),
            ('conv1', conv1),
            ('layer1', layer1),
            ('layer2', layer2),
            ('layer3', layer3),
            ('layer4', layer4)]))

        # default head
        self.num_features = (self.planes * 8) * Bottleneck.expansion

        fc = nn.Linear(self.num_features * global_pool_layer.feat_mult(), num_classes)

        self.global_pool = nn.Sequential(OrderedDict([('global_pool_layer', global_pool_layer)]))

        self.head = nn.Sequential(OrderedDict([('fc', fc)]))

        self.embeddings = []

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='leaky_relu')
            elif isinstance(m, nn.BatchNorm2d) or isinstance(m, InPlaceABN):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

        # initilize resnet in a magic way
        for m in self.modules():
            if isinstance(m, BasicBlock):
                m.conv2[1].weight = nn.Parameter(torch.zeros_like(m.conv2[1].weight))  # BN to zero
            if isinstance(m, Bottleneck):
                m.conv3[1].weight = nn.Parameter(torch.zeros_like(m.conv3[1].weight))  # BN to zero
            if isinstance(m, nn.Linear): m.weight.data.normal_(0, 0.01)

    def _make_layer(self, block, planes, blocks, stride=1, use_se=True, anti_alias_layer=None):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            layers = []
            if stride == 2:
                # avg pooling before 1x1 conv
                layers.append(
                    nn.AvgPool2d(kernel_size=2, stride=2, ceil_mode=True, count_include_pad=False))
            layers += [
                conv2d_ABN(self.inplanes, planes * block.expansion, kernel_size=1, stride=1,
                           activation="identity")]
            downsample = nn.Sequential(*layers)

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample, use_se=use_se,
                            anti_alias_layer=anti_alias_layer))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks): layers.append(
            block(self.inplanes, planes, use_se=use_se, anti_alias_layer=anti_alias_layer))
        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.body(x)
        self.embeddings = self.global_pool(x)

        logits = self.head(self.embeddings)

        return logits


def TResnetL_V2(model_params):
    """Constructs a large TResnet model.
    """
    in_chans = 3
    num_classes = model_params['num_classes']
    remove_model_jit = False


    layers_list = [3, 4, 23, 3]
    width_factor = 1.0
    global attn_layer
    global stem_type
    global layer_1_2_type
    attn_layer = SEModule
    stem_type = SpaceToDepthModule
    layer_1_2_type = Bottleneck
    model = TResNetV2(layers=layers_list, num_classes=num_classes, in_chans=in_chans,
                      width_factor=width_factor, remove_model_jit=remove_model_jit)
    return model