Update modeling_maelm.py

modeling_maelm.py

@@ -192,9 +192,9 @@ class MAEForCausalLM(PreTrainedModel):
         if bk_name == 'MAEViT':
             ckpt_path = backbone.pop('ckpt') if 'ckpt' in backbone else None
             self.backbone = MAEViT(**backbone)
-            if ckpt_path is not None:
-                ckpt = torch.load( ckpt_path,'cpu')
-                self.backbone.load_state_dict(ckpt['state_dict'])
+            #if ckpt_path is not None:
+            #    ckpt = torch.load( ckpt_path,'cpu')
+            #    self.backbone.load_state_dict(ckpt['state_dict'])
                 
         elif bk_name == 'HTSAT':
             ckpt_path = backbone.pop('ckpt') if 'ckpt' in backbone else None
@@ -239,16 +239,16 @@ class MAEForCausalLM(PreTrainedModel):
         # float32 --> bfloat16
         for p in self.parameters():
             p.data = p.data.to(torch.bfloat16)
-        if config.resume_from_checkpoint is not None:  
-            drain_loader = True
-            accelerator.load_state(config.resume_from_checkpoint, load_module_strict=False)
-            # start_epoch, start_step, all_step = [int(_.split('_')[1]) for _ in args.resume_from_checkpoint.split('/')[-2].split('-')]
-        elif config.resume_from_pth is not None:
-            print(f'###########loading##########{config.resume_from_pth}###########loading##########')
-            ckpt = torch.load(config.resume_from_pth, map_location='cpu')
-            ckpt_copy = {k[7:]: v for k, v in ckpt.items()}
-            self.load_state_dict(ckpt_copy, strict=False)
-            print(f'###########loaded##########{config.resume_from_pth}###########loaded##########')
+        #if config.resume_from_checkpoint is not None:  
+        #    drain_loader = True
+        #    accelerator.load_state(config.resume_from_checkpoint, load_module_strict=False)
+        #    # start_epoch, start_step, all_step = [int(_.split('_')[1]) for _ in args.resume_from_checkpoint.split('/')[-2].split('-')]
+        #elif config.resume_from_pth is not None:
+        #    print(f'###########loading##########{config.resume_from_pth}###########loading##########')
+        #    ckpt = torch.load(config.resume_from_pth, map_location='cpu')
+        #    ckpt_copy = {k[7:]: v for k, v in ckpt.items()}
+        #    self.load_state_dict(ckpt_copy, strict=False)
+        #    print(f'###########loaded##########{config.resume_from_pth}###########loaded##########')
 
         if False:
             self.patch_llm()