Spaces:
Runtime error
Runtime error
| """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 | |