| |
| |
| |
| import math |
| from os.path import join |
| import numpy as np |
| import copy |
| from functools import partial |
|
|
| import torch |
| from torch import nn |
| import torch.utils.model_zoo as model_zoo |
| import torch.nn.functional as F |
| import fvcore.nn.weight_init as weight_init |
|
|
| from detectron2.modeling.backbone import FPN |
| from detectron2.modeling.backbone.build import BACKBONE_REGISTRY |
| from detectron2.layers.batch_norm import get_norm, FrozenBatchNorm2d |
| from detectron2.modeling.backbone import Backbone |
|
|
| from timm import create_model |
| from timm.models.helpers import build_model_with_cfg |
| from timm.models.registry import register_model |
| from timm.models.resnet import ResNet, Bottleneck |
| from timm.models.resnet import default_cfgs as default_cfgs_resnet |
|
|
|
|
| class CustomResNet(ResNet): |
| def __init__(self, **kwargs): |
| self.out_indices = kwargs.pop('out_indices') |
| super().__init__(**kwargs) |
|
|
|
|
| def forward(self, x): |
| x = self.conv1(x) |
| x = self.bn1(x) |
| x = self.act1(x) |
| x = self.maxpool(x) |
| ret = [x] |
| x = self.layer1(x) |
| ret.append(x) |
| x = self.layer2(x) |
| ret.append(x) |
| x = self.layer3(x) |
| ret.append(x) |
| x = self.layer4(x) |
| ret.append(x) |
| return [ret[i] for i in self.out_indices] |
|
|
|
|
| def load_pretrained(self, cached_file): |
| data = torch.load(cached_file, map_location='cpu') |
| if 'state_dict' in data: |
| self.load_state_dict(data['state_dict']) |
| else: |
| self.load_state_dict(data) |
|
|
|
|
| model_params = { |
| 'resnet50': dict(block=Bottleneck, layers=[3, 4, 6, 3]), |
| 'resnet50_in21k': dict(block=Bottleneck, layers=[3, 4, 6, 3]), |
| } |
|
|
|
|
| def create_timm_resnet(variant, out_indices, pretrained=False, **kwargs): |
| params = model_params[variant] |
| default_cfgs_resnet['resnet50_in21k'] = \ |
| copy.deepcopy(default_cfgs_resnet['resnet50']) |
| default_cfgs_resnet['resnet50_in21k']['url'] = \ |
| 'https://miil-public-eu.oss-eu-central-1.aliyuncs.com/model-zoo/ImageNet_21K_P/models/resnet50_miil_21k.pth' |
| default_cfgs_resnet['resnet50_in21k']['num_classes'] = 11221 |
|
|
| return build_model_with_cfg( |
| CustomResNet, variant, pretrained, |
| default_cfg=default_cfgs_resnet[variant], |
| out_indices=out_indices, |
| pretrained_custom_load=True, |
| **params, |
| **kwargs) |
|
|
|
|
| class LastLevelP6P7_P5(nn.Module): |
| """ |
| """ |
| def __init__(self, in_channels, out_channels): |
| super().__init__() |
| self.num_levels = 2 |
| self.in_feature = "p5" |
| self.p6 = nn.Conv2d(in_channels, out_channels, 3, 2, 1) |
| self.p7 = nn.Conv2d(out_channels, out_channels, 3, 2, 1) |
| for module in [self.p6, self.p7]: |
| weight_init.c2_xavier_fill(module) |
|
|
| def forward(self, c5): |
| p6 = self.p6(c5) |
| p7 = self.p7(F.relu(p6)) |
| return [p6, p7] |
|
|
|
|
| def freeze_module(x): |
| """ |
| """ |
| for p in x.parameters(): |
| p.requires_grad = False |
| FrozenBatchNorm2d.convert_frozen_batchnorm(x) |
| return x |
|
|
|
|
| class TIMM(Backbone): |
| def __init__(self, base_name, out_levels, freeze_at=0, norm='FrozenBN'): |
| super().__init__() |
| out_indices = [x - 1 for x in out_levels] |
| if 'resnet' in base_name: |
| self.base = create_timm_resnet( |
| base_name, out_indices=out_indices, |
| pretrained=False) |
| elif 'eff' in base_name: |
| self.base = create_model( |
| base_name, features_only=True, |
| out_indices=out_indices, pretrained=True) |
| else: |
| assert 0, base_name |
| feature_info = [dict(num_chs=f['num_chs'], reduction=f['reduction']) \ |
| for i, f in enumerate(self.base.feature_info)] |
| self._out_features = ['layer{}'.format(x) for x in out_levels] |
| self._out_feature_channels = { |
| 'layer{}'.format(l): feature_info[l - 1]['num_chs'] for l in out_levels} |
| self._out_feature_strides = { |
| 'layer{}'.format(l): feature_info[l - 1]['reduction'] for l in out_levels} |
| self._size_divisibility = max(self._out_feature_strides.values()) |
| if 'resnet' in base_name: |
| self.freeze(freeze_at) |
| if norm == 'FrozenBN': |
| self = FrozenBatchNorm2d.convert_frozen_batchnorm(self) |
|
|
| def freeze(self, freeze_at=0): |
| """ |
| """ |
| if freeze_at >= 1: |
| print('Frezing', self.base.conv1) |
| self.base.conv1 = freeze_module(self.base.conv1) |
| if freeze_at >= 2: |
| print('Frezing', self.base.layer1) |
| self.base.layer1 = freeze_module(self.base.layer1) |
|
|
| def forward(self, x): |
| features = self.base(x) |
| ret = {k: v for k, v in zip(self._out_features, features)} |
| return ret |
| |
| @property |
| def size_divisibility(self): |
| return self._size_divisibility |
|
|
|
|
| @BACKBONE_REGISTRY.register() |
| def build_timm_backbone(cfg, input_shape): |
| model = TIMM( |
| cfg.MODEL.TIMM.BASE_NAME, |
| cfg.MODEL.TIMM.OUT_LEVELS, |
| freeze_at=cfg.MODEL.TIMM.FREEZE_AT, |
| norm=cfg.MODEL.TIMM.NORM, |
| ) |
| return model |
|
|
|
|
| @BACKBONE_REGISTRY.register() |
| def build_p67_timm_fpn_backbone(cfg, input_shape): |
| """ |
| """ |
| bottom_up = build_timm_backbone(cfg, input_shape) |
| in_features = cfg.MODEL.FPN.IN_FEATURES |
| out_channels = cfg.MODEL.FPN.OUT_CHANNELS |
| backbone = FPN( |
| bottom_up=bottom_up, |
| in_features=in_features, |
| out_channels=out_channels, |
| norm=cfg.MODEL.FPN.NORM, |
| top_block=LastLevelP6P7_P5(out_channels, out_channels), |
| fuse_type=cfg.MODEL.FPN.FUSE_TYPE, |
| ) |
| return backbone |
|
|
| @BACKBONE_REGISTRY.register() |
| def build_p35_timm_fpn_backbone(cfg, input_shape): |
| """ |
| """ |
| bottom_up = build_timm_backbone(cfg, input_shape) |
| |
| in_features = cfg.MODEL.FPN.IN_FEATURES |
| out_channels = cfg.MODEL.FPN.OUT_CHANNELS |
| backbone = FPN( |
| bottom_up=bottom_up, |
| in_features=in_features, |
| out_channels=out_channels, |
| norm=cfg.MODEL.FPN.NORM, |
| top_block=None, |
| fuse_type=cfg.MODEL.FPN.FUSE_TYPE, |
| ) |
| return backbone |
