update

嘉渊

2023-04-25 2a3336d9a8a9d74662be805510abdd0a29414ae5

 funasr/bin/train.py

@@ -59,16 +59,19 @@
    )
    parser.add_argument(
        "--dist_world_size",
        type=int,
        default=1,
        help="number of nodes for distributed training",
    )
    parser.add_argument(
        "--dist_rank",
        type=int,
        default=None,
        help="node rank for distributed training",
    )
    parser.add_argument(
        "--local_rank",
        type=int,
        default=None,
        help="local rank for distributed training",
    )
@@ -238,6 +241,12 @@
        help="Enable resuming if checkpoint is existing",
    )
    parser.add_argument(
        "--train_dtype",
        default="float32",
        choices=["float16", "float32", "float64"],
        help="Data type for training.",
    )
    parser.add_argument(
        "--use_amp",
        type=str2bool,
        default=False,
@@ -249,6 +258,12 @@
        help="Show the logs every the number iterations in each epochs at the "
             "training phase. If None is given, it is decided according the number "
             "of training samples automatically .",
    )
    parser.add_argument(
        "--use_tensorboard",
        type=str2bool,
        default=True,
        help="Enable tensorboard logging",
    )

    # pretrained model related
@@ -294,39 +309,30 @@
        help="whether to use dataloader for large dataset",
    )
    parser.add_argument(
        "--train_data_file",
        "--dataset_conf",
        action=NestedDictAction,
        default=dict(),
        help=f"The keyword arguments for dataset",
    )
    parser.add_argument(
        "--data_dir",
        type=str,
        default=None,
        help="train_list for large dataset",
        help="root path of data",
    )
    parser.add_argument(
        "--valid_data_file",
        "--train_set",
        type=str,
        default=None,
        help="valid_list for large dataset",
        default="train",
        help="train dataset",
    )
    parser.add_argument(
        "--train_data_path_and_name_and_type",
        action="append",
        default=[],
        help="e.g. '--train_data_path_and_name_and_type some/path/a.scp,foo,sound'. ",
    )
    parser.add_argument(
        "--valid_data_path_and_name_and_type",
        action="append",
        default=[],
    )
    parser.add_argument(
        "--train_shape_file",
        type=str, action="append",
        default=[],
    )
    parser.add_argument(
        "--valid_shape_file",
        "--valid_set",
        type=str,
        action="append",
        default=[],
        default="validation",
        help="dev dataset",
    )

    parser.add_argument(
        "--use_preprocessor",
        type=str2bool,
@@ -462,9 +468,9 @@

if __name__ == '__main__':
    parser = get_parser()
    common_args, extra_task_params = parser.parse_known_args()
    args, extra_task_params = parser.parse_known_args()
    if extra_task_params:
        args = build_args(common_args, parser, extra_task_params)
        args = build_args(args, parser, extra_task_params)

    # set random seed
    set_all_random_seed(args.seed)
@@ -495,6 +501,10 @@
    prepare_data(args, distributed_option)

    model = build_model(args)
    model = model.to(
        dtype=getattr(torch, args.train_dtype),
        device="cuda" if args.ngpu > 0 else "cpu",
    )
    optimizers = build_optimizer(args, model=model)
    schedulers = build_scheduler(args, optimizers)

@@ -502,6 +512,7 @@
                                                                   distributed_option.dist_rank,
                                                                   distributed_option.local_rank))
    logging.info(pytorch_cudnn_version())
    logging.info("Args: {}".format(args))
    logging.info(model_summary(model))
    logging.info("Optimizer: {}".format(optimizers))
    logging.info("Scheduler: {}".format(schedulers))