Merge pull request #80 from SegmentationBLWX/dev

merge from dev

Author: CharlesPikachu

ssseg/modules/models/backbones/mae.py

@@ -30,6 +30,9 @@ def initweights(self):
 class MAETransformerEncoderLayer(BEiTTransformerEncoderLayer):
     '''buildattn'''
     def buildattn(self, attn_cfg):
+        valid_keys = ['embed_dims', 'num_heads', 'window_size', 'bias', 'qk_scale', 'attn_drop_rate', 'proj_drop_rate']
+        for key in list(attn_cfg.keys()):
+            if key not in valid_keys: attn_cfg.pop(key)
         self.attn = MAEAttention(**attn_cfg)
 
 

tests/test_backbones/test_mae.py

@@ -0,0 +1,32 @@
+'''
+Function:
+    Implementation of Testing MAE
+Author:
+    Zhenchao Jin
+'''
+from ssseg.modules import BuildBackbone, loadpretrainedweights
+from ssseg.modules.models.backbones.mae import DEFAULT_MODEL_URLS
+
+
+'''MAEs'''
+cfgs = [{
+    'type': 'MAE', 'structure_type': 'mae_pretrain_vit_base', 'pretrained': True, 
+    'img_size': (512, 512), 'patch_size': 16, 'embed_dims': 768, 'num_layers': 12,
+    'num_heads': 12, 'mlp_ratio': 4, 'init_values': 1.0, 'drop_path_rate': 0.1,
+    'selected_indices': (0, 1, 2, 3), 'norm_cfg': {'type': 'LayerNorm', 'eps': 1e-6},
+}]
+for cfg in cfgs:
+    mae = BuildBackbone(cfg)
+    state_dict = loadpretrainedweights(
+        structure_type=cfg['structure_type'], pretrained_model_path='', default_model_urls=DEFAULT_MODEL_URLS
+    )
+    state_dict = mae.beitconvert(state_dict)
+    state_dict = mae.resizerelposembed(state_dict)
+    try:
+        mae.load_state_dict(state_dict, strict=False)
+    except Exception as err:
+        print(err)
+    try:
+        mae.load_state_dict(state_dict, strict=True)
+    except Exception as err:
+        print(err)

tests/test_backbones/test_resnest.py

@@ -0,0 +1,26 @@
+'''
+Function:
+    Implementation of Testing ResNeSt
+Author:
+    Zhenchao Jin
+'''
+from ssseg.modules import BuildBackbone, loadpretrainedweights
+from ssseg.modules.models.backbones.resnest import DEFAULT_MODEL_URLS
+
+
+'''resnests'''
+for depth in [50, 101, 200]:
+    resnest = BuildBackbone(backbone_cfg={
+        'type': 'ResNeSt', 'depth': depth, 'structure_type': f'resnest{depth}', 'pretrained': True, 'outstride': 8, 'selected_indices': (0, 1, 2, 3), 'stem_channels': {50: 64, 101: 128, 200: 128}[depth]
+    })
+    state_dict = loadpretrainedweights(
+        structure_type=f'resnest{depth}', pretrained_model_path='', default_model_urls=DEFAULT_MODEL_URLS
+    )
+    try:
+        resnest.load_state_dict(state_dict, strict=False)
+    except Exception as err:
+        print(err)
+    try:
+        resnest.load_state_dict(state_dict, strict=True)
+    except Exception as err:
+        print(err)