zhifu gao
2024-03-04 d2c1204d91d7c98be7998e3966bd82e22750293b
funasr/train_utils/load_pretrained_model.py
@@ -68,9 +68,9 @@
   else:
      buffer = BytesIO(oss_bucket.get_object(path).read())
      src_state = torch.load(buffer, map_location=map_location)
   if "state_dict" in src_state:
      src_state = src_state["state_dict"]
   src_state = src_state["state_dict"] if "state_dict" in src_state else src_state
   src_state = src_state["model_state_dict"] if "model_state_dict" in src_state else src_state
   src_state = src_state["model"] if "model" in src_state else src_state
   
   if isinstance(scope_map, str):
@@ -90,10 +90,12 @@
            
            if dst_prefix == "" and (src_prefix + k) in src_state.keys():
               k_src = src_prefix + k
               print(f"init param, map: {k} from {k_src} in ckpt")
               if not k_src.startswith("module."):
                  print(f"init param, map: {k} from {k_src} in ckpt")
            elif k.startswith(dst_prefix) and k.replace(dst_prefix, src_prefix, 1) in src_state.keys():
               k_src = k.replace(dst_prefix, src_prefix, 1)
               print(f"init param, map: {k} from {k_src} in ckpt")
               if not k_src.startswith("module."):
                  print(f"init param, map: {k} from {k_src} in ckpt")
               
      if k_src in src_state.keys():
         if ignore_init_mismatch and dst_state[k].shape != src_state[k_src].shape: