models/aios/backbones/backbone.py

"""Backbone modules."""
from collections import OrderedDict
import os

import torch
import torch.nn.functional as F
import torchvision
from torch import nn
from torchvision.models._utils import IntermediateLayerGetter
from typing import Dict, List

from util.misc import NestedTensor, clean_state_dict, is_main_process

from ..position_encoding import build_position_encoding
from .swin_transformer import build_swin_transformer


class FrozenBatchNorm2d(torch.nn.Module):
    """BatchNorm2d where the batch statistics and the affine parameters are
    fixed.

    Copy-paste from torchvision.misc.ops with added eps before rqsrt, without
    which any other models than torchvision.models.resnet[18,34,50,101] produce
    nans.
    """
    def __init__(self, n):
        super(FrozenBatchNorm2d, self).__init__()
        self.register_buffer('weight', torch.ones(n))
        self.register_buffer('bias', torch.zeros(n))
        self.register_buffer('running_mean', torch.zeros(n))
        self.register_buffer('running_var', torch.ones(n))

    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
                              missing_keys, unexpected_keys, error_msgs):
        num_batches_tracked_key = prefix + 'num_batches_tracked'
        if num_batches_tracked_key in state_dict:
            del state_dict[num_batches_tracked_key]

        super(FrozenBatchNorm2d,
              self)._load_from_state_dict(state_dict, prefix, local_metadata,
                                          strict, missing_keys,
                                          unexpected_keys, error_msgs)

    def forward(self, x):
        w = self.weight.reshape(1, -1, 1, 1)
        b = self.bias.reshape(1, -1, 1, 1)
        rv = self.running_var.reshape(1, -1, 1, 1)
        rm = self.running_mean.reshape(1, -1, 1, 1)
        eps = 1e-5
        scale = w * (rv + eps).rsqrt()
        bias = b - rm * scale
        return x * scale + bias


class BackboneBase(nn.Module):
    def __init__(self, backbone: nn.Module, train_backbone: bool,
                 num_channels: int, return_interm_indices: list):
        super().__init__()
        for name, parameter in backbone.named_parameters():
            if not train_backbone or 'layer0' not in name and 'layer1' not in name and 'layer2' not in name and 'layer3' not in name and 'layer4' not in name:
                parameter.requires_grad_(False)

        return_layers = {}
        for idx, layer_index in enumerate(return_interm_indices):
            return_layers.update({
                'layer{}'.format(5 - len(return_interm_indices) + idx):
                '{}'.format(layer_index)
            })

        self.body = IntermediateLayerGetter(backbone,
                                            return_layers=return_layers)
        self.num_channels = num_channels

    def forward(self, tensor_list: NestedTensor):
        xs = self.body(tensor_list.tensors)
        out: Dict[str, NestedTensor] = {}
        for name, x in xs.items():
            m = tensor_list.mask
            assert m is not None
            mask = F.interpolate(m[None].float(),
                                 size=x.shape[-2:]).to(torch.bool)[0]
            out[name] = NestedTensor(x, mask)
        return out


class Backbone(BackboneBase):
    """ResNet backbone with frozen BatchNorm."""
    def __init__(
        self,
        name: str,
        train_backbone: bool,
        dilation: bool,
        return_interm_indices: list,
        batch_norm=FrozenBatchNorm2d,
    ):
        if name in ['resnet18', 'resnet34', 'resnet50', 'resnet101']:
            # backbone = getattr(torchvision.models, name)(
            #     replace_stride_with_dilation=[False, False, dilation],
            #     pretrained=is_main_process(), norm_layer=batch_norm)
            backbone = getattr(torchvision.models, name)(
                replace_stride_with_dilation=[False, False, dilation],
                pretrained=False,
                norm_layer=batch_norm)
        else:
            raise NotImplementedError(
                'Why you can get here with name {}'.format(name))

        assert name not in (
            'resnet18',
            'resnet34'), 'Only resnet50 and resnet101 are available.'
        assert return_interm_indices in [[0, 1, 2, 3], [1, 2, 3], [3]]
        num_channels_all = [256, 512, 1024, 2048]
        num_channels = num_channels_all[4 - len(return_interm_indices):]
        super().__init__(backbone, train_backbone, num_channels,
                         return_interm_indices)


