Merge branch 'master' of git.nju.edu.cn:wanggh/Swin-Transformer-Object-Detection

configs/swin/mask_rcnn_swin_tiny_patch4_window7_mstrain_480-800_adamw_1x_coco.py

@@ -68,7 +68,7 @@ optimizer = dict(_delete_=True, type='AdamW', lr=0.0001, betas=(0.9, 0.999), wei
 lr_config = dict(step=[8, 11])
 #runner = dict(type='EpochBasedRunnerAmp', max_epochs=12)
 runner = dict(type='EpochBasedRunner', max_epochs=12)
- 
+
 # do not use mmdet version fp16
 # fp16 = None
 # optimizer_config = dict(

tools/gen_checkpoint.py

@@ -29,7 +29,6 @@ def merge(backbone, head):
     tsd = target['state_dict']
     bsd = backbone['state_dict']
     hsd = head['state_dict']
-
     for key in bsd.keys():
         if 'backbone' in key:
             tsd[key] = bsd[key]
@@ -67,6 +66,23 @@ def gen_backbone(backbone, new_backbone):
         tsd[key] = nbsd[nk]
     return target
 
+def gen_imagenet_h(backbone, head):
+    target = dict()
+    target['state_dict'] = dict()
+    tsd = target['state_dict']
+    bsd = backbone['model']
+    hsd = head['state_dict']
+    for key in hsd.keys():
+        if 'backbone' not in key:
+            tsd[key] = hsd[key]
+        else:
+            bkey = key[9:]
+            if bkey not in bsd:
+                print("{} not load".format(key))
+                continue
+            tsd[key] = bsd[bkey]
+    return target
+
 def main():
     args = parse_args()
     print("generate checkpoint")
@@ -80,6 +96,10 @@ def main():
         head = get_sd(args.head, return_sd=False)
         print("backbone+neck:{} head:{}".format(args.backbone_neck, args.head))
         target = merge_bn_h(backbone, head)
+    elif args.new_backbone and args.head:
+        backbone = get_sd(args.new_backbone, return_sd=False)
+        head = get_sd(args.head, return_sd=False)
+        target = gen_imagenet_h(backbone, head)
     elif args.new_backbone:
         backbone = get_sd(args.backbone, return_sd=False)
         nb = get_sd(args.new_backbone, return_sd=False)