ZhengPeng7/BiRefNet_lite
21261k
1### config.py2 3import os4import math5from transformers import PretrainedConfig6 7 8class Config(PretrainedConfig):9 def __init__(self) -> None:10 # Compatible with the latest version of transformers.11 # Error source: https://github.com/huggingface/transformers/commit/9568b506ed511c76ab4d0c6ed591c7fce8e048a512 # Previous solution in the users' end: https://github.com/ZhengPeng7/BiRefNet/issues/189#issuecomment-271668868813 super().__init__()14 15 # PATH settings16 self.sys_home_dir = os.path.expanduser('~') # Make up your file system as: SYS_HOME_DIR/codes/dis/BiRefNet, SYS_HOME_DIR/datasets/dis/xx, SYS_HOME_DIR/weights/xx17 18 # TASK settings19 self.task = ['DIS5K', 'COD', 'HRSOD', 'DIS5K+HRSOD+HRS10K', 'P3M-10k'][0]20 self.training_set = {21 'DIS5K': ['DIS-TR', 'DIS-TR+DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4'][0],22 'COD': 'TR-COD10K+TR-CAMO',23 'HRSOD': ['TR-DUTS', 'TR-HRSOD', 'TR-UHRSD', 'TR-DUTS+TR-HRSOD', 'TR-DUTS+TR-UHRSD', 'TR-HRSOD+TR-UHRSD', 'TR-DUTS+TR-HRSOD+TR-UHRSD'][5],24 'DIS5K+HRSOD+HRS10K': 'DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4+DIS-TR+TE-HRS10K+TE-HRSOD+TE-UHRSD+TR-HRS10K+TR-HRSOD+TR-UHRSD', # leave DIS-VD for evaluation.25 'P3M-10k': 'TR-P3M-10k',26 }[self.task]27 self.prompt4loc = ['dense', 'sparse'][0]28 29 # Faster-Training settings30 self.load_all = True31 self.compile = True # 1. Trigger CPU memory leak in some extend, which is an inherent problem of PyTorch.32 # Machines with > 70GB CPU memory can run the whole training on DIS5K with default setting.33 # 2. Higher PyTorch version may fix it: https://github.com/pytorch/pytorch/issues/119607.34 # 3. But compile in Pytorch > 2.0.1 seems to bring no acceleration for training.35 self.precisionHigh = True36 37 # MODEL settings38 self.ms_supervision = True39 self.out_ref = self.ms_supervision and True40 self.dec_ipt = True41 self.dec_ipt_split = True42 self.cxt_num = [0, 3][1] # multi-scale skip connections from encoder43 self.mul_scl_ipt = ['', 'add', 'cat'][2]44 self.dec_att = ['', 'ASPP', 'ASPPDeformable'][2]45 self.squeeze_block = ['', 'BasicDecBlk_x1', 'ResBlk_x4', 'ASPP_x3', 'ASPPDeformable_x3'][1]46 self.dec_blk = ['BasicDecBlk', 'ResBlk', 'HierarAttDecBlk'][0]47 48 # TRAINING settings49 self.batch_size = 450 self.IoU_finetune_last_epochs = [51 0,52 {53 'DIS5K': -50,54 'COD': -20,55 'HRSOD': -20,56 'DIS5K+HRSOD+HRS10K': -20,57 'P3M-10k': -20,58 }[self.task]59 ][1] # choose 0 to skip60 self.lr = (1e-4 if 'DIS5K' in self.task else 1e-5) * math.sqrt(self.batch_size / 4) # DIS needs high lr to converge faster. Adapt the lr linearly61 self.size = 102462 self.num_workers = max(4, self.batch_size) # will be decrease to min(it, batch_size) at the initialization of the data_loader63 64 # Backbone settings65 self.bb = [66 'vgg16', 'vgg16bn', 'resnet50', # 0, 1, 267 'swin_v1_t', 'swin_v1_s', # 3, 468 'swin_v1_b', 'swin_v1_l', # 5-bs9, 6-bs469 'pvt_v2_b0', 'pvt_v2_b1', # 7, 870 'pvt_v2_b2', 'pvt_v2_b5', # 9-bs10, 10-bs571 ][3]72 self.lateral_channels_in_collection = {73 'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64],74 'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64],75 'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192],76 'swin_v1_t': [768, 384, 192, 96], 'swin_v1_s': [768, 384, 192, 96],77 'pvt_v2_b0': [256, 160, 64, 32], 'pvt_v2_b1': [512, 320, 128, 64],78 }[self.bb]79 if self.mul_scl_ipt == 'cat':80 self.lateral_channels_in_collection = [channel * 2 for channel in self.lateral_channels_in_collection]81 self.cxt = self.lateral_channels_in_collection[1:][::-1][-self.cxt_num:] if self.cxt_num else []82 83 # MODEL settings - inactive84 self.lat_blk = ['BasicLatBlk'][0]85 self.dec_channels_inter = ['fixed', 'adap'][0]86 self.refine = ['', 'itself', 'RefUNet', 'Refiner', 'RefinerPVTInChannels4'][0]87 self.progressive_ref = self.refine and True88 self.ender = self.progressive_ref and False89 self.scale = self.progressive_ref and 290 self.auxiliary_classification = False # Only for DIS5K, where class labels are saved in `dataset.py`.91 self.refine_iteration = 192 self.freeze_bb = False93 self.model = [94 'BiRefNet',95 ][0]96 if self.dec_blk == 'HierarAttDecBlk':97 self.batch_size = 2 ** [0, 1, 2, 3, 4][2]98 99 # TRAINING settings - inactive100 self.preproc_methods = ['flip', 'enhance', 'rotate', 'pepper', 'crop'][:4]101 self.optimizer = ['Adam', 'AdamW'][1]102 self.lr_decay_epochs = [1e5] # Set to negative N to decay the lr in the last N-th epoch.103 self.lr_decay_rate = 0.5104 # Loss105 self.lambdas_pix_last = {106 # not 0 means opening this loss107 # original rate -- 1 : 30 : 1.5 : 0.2, bce x 30108 'bce': 30 * 1, # high performance109 'iou': 0.5 * 1, # 0 / 255110 'iou_patch': 0.5 * 0, # 0 / 255, win_size = (64, 64)111 'mse': 150 * 0, # can smooth the saliency map112 'triplet': 3 * 0,113 'reg': 100 * 0,114 'ssim': 10 * 1, # help contours,115 'cnt': 5 * 0, # help contours116 'structure': 5 * 0, # structure loss from codes of MVANet. A little improvement on DIS-TE[1,2,3], a bit more decrease on DIS-TE4.117 }118 self.lambdas_cls = {119 'ce': 5.0120 }121 # Adv122 self.lambda_adv_g = 10. * 0 # turn to 0 to avoid adv training123 self.lambda_adv_d = 3. * (self.lambda_adv_g > 0)124 125 # PATH settings - inactive126 self.data_root_dir = os.path.join(self.sys_home_dir, 'datasets/dis')127 self.weights_root_dir = os.path.join(self.sys_home_dir, 'weights')128 self.weights = {129 'pvt_v2_b2': os.path.join(self.weights_root_dir, 'pvt_v2_b2.pth'),130 'pvt_v2_b5': os.path.join(self.weights_root_dir, ['pvt_v2_b5.pth', 'pvt_v2_b5_22k.pth'][0]),131 'swin_v1_b': os.path.join(self.weights_root_dir, ['swin_base_patch4_window12_384_22kto1k.pth', 'swin_base_patch4_window12_384_22k.pth'][0]),132 'swin_v1_l': os.path.join(self.weights_root_dir, ['swin_large_patch4_window12_384_22kto1k.pth', 'swin_large_patch4_window12_384_22k.pth'][0]),133 'swin_v1_t': os.path.join(self.weights_root_dir, ['swin_tiny_patch4_window7_224_22kto1k_finetune.pth'][0]),134 'swin_v1_s': os.path.join(self.weights_root_dir, ['swin_small_patch4_window7_224_22kto1k_finetune.pth'][0]),135 'pvt_v2_b0': os.path.join(self.weights_root_dir, ['pvt_v2_b0.pth'][0]),136 'pvt_v2_b1': os.path.join(self.weights_root_dir, ['pvt_v2_b1.pth'][0]),137 }138 139 # Callbacks - inactive140 self.verbose_eval = True141 self.only_S_MAE = False142 self.use_fp16 = False # Bugs. It may cause nan in training.143 self.SDPA_enabled = False # Bugs. Slower and errors occur in multi-GPUs144 145 # others146 self.device = [0, 'cpu'][0] # .to(0) == .to('cuda:0')147 148 self.batch_size_valid = 1149 self.rand_seed = 7150 # run_sh_file = [f for f in os.listdir('.') if 'train.sh' == f] + [os.path.join('..', f) for f in os.listdir('..') if 'train.sh' == f]151 # with open(run_sh_file[0], 'r') as f:152 # lines = f.readlines()153 # self.save_last = int([l.strip() for l in lines if '"{}")'.format(self.task) in l and 'val_last=' in l][0].split('val_last=')[-1].split()[0])154 # self.save_step = int([l.strip() for l in lines if '"{}")'.format(self.task) in l and 'step=' in l][0].split('step=')[-1].split()[0])155 # self.val_step = [0, self.save_step][0]156 157 def print_task(self) -> None:158 # Return task for choosing settings in shell scripts.159 print(self.task)160 161 162 163### models/backbones/pvt_v2.py164 165import torch166import torch.nn as nn167from functools import partial168 169from timm.layers import DropPath, to_2tuple, trunc_normal_170 171 172import math173 174# from config import Config175 176# config = Config()177 178class Mlp(nn.Module):179 def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):180 super().__init__()181 out_features = out_features or in_features182 hidden_features = hidden_features or in_features183 self.fc1 = nn.Linear(in_features, hidden_features)184 self.dwconv = DWConv(hidden_features)185 self.act = act_layer()186 self.fc2 = nn.Linear(hidden_features, out_features)187 self.drop = nn.Dropout(drop)188 189 self.apply(self._init_weights)190 191 def _init_weights(self, m):192 if isinstance(m, nn.Linear):193 trunc_normal_(m.weight, std=.02)194 if isinstance(m, nn.Linear) and m.bias is not None:195 nn.init.constant_(m.bias, 0)196 elif isinstance(m, nn.LayerNorm):197 nn.init.constant_(m.bias, 0)198 nn.init.constant_(m.weight, 1.0)199 elif isinstance(m, nn.Conv2d):200 fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels201 fan_out //= m.groups202 m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))203 if m.bias is not None:204 m.bias.data.zero_()205 206 def forward(self, x, H, W):207 x = self.fc1(x)208 x = self.dwconv(x, H, W)209 x = self.act(x)210 x = self.drop(x)211 x = self.fc2(x)212 x = self.drop(x)213 return x214 215 216class Attention(nn.Module):217 def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., sr_ratio=1):218 super().__init__()219 assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}."220 221 self.dim = dim222 self.num_heads = num_heads223 head_dim = dim // num_heads224 self.scale = qk_scale or head_dim ** -0.5225 226 self.q = nn.Linear(dim, dim, bias=qkv_bias)227 self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias)228 self.attn_drop_prob = attn_drop229 self.attn_drop = nn.Dropout(attn_drop)230 self.proj = nn.Linear(dim, dim)231 self.proj_drop = nn.Dropout(proj_drop)232 233 self.sr_ratio = sr_ratio234 if sr_ratio > 1:235 self.sr = nn.Conv2d(dim, dim, kernel_size=sr_ratio, stride=sr_ratio)236 self.norm = nn.LayerNorm(dim)237 238 self.apply(self._init_weights)239 240 def _init_weights(self, m):241 if isinstance(m, nn.Linear):242 trunc_normal_(m.weight, std=.02)243 if isinstance(m, nn.Linear) and m.bias is not None:244 nn.init.constant_(m.bias, 0)245 elif isinstance(m, nn.LayerNorm):246 nn.init.constant_(m.bias, 0)247 nn.init.constant_(m.weight, 1.0)248 elif isinstance(m, nn.Conv2d):249 fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels250 fan_out //= m.groups251 m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))252 if m.bias is not None:253 m.bias.data.zero_()254 255 def forward(self, x, H, W):256 B, N, C = x.shape257 q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)258 259 if self.sr_ratio > 1:260 x_ = x.permute(0, 2, 1).reshape(B, C, H, W)261 x_ = self.sr(x_).reshape(B, C, -1).permute(0, 2, 1)262 x_ = self.norm(x_)263 kv = self.kv(x_).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)264 else:265 kv = self.kv(x).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)266 k, v = kv[0], kv[1]267 268 if config.SDPA_enabled:269 x = torch.nn.functional.scaled_dot_product_attention(270 q, k, v,271 attn_mask=None, dropout_p=self.attn_drop_prob, is_causal=False272 ).transpose(1, 2).reshape(B, N, C)273 else:274 attn = (q @ k.transpose(-2, -1)) * self.scale275 attn = attn.softmax(dim=-1)276 attn = self.attn_drop(attn)277 278 x = (attn @ v).transpose(1, 2).reshape(B, N, C)279 x = self.proj(x)280 x = self.proj_drop(x)281 282 return x283 284 285class Block(nn.Module):286 287 def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,288 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, sr_ratio=1):289 super().__init__()290 self.norm1 = norm_layer(dim)291 self.attn = Attention(292 dim,293 num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,294 attn_drop=attn_drop, proj_drop=drop, sr_ratio=sr_ratio)295 # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here296 self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()297 self.norm2 = norm_layer(dim)298 mlp_hidden_dim = int(dim * mlp_ratio)299 self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)300 301 self.apply(self._init_weights)302 303 def _init_weights(self, m):304 if isinstance(m, nn.Linear):305 trunc_normal_(m.weight, std=.02)306 if isinstance(m, nn.Linear) and m.bias is not None:307 nn.init.constant_(m.bias, 0)308 elif isinstance(m, nn.LayerNorm):309 nn.init.constant_(m.bias, 0)310 nn.init.constant_(m.weight, 1.0)311 elif isinstance(m, nn.Conv2d):312 fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels313 fan_out //= m.groups314 m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))315 if m.bias is not None:316 m.bias.data.zero_()317 318 def forward(self, x, H, W):319 x = x + self.drop_path(self.attn(self.norm1(x), H, W))320 x = x + self.drop_path(self.mlp(self.norm2(x), H, W))321 322 return x323 324 325class OverlapPatchEmbed(nn.Module):326 """ Image to Patch Embedding327 """328 329 def __init__(self, img_size=224, patch_size=7, stride=4, in_channels=3, embed_dim=768):330 super().__init__()331 img_size = to_2tuple(img_size)332 patch_size = to_2tuple(patch_size)333 334 self.img_size = img_size335 self.patch_size = patch_size336 self.H, self.W = img_size[0] // patch_size[0], img_size[1] // patch_size[1]337 self.num_patches = self.H * self.W338 self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=stride,339 padding=(patch_size[0] // 2, patch_size[1] // 2))340 self.norm = nn.LayerNorm(embed_dim)341 342 self.apply(self._init_weights)343 344 def _init_weights(self, m):345 if isinstance(m, nn.Linear):346 trunc_normal_(m.weight, std=.02)347 if isinstance(m, nn.Linear) and m.bias is not None:348 nn.init.constant_(m.bias, 0)349 elif isinstance(m, nn.LayerNorm):350 nn.init.constant_(m.bias, 0)351 nn.init.constant_(m.weight, 1.0)352 elif isinstance(m, nn.Conv2d):353 fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels354 fan_out //= m.groups355 m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))356 if m.bias is not None:357 m.bias.data.zero_()358 359 def forward(self, x):360 x = self.proj(x)361 _, _, H, W = x.shape362 x = x.flatten(2).transpose(1, 2)363 x = self.norm(x)364 365 return x, H, W366 367 368class PyramidVisionTransformerImpr(nn.Module):369 def __init__(self, img_size=224, patch_size=16, in_channels=3, num_classes=1000, embed_dims=[64, 128, 256, 512],370 num_heads=[1, 2, 4, 8], mlp_ratios=[4, 4, 4, 4], qkv_bias=False, qk_scale=None, drop_rate=0.,371 attn_drop_rate=0., drop_path_rate=0., norm_layer=nn.LayerNorm,372 depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1]):373 super().__init__()374 self.num_classes = num_classes375 self.depths = depths376 377 # patch_embed378 self.patch_embed1 = OverlapPatchEmbed(img_size=img_size, patch_size=7, stride=4, in_channels=in_channels,379 embed_dim=embed_dims[0])380 self.patch_embed2 = OverlapPatchEmbed(img_size=img_size // 4, patch_size=3, stride=2, in_channels=embed_dims[0],381 embed_dim=embed_dims[1])382 self.patch_embed3 = OverlapPatchEmbed(img_size=img_size // 8, patch_size=3, stride=2, in_channels=embed_dims[1],383 embed_dim=embed_dims[2])384 self.patch_embed4 = OverlapPatchEmbed(img_size=img_size // 16, patch_size=3, stride=2, in_channels=embed_dims[2],385 embed_dim=embed_dims[3])386 387 # transformer encoder388 dpr = np.linspace(0, drop_path_rate, sum(depths)).tolist() # stochastic depth decay rule389 cur = 0390 self.block1 = nn.ModuleList([Block(391 dim=embed_dims[0], num_heads=num_heads[0], mlp_ratio=mlp_ratios[0], qkv_bias=qkv_bias, qk_scale=qk_scale,392 drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer,393 sr_ratio=sr_ratios[0])394 for i in range(depths[0])])395 self.norm1 = norm_layer(embed_dims[0])396 397 cur += depths[0]398 self.block2 = nn.ModuleList([Block(399 dim=embed_dims[1], num_heads=num_heads[1], mlp_ratio=mlp_ratios[1], qkv_bias=qkv_bias, qk_scale=qk_scale,400 drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer,401 sr_ratio=sr_ratios[1])402 for i in range(depths[1])])403 self.norm2 = norm_layer(embed_dims[1])404 405 cur += depths[1]406 self.block3 = nn.ModuleList([Block(407 dim=embed_dims[2], num_heads=num_heads[2], mlp_ratio=mlp_ratios[2], qkv_bias=qkv_bias, qk_scale=qk_scale,408 drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer,409 sr_ratio=sr_ratios[2])410 for i in range(depths[2])])411 self.norm3 = norm_layer(embed_dims[2])412 413 cur += depths[2]414 self.block4 = nn.ModuleList([Block(415 dim=embed_dims[3], num_heads=num_heads[3], mlp_ratio=mlp_ratios[3], qkv_bias=qkv_bias, qk_scale=qk_scale,416 drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer,417 sr_ratio=sr_ratios[3])418 for i in range(depths[3])])419 self.norm4 = norm_layer(embed_dims[3])420 421 # classification head422 # self.head = nn.Linear(embed_dims[3], num_classes) if num_classes > 0 else nn.Identity()423 424 self.apply(self._init_weights)425 426 def _init_weights(self, m):427 if isinstance(m, nn.Linear):428 trunc_normal_(m.weight, std=.02)429 if isinstance(m, nn.Linear) and m.bias is not None:430 nn.init.constant_(m.bias, 0)431 elif isinstance(m, nn.LayerNorm):432 nn.init.constant_(m.bias, 0)433 nn.init.constant_(m.weight, 1.0)434 elif isinstance(m, nn.Conv2d):435 fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels436 fan_out //= m.groups437 m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))438 if m.bias is not None:439 m.bias.data.zero_()440 441 def init_weights(self, pretrained=None):442 if isinstance(pretrained, str):443 logger = 1444 #load_checkpoint(self, pretrained, map_location='cpu', strict=False, logger=logger)445 446 def reset_drop_path(self, drop_path_rate):447 dpr = np.linspace(0, drop_path_rate, sum(self.depths)).tolist()448 cur = 0449 for i in range(self.depths[0]):450 self.block1[i].drop_path.drop_prob = dpr[cur + i]451 452 cur += self.depths[0]453 for i in range(self.depths[1]):454 self.block2[i].drop_path.drop_prob = dpr[cur + i]455 456 cur += self.depths[1]457 for i in range(self.depths[2]):458 self.block3[i].drop_path.drop_prob = dpr[cur + i]459 460 cur += self.depths[2]461 for i in range(self.depths[3]):462 self.block4[i].drop_path.drop_prob = dpr[cur + i]463 464 def freeze_patch_emb(self):465 self.patch_embed1.requires_grad = False466 467 @torch.jit.ignore468 def no_weight_decay(self):469 return {'pos_embed1', 'pos_embed2', 'pos_embed3', 'pos_embed4', 'cls_token'} # has pos_embed may be better470 471 def get_classifier(self):472 return self.head473 474 def reset_classifier(self, num_classes, global_pool=''):475 self.num_classes = num_classes476 self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()477 478 def forward_features(self, x):479 B = x.shape[0]480 outs = []481 482 # stage 1483 x, H, W = self.patch_embed1(x)484 for i, blk in enumerate(self.block1):485 x = blk(x, H, W)486 x = self.norm1(x)487 x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()488 outs.append(x)489 490 # stage 2491 x, H, W = self.patch_embed2(x)492 for i, blk in enumerate(self.block2):493 x = blk(x, H, W)494 x = self.norm2(x)495 x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()496 outs.append(x)497 498 # stage 3499 x, H, W = self.patch_embed3(x)500 for i, blk in enumerate(self.block3):501 x = blk(x, H, W)502 x = self.norm3(x)503 x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()504 outs.append(x)505 506 # stage 4507 x, H, W = self.patch_embed4(x)508 for i, blk in enumerate(self.block4):509 x = blk(x, H, W)510 x = self.norm4(x)511 x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()512 outs.append(x)513 514 return outs515 516 # return x.mean(dim=1)517 518 def forward(self, x):519 x = self.forward_features(x)520 # x = self.head(x)521 522 return x523 524 525class DWConv(nn.Module):526 def __init__(self, dim=768):527 super(DWConv, self).__init__()528 self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim)529 530 def forward(self, x, H, W):531 B, N, C = x.shape532 x = x.transpose(1, 2).view(B, C, H, W).contiguous()533 x = self.dwconv(x)534 x = x.flatten(2).transpose(1, 2)535 536 return x537 538 539def _conv_filter(state_dict, patch_size=16):540 """ convert patch embedding weight from manual patchify + linear proj to conv"""541 out_dict = {}542 for k, v in state_dict.items():543 if 'patch_embed.proj.weight' in k:544 v = v.reshape((v.shape[0], 3, patch_size, patch_size))545 out_dict[k] = v546 547 return out_dict548 549 550class pvt_v2_b0(PyramidVisionTransformerImpr):551 def __init__(self, **kwargs):552 super(pvt_v2_b0, self).__init__(553 patch_size=4, embed_dims=[32, 64, 160, 256], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4],554 qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1],555 drop_rate=0.0, drop_path_rate=0.1)556 557 558 559class pvt_v2_b1(PyramidVisionTransformerImpr):560 def __init__(self, **kwargs):561 super(pvt_v2_b1, self).__init__(562 patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4],563 qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1],564 drop_rate=0.0, drop_path_rate=0.1)565 566class pvt_v2_b2(PyramidVisionTransformerImpr):567 def __init__(self, in_channels=3, **kwargs):568 super(pvt_v2_b2, self).__init__(569 patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4],570 qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1],571 drop_rate=0.0, drop_path_rate=0.1, in_channels=in_channels)572 573class pvt_v2_b3(PyramidVisionTransformerImpr):574 def __init__(self, **kwargs):575 super(pvt_v2_b3, self).__init__(576 patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4],577 qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 18, 3], sr_ratios=[8, 4, 2, 1],578 drop_rate=0.0, drop_path_rate=0.1)579 580class pvt_v2_b4(PyramidVisionTransformerImpr):581 def __init__(self, **kwargs):582 super(pvt_v2_b4, self).__init__(583 patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4],584 qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 8, 27, 3], sr_ratios=[8, 4, 2, 1],585 drop_rate=0.0, drop_path_rate=0.1)586 587 588class pvt_v2_b5(PyramidVisionTransformerImpr):589 def __init__(self, **kwargs):590 super(pvt_v2_b5, self).__init__(591 patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[4, 4, 4, 4],592 qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 6, 40, 3], sr_ratios=[8, 4, 2, 1],593 drop_rate=0.0, drop_path_rate=0.1)594 595 596 597### models/backbones/swin_v1.py598 599# --------------------------------------------------------600# Swin Transformer601# Copyright (c) 2021 Microsoft602# Licensed under The MIT License [see LICENSE for details]603# Written by Ze Liu, Yutong Lin, Yixuan Wei604# --------------------------------------------------------605 606import torch607import torch.nn as nn608import torch.nn.functional as F609import torch.utils.checkpoint as checkpoint610import numpy as np611from timm.layers import DropPath, to_2tuple, trunc_normal_612 613# from config import Config614 615 616# config = Config()617 618 619class Mlp(nn.Module):620 """ Multilayer perceptron."""621 622 def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):623 super().__init__()624 out_features = out_features or in_features625 hidden_features = hidden_features or in_features626 self.fc1 = nn.Linear(in_features, hidden_features)627 self.act = act_layer()628 self.fc2 = nn.Linear(hidden_features, out_features)629 self.drop = nn.Dropout(drop)630 631 def forward(self, x):632 x = self.fc1(x)633 x = self.act(x)634 x = self.drop(x)635 x = self.fc2(x)636 x = self.drop(x)637 return x638 639 640def window_partition(x, window_size):641 """642 Args:643 x: (B, H, W, C)644 window_size (int): window size645 646 Returns:647 windows: (num_windows*B, window_size, window_size, C)648 """649 B, H, W, C = x.shape650 x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)651 windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)652 return windows653 654 655def window_reverse(windows, window_size, H, W):656 """657 Args:658 windows: (num_windows*B, window_size, window_size, C)659 window_size (int): Window size660 H (int): Height of image661 W (int): Width of image662 663 Returns:664 x: (B, H, W, C)665 """666 B = int(windows.shape[0] / (H * W / window_size / window_size))667 x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)668 x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)669 return x670 671 672class WindowAttention(nn.Module):673 """ Window based multi-head self attention (W-MSA) module with relative position bias.674 It supports both of shifted and non-shifted window.675 676 Args:677 dim (int): Number of input channels.678 window_size (tuple[int]): The height and width of the window.679 num_heads (int): Number of attention heads.680 qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True681 qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set682 attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0683 proj_drop (float, optional): Dropout ratio of output. Default: 0.0684 """685 686 def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):687 688 super().__init__()689 self.dim = dim690 self.window_size = window_size # Wh, Ww691 self.num_heads = num_heads692 head_dim = dim // num_heads693 self.scale = qk_scale or head_dim ** -0.5694 695 # define a parameter table of relative position bias696 self.relative_position_bias_table = nn.Parameter(697 torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 2*Wh-1 * 2*Ww-1, nH698 699 # get pair-wise relative position index for each token inside the window700 coords_h = torch.arange(self.window_size[0])701 coords_w = torch.arange(self.window_size[1])702 coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing='ij')) # 2, Wh, Ww703 coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww704 relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww705 relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2706 relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0707 relative_coords[:, :, 1] += self.window_size[1] - 1708 relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1709 relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww710 self.register_buffer("relative_position_index", relative_position_index)711 712 self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)713 self.attn_drop_prob = attn_drop714 self.attn_drop = nn.Dropout(attn_drop)715 self.proj = nn.Linear(dim, dim)716 self.proj_drop = nn.Dropout(proj_drop)717 718 trunc_normal_(self.relative_position_bias_table, std=.02)719 self.softmax = nn.Softmax(dim=-1)720 721 def forward(self, x, mask=None):722 """ Forward function.723 724 Args:725 x: input features with shape of (num_windows*B, N, C)726 mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None727 """728 B_, N, C = x.shape729 qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)730 q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple)731 732 q = q * self.scale733 734 if config.SDPA_enabled:735 x = torch.nn.functional.scaled_dot_product_attention(736 q, k, v,737 attn_mask=None, dropout_p=self.attn_drop_prob, is_causal=False738 ).transpose(1, 2).reshape(B_, N, C)739 else:740 attn = (q @ k.transpose(-2, -1))741 742 relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(743 self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1744 ) # Wh*Ww, Wh*Ww, nH745 relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww746 attn = attn + relative_position_bias.unsqueeze(0)747 748 if mask is not None:749 nW = mask.shape[0]750 attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)751 attn = attn.view(-1, self.num_heads, N, N)752 attn = self.softmax(attn)753 else:754 attn = self.softmax(attn)755 756 attn = self.attn_drop(attn)757 758 x = (attn @ v).transpose(1, 2).reshape(B_, N, C)759 x = self.proj(x)760 x = self.proj_drop(x)761 return x762 763 764class SwinTransformerBlock(nn.Module):765 """ Swin Transformer Block.766 767 Args:768 dim (int): Number of input channels.769 num_heads (int): Number of attention heads.770 window_size (int): Window size.771 shift_size (int): Shift size for SW-MSA.772 mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.773 qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True774 qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.775 drop (float, optional): Dropout rate. Default: 0.0776 attn_drop (float, optional): Attention dropout rate. Default: 0.0777 drop_path (float, optional): Stochastic depth rate. Default: 0.0778 act_layer (nn.Module, optional): Activation layer. Default: nn.GELU779 norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm780 """781 782 def __init__(self, dim, num_heads, window_size=7, shift_size=0,783 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,784 act_layer=nn.GELU, norm_layer=nn.LayerNorm):785 super().__init__()786 self.dim = dim787 self.num_heads = num_heads788 self.window_size = window_size789 self.shift_size = shift_size790 self.mlp_ratio = mlp_ratio791 assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"792 793 self.norm1 = norm_layer(dim)794 self.attn = WindowAttention(795 dim, window_size=to_2tuple(self.window_size), num_heads=num_heads,796 qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)797 798 self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()799 self.norm2 = norm_layer(dim)800 mlp_hidden_dim = int(dim * mlp_ratio)801 self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)802 803 self.H = None804 self.W = None805 806 def forward(self, x, mask_matrix):807 """ Forward function.808 809 Args:810 x: Input feature, tensor size (B, H*W, C).811 H, W: Spatial resolution of the input feature.812 mask_matrix: Attention mask for cyclic shift.813 """814 B, L, C = x.shape815 H, W = self.H, self.W816 assert L == H * W, "input feature has wrong size"817 818 shortcut = x819 x = self.norm1(x)820 x = x.view(B, H, W, C)821 822 # pad feature maps to multiples of window size823 pad_l = pad_t = 0824 pad_r = (self.window_size - W % self.window_size) % self.window_size825 pad_b = (self.window_size - H % self.window_size) % self.window_size826 x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b))827 _, Hp, Wp, _ = x.shape828 829 # cyclic shift830 if self.shift_size > 0:831 shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))832 attn_mask = mask_matrix833 else:834 shifted_x = x835 attn_mask = None836 837 # partition windows838 x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C839 x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C840 841 # W-MSA/SW-MSA842 attn_windows = self.attn(x_windows, mask=attn_mask) # nW*B, window_size*window_size, C843 844 # merge windows845 attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)846 shifted_x = window_reverse(attn_windows, self.window_size, Hp, Wp) # B H' W' C847 848 # reverse cyclic shift849 if self.shift_size > 0:850 x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))851 else:852 x = shifted_x853 854 if pad_r > 0 or pad_b > 0:855 x = x[:, :H, :W, :].contiguous()856 857 x = x.view(B, H * W, C)858 859 # FFN860 x = shortcut + self.drop_path(x)861 x = x + self.drop_path(self.mlp(self.norm2(x)))862 863 return x864 865 866class PatchMerging(nn.Module):867 """ Patch Merging Layer868 869 Args:870 dim (int): Number of input channels.871 norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm872 """873 def __init__(self, dim, norm_layer=nn.LayerNorm):874 super().__init__()875 self.dim = dim876 self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)877 self.norm = norm_layer(4 * dim)878 879 def forward(self, x, H, W):880 """ Forward function.881 882 Args:883 x: Input feature, tensor size (B, H*W, C).884 H, W: Spatial resolution of the input feature.885 """886 B, L, C = x.shape887 assert L == H * W, "input feature has wrong size"888 889 x = x.view(B, H, W, C)890 891 # padding892 pad_input = (H % 2 == 1) or (W % 2 == 1)893 if pad_input:894 x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2))895 896 x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C897 x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C898 x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C899 x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C900 x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C901 x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C902 903 x = self.norm(x)904 x = self.reduction(x)905 906 return x907 908 909class BasicLayer(nn.Module):910 """ A basic Swin Transformer layer for one stage.911 912 Args:913 dim (int): Number of feature channels914 depth (int): Depths of this stage.915 num_heads (int): Number of attention head.916 window_size (int): Local window size. Default: 7.917 mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.918 qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True919 qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.920 drop (float, optional): Dropout rate. Default: 0.0921 attn_drop (float, optional): Attention dropout rate. Default: 0.0922 drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0923 norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm924 downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None925 use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.926 """927 928 def __init__(self,929 dim,930 depth,931 num_heads,932 window_size=7,933 mlp_ratio=4.,934 qkv_bias=True,935 qk_scale=None,936 drop=0.,937 attn_drop=0.,938 drop_path=0.,939 norm_layer=nn.LayerNorm,940 downsample=None,941 use_checkpoint=False):942 super().__init__()943 self.window_size = window_size944 self.shift_size = window_size // 2945 self.depth = depth946 self.use_checkpoint = use_checkpoint947 948 # build blocks949 self.blocks = nn.ModuleList([950 SwinTransformerBlock(951 dim=dim,952 num_heads=num_heads,953 window_size=window_size,954 shift_size=0 if (i % 2 == 0) else window_size // 2,955 mlp_ratio=mlp_ratio,956 qkv_bias=qkv_bias,957 qk_scale=qk_scale,958 drop=drop,959 attn_drop=attn_drop,960 drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,961 norm_layer=norm_layer)962 for i in range(depth)])963 964 # patch merging layer965 if downsample is not None:966 self.downsample = downsample(dim=dim, norm_layer=norm_layer)967 else:968 self.downsample = None969 970 def forward(self, x, H, W):971 """ Forward function.972 973 Args:974 x: Input feature, tensor size (B, H*W, C).975 H, W: Spatial resolution of the input feature.976 """977 978 # calculate attention mask for SW-MSA979 # Turn int to torch.tensor for the compatiability with torch.compile in PyTorch 2.5.980 Hp = torch.ceil(torch.tensor(H) / self.window_size).to(torch.int64) * self.window_size981 Wp = torch.ceil(torch.tensor(W) / self.window_size).to(torch.int64) * self.window_size982 img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device) # 1 Hp Wp 1983 h_slices = (slice(0, -self.window_size),984 slice(-self.window_size, -self.shift_size),985 slice(-self.shift_size, None))986 w_slices = (slice(0, -self.window_size),987 slice(-self.window_size, -self.shift_size),988 slice(-self.shift_size, None))989 cnt = 0990 for h in h_slices:991 for w in w_slices:992 img_mask[:, h, w, :] = cnt993 cnt += 1994 995 mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1996 mask_windows = mask_windows.view(-1, self.window_size * self.window_size)997 attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)998 attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)).to(x.dtype)999 1000 for blk in self.blocks:1001 blk.H, blk.W = H, W1002 if self.use_checkpoint:1003 x = checkpoint.checkpoint(blk, x, attn_mask)1004 else:1005 x = blk(x, attn_mask)1006 if self.downsample is not None:1007 x_down = self.downsample(x, H, W)1008 Wh, Ww = (H + 1) // 2, (W + 1) // 21009 return x, H, W, x_down, Wh, Ww1010 else:1011 return x, H, W, x, H, W1012 1013 1014class PatchEmbed(nn.Module):1015 """ Image to Patch Embedding1016 1017 Args:1018 patch_size (int): Patch token size. Default: 4.1019 in_channels (int): Number of input image channels. Default: 3.1020 embed_dim (int): Number of linear projection output channels. Default: 96.1021 norm_layer (nn.Module, optional): Normalization layer. Default: None1022 """1023 1024 def __init__(self, patch_size=4, in_channels=3, embed_dim=96, norm_layer=None):1025 super().__init__()1026 patch_size = to_2tuple(patch_size)1027 self.patch_size = patch_size1028 1029 self.in_channels = in_channels1030 self.embed_dim = embed_dim1031 1032 self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size)1033 if norm_layer is not None:1034 self.norm = norm_layer(embed_dim)1035 else:1036 self.norm = None1037 1038 def forward(self, x):1039 """Forward function."""1040 # padding1041 _, _, H, W = x.size()1042 if W % self.patch_size[1] != 0:1043 x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1]))1044 if H % self.patch_size[0] != 0:1045 x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0]))1046 1047 x = self.proj(x) # B C Wh Ww1048 if self.norm is not None:1049 Wh, Ww = x.size(2), x.size(3)1050 x = x.flatten(2).transpose(1, 2)1051 x = self.norm(x)1052 x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww)1053 1054 return x1055 1056 1057class SwinTransformer(nn.Module):1058 """ Swin Transformer backbone.1059 A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows` -1060 https://arxiv.org/pdf/2103.140301061 1062 Args:1063 pretrain_img_size (int): Input image size for training the pretrained model,1064 used in absolute postion embedding. Default 224.1065 patch_size (int | tuple(int)): Patch size. Default: 4.1066 in_channels (int): Number of input image channels. Default: 3.1067 embed_dim (int): Number of linear projection output channels. Default: 96.1068 depths (tuple[int]): Depths of each Swin Transformer stage.1069 num_heads (tuple[int]): Number of attention head of each stage.1070 window_size (int): Window size. Default: 7.1071 mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.1072 qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True1073 qk_scale (float): Override default qk scale of head_dim ** -0.5 if set.1074 drop_rate (float): Dropout rate.1075 attn_drop_rate (float): Attention dropout rate. Default: 0.1076 drop_path_rate (float): Stochastic depth rate. Default: 0.2.1077 norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.1078 ape (bool): If True, add absolute position embedding to the patch embedding. Default: False.1079 patch_norm (bool): If True, add normalization after patch embedding. Default: True.1080 out_indices (Sequence[int]): Output from which stages.1081 frozen_stages (int): Stages to be frozen (stop grad and set eval mode).1082 -1 means not freezing any parameters.1083 use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.1084 """1085 1086 def __init__(self,1087 pretrain_img_size=224,1088 patch_size=4,1089 in_channels=3,1090 embed_dim=96,1091 depths=[2, 2, 6, 2],1092 num_heads=[3, 6, 12, 24],1093 window_size=7,1094 mlp_ratio=4.,1095 qkv_bias=True,1096 qk_scale=None,1097 drop_rate=0.,1098 attn_drop_rate=0.,1099 drop_path_rate=0.2,1100 norm_layer=nn.LayerNorm,1101 ape=False,1102 patch_norm=True,1103 out_indices=(0, 1, 2, 3),1104 frozen_stages=-1,1105 use_checkpoint=False):1106 super().__init__()1107 1108 self.pretrain_img_size = pretrain_img_size1109 self.num_layers = len(depths)1110 self.embed_dim = embed_dim1111 self.ape = ape1112 self.patch_norm = patch_norm1113 self.out_indices = out_indices1114 self.frozen_stages = frozen_stages1115 1116 # split image into non-overlapping patches1117 self.patch_embed = PatchEmbed(1118 patch_size=patch_size, in_channels=in_channels, embed_dim=embed_dim,1119 norm_layer=norm_layer if self.patch_norm else None)1120 1121 # absolute position embedding1122 if self.ape:1123 pretrain_img_size = to_2tuple(pretrain_img_size)1124 patch_size = to_2tuple(patch_size)1125 patches_resolution = [pretrain_img_size[0] // patch_size[0], pretrain_img_size[1] // patch_size[1]]1126 1127 self.absolute_pos_embed = nn.Parameter(torch.zeros(1, embed_dim, patches_resolution[0], patches_resolution[1]))1128 trunc_normal_(self.absolute_pos_embed, std=.02)1129 1130 self.pos_drop = nn.Dropout(p=drop_rate)1131 1132 # stochastic depth1133 dpr = np.linspace(0, drop_path_rate, sum(depths)).tolist() # stochastic depth decay rule1134 1135 # build layers1136 self.layers = nn.ModuleList()1137 for i_layer in range(self.num_layers):1138 layer = BasicLayer(1139 dim=int(embed_dim * 2 ** i_layer),1140 depth=depths[i_layer],1141 num_heads=num_heads[i_layer],1142 window_size=window_size,1143 mlp_ratio=mlp_ratio,1144 qkv_bias=qkv_bias,1145 qk_scale=qk_scale,1146 drop=drop_rate,1147 attn_drop=attn_drop_rate,1148 drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],1149 norm_layer=norm_layer,1150 downsample=PatchMerging if (i_layer < self.num_layers - 1) else None,1151 use_checkpoint=use_checkpoint)1152 self.layers.append(layer)1153 1154 num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)]1155 self.num_features = num_features1156 1157 # add a norm layer for each output1158 for i_layer in out_indices:1159 layer = norm_layer(num_features[i_layer])1160 layer_name = f'norm{i_layer}'1161 self.add_module(layer_name, layer)1162 1163 self._freeze_stages()1164 1165 def _freeze_stages(self):1166 if self.frozen_stages >= 0:1167 self.patch_embed.eval()1168 for param in self.patch_embed.parameters():1169 param.requires_grad = False1170 1171 if self.frozen_stages >= 1 and self.ape:1172 self.absolute_pos_embed.requires_grad = False1173 1174 if self.frozen_stages >= 2:1175 self.pos_drop.eval()1176 for i in range(0, self.frozen_stages - 1):1177 m = self.layers[i]1178 m.eval()1179 for param in m.parameters():1180 param.requires_grad = False1181 1182 1183 def forward(self, x):1184 """Forward function."""1185 x = self.patch_embed(x)1186 1187 Wh, Ww = x.size(2), x.size(3)1188 if self.ape:1189 # interpolate the position embedding to the corresponding size1190 absolute_pos_embed = F.interpolate(self.absolute_pos_embed, size=(Wh, Ww), mode='bicubic')1191 x = (x + absolute_pos_embed) # B Wh*Ww C1192 1193 outs = []#x.contiguous()]1194 x = x.flatten(2).transpose(1, 2)1195 x = self.pos_drop(x)1196 for i in range(self.num_layers):1197 layer = self.layers[i]1198 x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww)1199 1200 if i in self.out_indices: