xdecoder/Instruct-X-Decoder
163
1# --------------------------------------------------------2# FocalNet for Semantic Segmentation3# Copyright (c) 2022 Microsoft4# Licensed under The MIT License [see LICENSE for details]5# Written by Jianwei Yang6# --------------------------------------------------------7import math8import time9import numpy as np10import logging11import torch12import torch.nn as nn13import torch.nn.functional as F14import torch.utils.checkpoint as checkpoint15from timm.models.layers import DropPath, to_2tuple, trunc_normal_16 17from detectron2.utils.file_io import PathManager18from detectron2.modeling import BACKBONE_REGISTRY, Backbone, ShapeSpec19 20from .registry import register_backbone21 22logger = logging.getLogger(__name__)23 24class Mlp(nn.Module):25 """ Multilayer perceptron."""26 27 def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):28 super().__init__()29 out_features = out_features or in_features30 hidden_features = hidden_features or in_features31 self.fc1 = nn.Linear(in_features, hidden_features)32 self.act = act_layer()33 self.fc2 = nn.Linear(hidden_features, out_features)34 self.drop = nn.Dropout(drop)35 36 def forward(self, x):37 x = self.fc1(x)38 x = self.act(x)39 x = self.drop(x)40 x = self.fc2(x)41 x = self.drop(x)42 return x43 44class FocalModulation(nn.Module):45 """ Focal Modulation46 47 Args:48 dim (int): Number of input channels.49 proj_drop (float, optional): Dropout ratio of output. Default: 0.050 focal_level (int): Number of focal levels51 focal_window (int): Focal window size at focal level 152 focal_factor (int, default=2): Step to increase the focal window53 use_postln (bool, default=False): Whether use post-modulation layernorm54 """55 56 def __init__(self, dim, proj_drop=0., focal_level=2, focal_window=7, focal_factor=2, use_postln=False, use_postln_in_modulation=False, scaling_modulator=False):57 58 super().__init__()59 self.dim = dim60 61 # specific args for focalv362 self.focal_level = focal_level63 self.focal_window = focal_window64 self.focal_factor = focal_factor65 self.use_postln_in_modulation = use_postln_in_modulation66 self.scaling_modulator = scaling_modulator67 68 self.f = nn.Linear(dim, 2*dim+(self.focal_level+1), bias=True)69 self.h = nn.Conv2d(dim, dim, kernel_size=1, stride=1, padding=0, groups=1, bias=True)70 71 self.act = nn.GELU()72 self.proj = nn.Linear(dim, dim)73 self.proj_drop = nn.Dropout(proj_drop)74 self.focal_layers = nn.ModuleList()75 76 if self.use_postln_in_modulation:77 self.ln = nn.LayerNorm(dim)78 79 for k in range(self.focal_level):80 kernel_size = self.focal_factor*k + self.focal_window81 self.focal_layers.append(82 nn.Sequential(83 nn.Conv2d(dim, dim, kernel_size=kernel_size, stride=1, groups=dim, 84 padding=kernel_size//2, bias=False),85 nn.GELU(),86 )87 )88 89 def forward(self, x):90 """ Forward function.91 92 Args:93 x: input features with shape of (B, H, W, C)94 """95 B, nH, nW, C = x.shape96 x = self.f(x)97 x = x.permute(0, 3, 1, 2).contiguous()98 q, ctx, gates = torch.split(x, (C, C, self.focal_level+1), 1)99 100 ctx_all = 0101 for l in range(self.focal_level): 102 ctx = self.focal_layers[l](ctx)103 ctx_all = ctx_all + ctx*gates[:, l:l+1]104 ctx_global = self.act(ctx.mean(2, keepdim=True).mean(3, keepdim=True))105 ctx_all = ctx_all + ctx_global*gates[:,self.focal_level:]106 107 if self.scaling_modulator:108 ctx_all = ctx_all / (self.focal_level + 1)109 110 x_out = q * self.h(ctx_all)111 x_out = x_out.permute(0, 2, 3, 1).contiguous()112 if self.use_postln_in_modulation:113 x_out = self.ln(x_out) 114 x_out = self.proj(x_out)115 x_out = self.proj_drop(x_out)116 return x_out117 118class FocalModulationBlock(nn.Module):119 """ Focal Modulation Block.120 121 Args:122 dim (int): Number of input channels.123 mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.124 drop (float, optional): Dropout rate. Default: 0.0125 drop_path (float, optional): Stochastic depth rate. Default: 0.0126 act_layer (nn.Module, optional): Activation layer. Default: nn.GELU127 norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm128 focal_level (int): number of focal levels129 focal_window (int): focal kernel size at level 1130 """131 132 def __init__(self, dim, mlp_ratio=4., drop=0., drop_path=0., 133 act_layer=nn.GELU, norm_layer=nn.LayerNorm,134 focal_level=2, focal_window=9, 135 use_postln=False, use_postln_in_modulation=False,136 scaling_modulator=False, 137 use_layerscale=False, 138 layerscale_value=1e-4):139 super().__init__()140 self.dim = dim141 self.mlp_ratio = mlp_ratio142 self.focal_window = focal_window143 self.focal_level = focal_level144 self.use_postln = use_postln145 self.use_layerscale = use_layerscale146 147 self.dw1 = nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, groups=dim)148 self.norm1 = norm_layer(dim)149 self.modulation = FocalModulation(150 dim, focal_window=self.focal_window, focal_level=self.focal_level, proj_drop=drop, use_postln_in_modulation=use_postln_in_modulation, scaling_modulator=scaling_modulator151 ) 152 153 self.dw2 = nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, groups=dim)154 self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()155 self.norm2 = norm_layer(dim)156 mlp_hidden_dim = int(dim * mlp_ratio)157 self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)158 159 self.H = None160 self.W = None161 162 self.gamma_1 = 1.0163 self.gamma_2 = 1.0164 if self.use_layerscale:165 self.gamma_1 = nn.Parameter(layerscale_value * torch.ones((dim)), requires_grad=True)166 self.gamma_2 = nn.Parameter(layerscale_value * torch.ones((dim)), requires_grad=True)167 168 def forward(self, x):169 """ Forward function.170 171 Args:172 x: Input feature, tensor size (B, H*W, C).173 H, W: Spatial resolution of the input feature.174 """175 B, L, C = x.shape176 H, W = self.H, self.W177 assert L == H * W, "input feature has wrong size"178 179 x = x.view(B, H, W, C).permute(0, 3, 1, 2).contiguous()180 x = x + self.dw1(x)181 x = x.permute(0, 2, 3, 1).contiguous().view(B, L, C)182 183 shortcut = x184 if not self.use_postln:185 x = self.norm1(x)186 x = x.view(B, H, W, C)187 188 # FM189 x = self.modulation(x).view(B, H * W, C)190 x = shortcut + self.drop_path(self.gamma_1 * x)191 if self.use_postln:192 x = self.norm1(x)193 194 x = x.view(B, H, W, C).permute(0, 3, 1, 2).contiguous()195 x = x + self.dw2(x)196 x = x.permute(0, 2, 3, 1).contiguous().view(B, L, C)197 198 if not self.use_postln:199 x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x))) 200 else:201 x = x + self.drop_path(self.gamma_2 * self.mlp(x))202 x = self.norm2(x)203 204 return x205 206class BasicLayer(nn.Module):207 """ A basic focal modulation layer for one stage.208 209 Args:210 dim (int): Number of feature channels211 depth (int): Depths of this stage.212 mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.213 drop (float, optional): Dropout rate. Default: 0.0214 drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0215 norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm216 downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None217 focal_level (int): Number of focal levels218 focal_window (int): Focal window size at focal level 1219 use_conv_embed (bool): Use overlapped convolution for patch embedding or now. Default: False220 use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False221 """222 223 def __init__(self,224 dim,225 depth,226 mlp_ratio=4.,227 drop=0.,228 drop_path=0.,229 norm_layer=nn.LayerNorm,230 downsample=None,231 focal_window=9, 232 focal_level=2, 233 use_conv_embed=False, 234 use_postln=False, 235 use_postln_in_modulation=False, 236 scaling_modulator=False,237 use_layerscale=False, 238 use_checkpoint=False, 239 use_pre_norm=False, 240 ):241 super().__init__()242 self.depth = depth243 self.use_checkpoint = use_checkpoint244 245 # build blocks246 self.blocks = nn.ModuleList([247 FocalModulationBlock(248 dim=dim,249 mlp_ratio=mlp_ratio,250 drop=drop,251 drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,252 focal_window=focal_window, 253 focal_level=focal_level, 254 use_postln=use_postln, 255 use_postln_in_modulation=use_postln_in_modulation, 256 scaling_modulator=scaling_modulator,257 use_layerscale=use_layerscale, 258 norm_layer=norm_layer)259 for i in range(depth)])260 261 # patch merging layer262 if downsample is not None:263 self.downsample = downsample(264 patch_size=2,265 in_chans=dim, embed_dim=2*dim, 266 use_conv_embed=use_conv_embed, 267 norm_layer=norm_layer, 268 is_stem=False, 269 use_pre_norm=use_pre_norm270 )271 272 else:273 self.downsample = None274 275 def forward(self, x, H, W):276 """ Forward function.277 278 Args:279 x: Input feature, tensor size (B, H*W, C).280 H, W: Spatial resolution of the input feature.281 """282 for blk in self.blocks:283 blk.H, blk.W = H, W284 if self.use_checkpoint:285 x = checkpoint.checkpoint(blk, x)286 else:287 x = blk(x)288 if self.downsample is not None:289 x_reshaped = x.transpose(1, 2).view(x.shape[0], x.shape[-1], H, W)290 x_down = self.downsample(x_reshaped) 291 x_down = x_down.flatten(2).transpose(1, 2) 292 Wh, Ww = (H + 1) // 2, (W + 1) // 2293 return x, H, W, x_down, Wh, Ww294 else:295 return x, H, W, x, H, W296 297 298# class PatchEmbed(nn.Module):299# r""" Image to Patch Embedding300 301# Args:302# img_size (int): Image size. Default: 224.303# patch_size (int): Patch token size. Default: 4.304# in_chans (int): Number of input image channels. Default: 3.305# embed_dim (int): Number of linear projection output channels. Default: 96.306# norm_layer (nn.Module, optional): Normalization layer. Default: None307# """308 309# def __init__(self, img_size=(224, 224), patch_size=4, in_chans=3, embed_dim=96, 310# use_conv_embed=False, norm_layer=None, is_stem=False, use_pre_norm=False):311# super().__init__()312# patch_size = to_2tuple(patch_size)313# patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]314# self.img_size = img_size315# self.patch_size = patch_size316# self.patches_resolution = patches_resolution317# self.num_patches = patches_resolution[0] * patches_resolution[1]318 319# self.in_chans = in_chans320# self.embed_dim = embed_dim321# self.use_pre_norm = use_pre_norm322 323# if use_conv_embed:324# # if we choose to use conv embedding, then we treat the stem and non-stem differently325# if is_stem:326# kernel_size = 7; padding = 3; stride = 4327# else:328# kernel_size = 3; padding = 1; stride = 2329# self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding)330# else:331# self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)332 333# if self.use_pre_norm:334# if norm_layer is not None:335# self.norm = norm_layer(in_chans)336# else:337# self.norm = None338# else:339# if norm_layer is not None:340# self.norm = norm_layer(embed_dim)341# else:342# self.norm = None343 344# def forward(self, x):345# B, C, H, W = x.shape346# # FIXME look at relaxing size constraints347# assert H == self.img_size[0] and W == self.img_size[1], \348# f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."349 350# if self.use_pre_norm:351# if self.norm is not None:352# x = x.flatten(2).transpose(1, 2) # B Ph*Pw C353# x = self.norm(x).transpose(1, 2).view(B, C, H, W)354# x = self.proj(x).flatten(2).transpose(1, 2)355# else:356# x = self.proj(x).flatten(2).transpose(1, 2) # B Ph*Pw C357# if self.norm is not None:358# x = self.norm(x)359# return x360 361# def flops(self):362# Ho, Wo = self.patches_resolution363# flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])364# if self.norm is not None:365# flops += Ho * Wo * self.embed_dim366# return flops367 368class PatchEmbed(nn.Module):369 """ Image to Patch Embedding370 371 Args:372 patch_size (int): Patch token size. Default: 4.373 in_chans (int): Number of input image channels. Default: 3.374 embed_dim (int): Number of linear projection output channels. Default: 96.375 norm_layer (nn.Module, optional): Normalization layer. Default: None376 use_conv_embed (bool): Whether use overlapped convolution for patch embedding. Default: False377 is_stem (bool): Is the stem block or not. 378 """379 380 def __init__(self, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None, use_conv_embed=False, is_stem=False, use_pre_norm=False):381 super().__init__()382 patch_size = to_2tuple(patch_size)383 self.patch_size = patch_size384 385 self.in_chans = in_chans386 self.embed_dim = embed_dim387 self.use_pre_norm = use_pre_norm388 389 if use_conv_embed:390 # if we choose to use conv embedding, then we treat the stem and non-stem differently391 if is_stem:392 kernel_size = 7; padding = 3; stride = 4393 else:394 kernel_size = 3; padding = 1; stride = 2395 self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding) 396 else:397 self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)398 399 if self.use_pre_norm:400 if norm_layer is not None:401 self.norm = norm_layer(in_chans)402 else:403 self.norm = None 404 else:405 if norm_layer is not None:406 self.norm = norm_layer(embed_dim)407 else:408 self.norm = None409 410 def forward(self, x):411 """Forward function."""412 B, C, H, W = x.size()413 if W % self.patch_size[1] != 0:414 x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1]))415 if H % self.patch_size[0] != 0:416 x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0]))417 418 if self.use_pre_norm:419 if self.norm is not None:420 x = x.flatten(2).transpose(1, 2) # B Ph*Pw C421 x = self.norm(x).transpose(1, 2).view(B, C, H, W)422 x = self.proj(x)423 else:424 x = self.proj(x) # B C Wh Ww425 if self.norm is not None:426 Wh, Ww = x.size(2), x.size(3)427 x = x.flatten(2).transpose(1, 2)428 x = self.norm(x)429 x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww)430 431 return x432 433 434class FocalNet(nn.Module):435 """ FocalNet backbone.436 437 Args:438 pretrain_img_size (int): Input image size for training the pretrained model,439 used in absolute postion embedding. Default 224.440 patch_size (int | tuple(int)): Patch size. Default: 4.441 in_chans (int): Number of input image channels. Default: 3.442 embed_dim (int): Number of linear projection output channels. Default: 96.443 depths (tuple[int]): Depths of each Swin Transformer stage.444 mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.445 drop_rate (float): Dropout rate.446 drop_path_rate (float): Stochastic depth rate. Default: 0.2.447 norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.448 patch_norm (bool): If True, add normalization after patch embedding. Default: True.449 out_indices (Sequence[int]): Output from which stages.450 frozen_stages (int): Stages to be frozen (stop grad and set eval mode).451 -1 means not freezing any parameters.452 focal_levels (Sequence[int]): Number of focal levels at four stages453 focal_windows (Sequence[int]): Focal window sizes at first focal level at four stages454 use_conv_embed (bool): Whether use overlapped convolution for patch embedding455 use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.456 """457 458 def __init__(self,459 pretrain_img_size=1600,460 patch_size=4,461 in_chans=3,462 embed_dim=96,463 depths=[2, 2, 6, 2],464 mlp_ratio=4.,465 drop_rate=0.,466 drop_path_rate=0.2,467 norm_layer=nn.LayerNorm,468 patch_norm=True,469 out_indices=[0, 1, 2, 3],470 frozen_stages=-1,471 focal_levels=[2,2,2,2], 472 focal_windows=[9,9,9,9],473 use_pre_norms=[False, False, False, False], 474 use_conv_embed=False, 475 use_postln=False, 476 use_postln_in_modulation=False, 477 scaling_modulator=False,478 use_layerscale=False, 479 use_checkpoint=False, 480 ):481 super().__init__()482 483 self.pretrain_img_size = pretrain_img_size484 self.num_layers = len(depths)485 self.embed_dim = embed_dim486 self.patch_norm = patch_norm487 self.out_indices = out_indices488 self.frozen_stages = frozen_stages489 490 # split image into non-overlapping patches491 self.patch_embed = PatchEmbed(492 patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim,493 norm_layer=norm_layer if self.patch_norm else None, 494 use_conv_embed=use_conv_embed, is_stem=True, use_pre_norm=False)495 496 self.pos_drop = nn.Dropout(p=drop_rate)497 498 # stochastic depth499 dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule500 501 # build layers502 self.layers = nn.ModuleList()503 for i_layer in range(self.num_layers):504 layer = BasicLayer(505 dim=int(embed_dim * 2 ** i_layer),506 depth=depths[i_layer],507 mlp_ratio=mlp_ratio,508 drop=drop_rate,509 drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],510 norm_layer=norm_layer,511 downsample=PatchEmbed if (i_layer < self.num_layers - 1) else None,512 focal_window=focal_windows[i_layer], 513 focal_level=focal_levels[i_layer], 514 use_pre_norm=use_pre_norms[i_layer], 515 use_conv_embed=use_conv_embed,516 use_postln=use_postln, 517 use_postln_in_modulation=use_postln_in_modulation,518 scaling_modulator=scaling_modulator,519 use_layerscale=use_layerscale, 520 use_checkpoint=use_checkpoint)521 self.layers.append(layer)522 523 num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)]524 self.num_features = num_features 525 # self.norm = norm_layer(num_features[-1])526 527 # add a norm layer for each output528 for i_layer in self.out_indices:529 layer = norm_layer(num_features[i_layer])530 layer_name = f'norm{i_layer}'531 self.add_module(layer_name, layer)532 533 self._freeze_stages()534 535 def _freeze_stages(self):536 if self.frozen_stages >= 0:537 self.patch_embed.eval()538 for param in self.patch_embed.parameters():539 param.requires_grad = False540 541 if self.frozen_stages >= 2:542 self.pos_drop.eval()543 for i in range(0, self.frozen_stages - 1):544 m = self.layers[i]545 m.eval()546 for param in m.parameters():547 param.requires_grad = False548 549 def init_weights(self, pretrained=None):550 """Initialize the weights in backbone.551 552 Args:553 pretrained (str, optional): Path to pre-trained weights.554 Defaults to None.555 """556 557 def _init_weights(m):558 if isinstance(m, nn.Linear):559 trunc_normal_(m.weight, std=.02)560 if isinstance(m, nn.Linear) and m.bias is not None:561 nn.init.constant_(m.bias, 0)562 elif isinstance(m, nn.LayerNorm):563 nn.init.constant_(m.bias, 0)564 nn.init.constant_(m.weight, 1.0)565 566 if isinstance(pretrained, str):567 self.apply(_init_weights)568 logger = get_root_logger()569 load_checkpoint(self, pretrained, strict=False, logger=logger)570 elif pretrained is None:571 self.apply(_init_weights)572 else:573 raise TypeError('pretrained must be a str or None')574 575 def load_weights(self, pretrained_dict=None, pretrained_layers=[], verbose=True):576 model_dict = self.state_dict()577 578 missed_dict = [k for k in model_dict.keys() if k not in pretrained_dict]579 logger.info(f'=> Missed keys {missed_dict}')580 unexpected_dict = [k for k in pretrained_dict.keys() if k not in model_dict]581 logger.info(f'=> Unexpected keys {unexpected_dict}')582 583 pretrained_dict = {584 k: v for k, v in pretrained_dict.items()585 if k in model_dict.keys()586 }587 588 need_init_state_dict = {}589 for k, v in pretrained_dict.items():590 need_init = (591 (592 k.split('.')[0] in pretrained_layers593 or pretrained_layers[0] == '*'594 )595 and 'relative_position_index' not in k596 and 'attn_mask' not in k597 )598 599 if need_init:600 # if verbose:601 # logger.info(f'=> init {k} from {pretrained}')602 603 if ('pool_layers' in k) or ('focal_layers' in k) and v.size() != model_dict[k].size():604 table_pretrained = v605 table_current = model_dict[k]606 fsize1 = table_pretrained.shape[2]607 fsize2 = table_current.shape[2]608 609 # NOTE: different from interpolation used in self-attention, we use padding or clipping for focal conv610 if fsize1 < fsize2:611 table_pretrained_resized = torch.zeros(table_current.shape)612 table_pretrained_resized[:, :, (fsize2-fsize1)//2:-(fsize2-fsize1)//2, (fsize2-fsize1)//2:-(fsize2-fsize1)//2] = table_pretrained613 v = table_pretrained_resized614 elif fsize1 > fsize2:615 table_pretrained_resized = table_pretrained[:, :, (fsize1-fsize2)//2:-(fsize1-fsize2)//2, (fsize1-fsize2)//2:-(fsize1-fsize2)//2]616 v = table_pretrained_resized617 618 619 if ("modulation.f" in k or "pre_conv" in k): 620 table_pretrained = v621 table_current = model_dict[k]622 if table_pretrained.shape != table_current.shape:623 if len(table_pretrained.shape) == 2:624 dim = table_pretrained.shape[1]625 assert table_current.shape[1] == dim626 L1 = table_pretrained.shape[0]627 L2 = table_current.shape[0]628 629 if L1 < L2:630 table_pretrained_resized = torch.zeros(table_current.shape)631 # copy for linear project632 table_pretrained_resized[:2*dim] = table_pretrained[:2*dim]633 # copy for global token gating634 table_pretrained_resized[-1] = table_pretrained[-1]635 # copy for first multiple focal levels636 table_pretrained_resized[2*dim:2*dim+(L1-2*dim-1)] = table_pretrained[2*dim:-1]637 # reassign pretrained weights638 v = table_pretrained_resized639 elif L1 > L2:640 raise NotImplementedError641 elif len(table_pretrained.shape) == 1:642 dim = table_pretrained.shape[0]643 L1 = table_pretrained.shape[0]644 L2 = table_current.shape[0]645 if L1 < L2:646 table_pretrained_resized = torch.zeros(table_current.shape)647 # copy for linear project648 table_pretrained_resized[:dim] = table_pretrained[:dim]649 # copy for global token gating650 table_pretrained_resized[-1] = table_pretrained[-1]651 # copy for first multiple focal levels652 # table_pretrained_resized[dim:2*dim+(L1-2*dim-1)] = table_pretrained[2*dim:-1]653 # reassign pretrained weights654 v = table_pretrained_resized655 elif L1 > L2:656 raise NotImplementedError 657 658 need_init_state_dict[k] = v659 660 self.load_state_dict(need_init_state_dict, strict=False)661 662 663 def forward(self, x):664 """Forward function."""665 tic = time.time()666 x = self.patch_embed(x)667 Wh, Ww = x.size(2), x.size(3)668 669 x = x.flatten(2).transpose(1, 2)670 x = self.pos_drop(x)671 672 outs = {}673 for i in range(self.num_layers):674 layer = self.layers[i]675 x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww)676 if i in self.out_indices:677 norm_layer = getattr(self, f'norm{i}')678 x_out = norm_layer(x_out)679 680 out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()681 outs["res{}".format(i + 2)] = out682 683 if len(self.out_indices) == 0:684 outs["res5"] = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()685 686 toc = time.time()687 return outs688 689 def train(self, mode=True):690 """Convert the model into training mode while keep layers freezed."""691 super(FocalNet, self).train(mode)692 self._freeze_stages()693 694 695class D2FocalNet(FocalNet, Backbone):696 def __init__(self, cfg, input_shape):697 698 pretrain_img_size = cfg['BACKBONE']['FOCAL']['PRETRAIN_IMG_SIZE']699 patch_size = cfg['BACKBONE']['FOCAL']['PATCH_SIZE']700 in_chans = 3701 embed_dim = cfg['BACKBONE']['FOCAL']['EMBED_DIM']702 depths = cfg['BACKBONE']['FOCAL']['DEPTHS']703 mlp_ratio = cfg['BACKBONE']['FOCAL']['MLP_RATIO']704 drop_rate = cfg['BACKBONE']['FOCAL']['DROP_RATE']705 drop_path_rate = cfg['BACKBONE']['FOCAL']['DROP_PATH_RATE']706 norm_layer = nn.LayerNorm707 patch_norm = cfg['BACKBONE']['FOCAL']['PATCH_NORM']708 use_checkpoint = cfg['BACKBONE']['FOCAL']['USE_CHECKPOINT']709 out_indices = cfg['BACKBONE']['FOCAL']['OUT_INDICES']710 scaling_modulator = cfg['BACKBONE']['FOCAL'].get('SCALING_MODULATOR', False)711 712 super().__init__(713 pretrain_img_size,714 patch_size,715 in_chans,716 embed_dim,717 depths,718 mlp_ratio,719 drop_rate,720 drop_path_rate,721 norm_layer,722 patch_norm,723 out_indices,724 focal_levels=cfg['BACKBONE']['FOCAL']['FOCAL_LEVELS'],725 focal_windows=cfg['BACKBONE']['FOCAL']['FOCAL_WINDOWS'], 726 use_conv_embed=cfg['BACKBONE']['FOCAL']['USE_CONV_EMBED'], 727 use_postln=cfg['BACKBONE']['FOCAL']['USE_POSTLN'], 728 use_postln_in_modulation=cfg['BACKBONE']['FOCAL']['USE_POSTLN_IN_MODULATION'], 729 scaling_modulator=scaling_modulator,730 use_layerscale=cfg['BACKBONE']['FOCAL']['USE_LAYERSCALE'], 731 use_checkpoint=use_checkpoint,732 )733 734 self._out_features = cfg['BACKBONE']['FOCAL']['OUT_FEATURES']735 736 self._out_feature_strides = {737 "res2": 4,738 "res3": 8,739 "res4": 16,740 "res5": 32,741 }742 self._out_feature_channels = {743 "res2": self.num_features[0],744 "res3": self.num_features[1],745 "res4": self.num_features[2],746 "res5": self.num_features[3],747 }748 749 def forward(self, x):750 """751 Args:752 x: Tensor of shape (N,C,H,W). H, W must be a multiple of ``self.size_divisibility``.753 Returns:754 dict[str->Tensor]: names and the corresponding features755 """756 assert (757 x.dim() == 4758 ), f"SwinTransformer takes an input of shape (N, C, H, W). Got {x.shape} instead!"759 outputs = {}760 y = super().forward(x)761 for k in y.keys():762 if k in self._out_features:763 outputs[k] = y[k]764 return outputs765 766 def output_shape(self):767 return {768 name: ShapeSpec(769 channels=self._out_feature_channels[name], stride=self._out_feature_strides[name]770 )771 for name in self._out_features772 }773 774 @property775 def size_divisibility(self):776 return 32777 778@register_backbone779def get_focal_backbone(cfg):780 focal = D2FocalNet(cfg['MODEL'], 224) 781 782 if cfg['MODEL']['BACKBONE']['LOAD_PRETRAINED'] is True:783 filename = cfg['MODEL']['BACKBONE']['PRETRAINED']784 logger.info(f'=> init from {filename}')785 with PathManager.open(filename, "rb") as f:786 ckpt = torch.load(f)['model']787 focal.load_weights(ckpt, cfg['MODEL']['BACKBONE']['FOCAL'].get('PRETRAINED_LAYERS', ['*']), cfg['VERBOSE'])788 789 return focal