update cmakelist

雾聪

2024-03-14 0cf5dfec2c8313fc2ed2aab8d10bf3dc4b9c283f

 funasr/train_utils/load_pretrained_model.py

@@ -7,7 +7,7 @@
import torch
import torch.nn
import torch.optim

import pdb

def filter_state_dict(
   dst_state: Dict[str, Union[float, torch.Tensor]],
@@ -63,14 +63,15 @@
   dst_state = obj.state_dict()
   
   print(f"ckpt: {path}")

   if oss_bucket is None:
      src_state = torch.load(path, map_location=map_location)
   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 +91,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: