Update load_pretrained_model.py

zhifu gao

2024-01-17 e67ed1d45d5a9d7fb7bb22d15fd2bfef17e9076f

 funasr/bin/train.py

@@ -40,8 +40,7 @@


def main(**kwargs):
    # preprocess_config(kwargs)
    # import pdb; pdb.set_trace()
    
    # set random seed
    tables.print()
    set_all_random_seed(kwargs.get("seed", 0))
@@ -96,9 +95,11 @@
            logging.info(f"Loading pretrained params from {p}")
            load_pretrained_model(
                model=model,
                init_param=p,
                path=p,
                ignore_init_mismatch=kwargs.get("ignore_init_mismatch", True),
                oss_bucket=kwargs.get("oss_bucket", None),
                scope_map=kwargs.get("scope_map", None),
                excludes=kwargs.get("excludes", None),
            )
    else:
        initialize(model, kwargs.get("init", "kaiming_normal"))
@@ -167,6 +168,8 @@
        local_rank=local_rank,
        use_ddp=use_ddp,
        use_fsdp=use_fsdp,
        output_dir=kwargs.get("output_dir", "./exp"),
        resume=kwargs.get("resume", True),
        **kwargs.get("train_conf"),
    )
    trainer.run()