class Joiner(nn.Sequential):
    def __init__(self, backbone, position_embedding):
        super().__init__(backbone, position_embedding)

    def forward(self, tensor_list: NestedTensor):
        xs = self[0](tensor_list)
        out: List[NestedTensor] = []
        pos = []
        for name, x in xs.items():
            out.append(x)

            pos.append(self[1](x).to(x.tensors.dtype))

        return out, pos


def build_backbone(args):
    """Useful args:

    - backbone: backbone name
    - lr_backbone:
    - dilation
    - return_interm_indices: available: [0,1,2,3], [1,2,3], [3]
    - backbone_freeze_keywords:
    - use_checkpoint: for swin only for now
    """
    position_embedding = build_position_encoding(args)
    train_backbone = args.lr_backbone > 0
    if not train_backbone:
        raise ValueError('Please set lr_backbone > 0')
    return_interm_indices = args.return_interm_indices
    assert return_interm_indices in [[0, 1, 2, 3], [1, 2, 3], [3]]  #  [1,2,3]
    backbone_freeze_keywords = args.backbone_freeze_keywords  # None
    use_checkpoint = getattr(args, 'use_checkpoint', False)  # False

    if args.backbone in ['resnet50', 'resnet101']:
        backbone = Backbone(args.backbone,
                            train_backbone,
                            args.dilation,
                            return_interm_indices,
                            batch_norm=FrozenBatchNorm2d)
        bb_num_channels = backbone.num_channels

    elif args.backbone in [
            'swin_T_224_1k', 'swin_B_224_22k', 'swin_B_384_22k',
            'swin_L_224_22k', 'swin_L_384_22k'
    ]:
        pretrain_img_size = int(args.backbone.split('_')[-2])
        backbone = build_swin_transformer(
            args.backbone,
            pretrain_img_size=pretrain_img_size,
            out_indices=tuple(return_interm_indices),
            dilation=args.dilation,
            use_checkpoint=use_checkpoint)
        # freeze some layers
        if backbone_freeze_keywords is not None:
            for name, parameter in backbone.named_parameters():
                for keyword in backbone_freeze_keywords:
                    if keyword in name:
                        parameter.requires_grad_(False)
                        break
        pretrained_dir = os.environ.get('pretrain_model_path')
        # import pdb
        # pdb.set_trace()
        PTDICT = {
            'swin_T_224_1k': 'swin_tiny_patch4_window7_224.pth',
            'swin_B_384_22k': 'swin_base_patch4_window12_384.pth',
            'swin_L_384_22k': 'swin_large_patch4_window12_384_22k.pth',
        }
        pretrainedpath = os.path.join(pretrained_dir, PTDICT[args.backbone])
        checkpoint = torch.load(pretrainedpath, map_location='cpu')['model']
        from collections import OrderedDict

        def key_select_function(keyname):
            if 'head' in keyname:
                return False
            if args.dilation and 'layers.3' in keyname:
                return False
            return True

        _tmp_st = OrderedDict({
            k: v
            for k, v in clean_state_dict(checkpoint).items()
            if key_select_function(k)
        })
        _tmp_st_output = backbone.load_state_dict(_tmp_st, strict=False)
        print(str(_tmp_st_output))

        bb_num_channels = backbone.num_features[4 -
                                                len(return_interm_indices):]
    else:
        raise NotImplementedError('Unknown backbone {}'.format(args.backbone))

    assert len(bb_num_channels) == len(
        return_interm_indices
    ), f'len(bb_num_channels) {len(bb_num_channels)} != len(return_interm_indices) {len(return_interm_indices)}'

    model = Joiner(backbone, position_embedding)
    model.num_channels = bb_num_channels
    assert isinstance(
        bb_num_channels,
        List), 'bb_num_channels is expected to be a List but {}'.format(
            type(bb_num_channels))
    # import pdb; pdb.set_trace()
    return model