bugfix

游雁

2024-02-21 497edf4c9d6c1565a4bcf1a3edfcd47ffec8c10d

bugfix

 examples/aishell/conformer/infer.sh

New file
@@ -0,0 +1,12 @@


python funasr/bin/inference.py \
--config-path="/mnt/workspace/FunASR/examples/aishell/paraformer/exp/baseline_paraformer_conformer_12e_6d_2048_256_zh_char_exp3" \
--config-name="config.yaml" \
++init_param="/mnt/workspace/FunASR/examples/aishell/paraformer/exp/baseline_paraformer_conformer_12e_6d_2048_256_zh_char_exp3/model.pt.ep38" \
++tokenizer_conf.token_list="/mnt/nfs/zhifu.gzf/data/AISHELL-1-feats/DATA/data/zh_token_list/char/tokens.txt" \
++frontend_conf.cmvn_file="/mnt/nfs/zhifu.gzf/data/AISHELL-1-feats/DATA/data/train/am.mvn" \
++input="/mnt/nfs/zhifu.gzf/data/AISHELL-1/data_aishell/wav/train/S0002/BAC009S0002W0122.wav" \
++output_dir="./outputs/debug" \
++device="cpu" \


 funasr/train_utils/load_pretrained_model.py

@@ -75,6 +75,7 @@
    
    return assignment_map


def load_pretrained_model(
    path: str,
    model: torch.nn.Module,
@@ -94,25 +95,69 @@
    """
    
    obj = model
	
    dst_state = obj.state_dict()
    # import pdb;
    # pdb.set_trace()
    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)
    src_state = src_state["model"] if "model" in src_state else src_state
    if "state_dict" in src_state:
        src_state = src_state["state_dict"]
    
    if excludes is not None:
        for e in excludes.split(","):
            src_state = {k: v for k, v in src_state.items() if not k.startswith(e)}
    for k in dst_state.keys():
        if not k.startswith("module.") and "module." + k in src_state.keys():
            k_ddp = "module." + k
        else:
            k_ddp = k
        if k_ddp in src_state:
            dst_state[k] = src_state[k_ddp]
        else:
            print(f"Miss key in ckpt: model: {k}, ckpt: {k_ddp}")
    
    dst_state = obj.state_dict()
    src_state = assigment_scope_map(dst_state, src_state, scope_map)
    flag = obj.load_state_dict(dst_state, strict=True)
    print(flag)
    
    if ignore_init_mismatch:
        src_state = filter_state_dict(dst_state, src_state)
	
    logging.debug("Loaded src_state keys: {}".format(src_state.keys()))
    logging.debug("Loaded dst_state keys: {}".format(dst_state.keys()))
    dst_state.update(src_state)
    obj.load_state_dict(dst_state, strict=True)
# def load_pretrained_model(
#     path: str,
#     model: torch.nn.Module,
#     ignore_init_mismatch: bool,
#     map_location: str = "cpu",
#     oss_bucket=None,
#     scope_map=None,
#     excludes=None,
# ):
#     """Load a model state and set it to the model.
#
#     Args:
#         init_param: <file_path>:<src_key>:<dst_key>:<exclude_Keys>
#
#     Examples:
#
#     """
#
#     obj = model
#
#     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)
#     src_state = src_state["model"] if "model" in src_state else src_state
#
#     if excludes is not None:
#         for e in excludes.split(","):
#             src_state = {k: v for k, v in src_state.items() if not k.startswith(e)}
#
#     dst_state = obj.state_dict()
#     src_state = assigment_scope_map(dst_state, src_state, scope_map)
#
#     if ignore_init_mismatch:
#         src_state = filter_state_dict(dst_state, src_state)
#
#     logging.debug("Loaded src_state keys: {}".format(src_state.keys()))
#     logging.debug("Loaded dst_state keys: {}".format(dst_state.keys()))
#     dst_state.update(src_state)
#     obj.load_state_dict(dst_state, strict=True)

 funasr/train_utils/trainer.py

@@ -128,7 +128,20 @@
        if os.path.isfile(ckpt):
            checkpoint = torch.load(ckpt)
            self.start_epoch = checkpoint['epoch'] + 1
            self.model.load_state_dict(checkpoint['state_dict'])
            # self.model.load_state_dict(checkpoint['state_dict'])
            src_state = checkpoint['state_dict']
            dst_state = self.model.state_dict()
            for k in dst_state.keys():
                if not k.startswith("module.") and "module."+k in src_state.keys():
                    k_ddp = "module."+k
                else:
                    k_ddp = k
                if k_ddp in src_state.keys():
                    dst_state[k] = src_state[k_ddp]
                else:
                    print(f"Miss key in ckpt: model: {k}, ckpt: {k_ddp}")

            self.model.load_state_dict(dst_state)
            self.optim.load_state_dict(checkpoint['optimizer'])
            self.scheduler.load_state_dict(checkpoint['scheduler'])
            print(f"Checkpoint loaded successfully from '{ckpt}'")