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
(args)
| 132 | |
| 133 | |
| 134 | def build_backbone(args): |
| 135 | """Useful args: |
| 136 | |
| 137 | - backbone: backbone name |
| 138 | - lr_backbone: |
| 139 | - dilation |
| 140 | - return_interm_indices: available: [0,1,2,3], [1,2,3], [3] |
| 141 | - backbone_freeze_keywords: |
| 142 | - use_checkpoint: for swin only for now |
| 143 | """ |
| 144 | position_embedding = build_position_encoding(args) |
| 145 | train_backbone = args.lr_backbone > 0 |
| 146 | if not train_backbone: |
| 147 | raise ValueError('Please set lr_backbone > 0') |
| 148 | return_interm_indices = args.return_interm_indices |
| 149 | assert return_interm_indices in [[0, 1, 2, 3], [1, 2, 3], [3]] # [1,2,3] |
| 150 | backbone_freeze_keywords = args.backbone_freeze_keywords # None |
| 151 | use_checkpoint = getattr(args, 'use_checkpoint', False) # False |
| 152 | |
| 153 | if args.backbone in ['resnet50', 'resnet101']: |
| 154 | backbone = Backbone(args.backbone, |
| 155 | train_backbone, |
| 156 | args.dilation, |
| 157 | return_interm_indices, |
| 158 | batch_norm=FrozenBatchNorm2d) |
| 159 | bb_num_channels = backbone.num_channels |
| 160 | |
| 161 | elif args.backbone in [ |
| 162 | 'swin_T_224_1k', 'swin_B_224_22k', 'swin_B_384_22k', |
| 163 | 'swin_L_224_22k', 'swin_L_384_22k' |
| 164 | ]: |
| 165 | pretrain_img_size = int(args.backbone.split('_')[-2]) |
| 166 | backbone = build_swin_transformer( |
| 167 | args.backbone, |
| 168 | pretrain_img_size=pretrain_img_size, |
| 169 | out_indices=tuple(return_interm_indices), |
| 170 | dilation=args.dilation, |
| 171 | use_checkpoint=use_checkpoint) |
| 172 | # freeze some layers |
| 173 | if backbone_freeze_keywords is not None: |
| 174 | for name, parameter in backbone.named_parameters(): |
| 175 | for keyword in backbone_freeze_keywords: |
| 176 | if keyword in name: |
| 177 | parameter.requires_grad_(False) |
| 178 | break |
| 179 | pretrained_dir = os.environ.get('pretrain_model_path') |
| 180 | # import pdb |
| 181 | # pdb.set_trace() |
| 182 | PTDICT = { |
| 183 | 'swin_T_224_1k': 'swin_tiny_patch4_window7_224.pth', |
| 184 | 'swin_B_384_22k': 'swin_base_patch4_window12_384.pth', |
| 185 | 'swin_L_384_22k': 'swin_large_patch4_window12_384_22k.pth', |
| 186 | } |
| 187 | pretrainedpath = os.path.join(pretrained_dir, PTDICT[args.backbone]) |
| 188 | checkpoint = torch.load(pretrainedpath, map_location='cpu')['model'] |
| 189 | from collections import OrderedDict |
| 190 | |
| 191 | def key_select_function(keyname): |
no test coverage detected