雾聪
2023-05-17 8706e767affc6bdc8cb7a67ca3a20a62779ff048
funasr/tasks/abs_task.py
@@ -445,6 +445,12 @@
            help='Perform on "collect stats" mode',
        )
        group.add_argument(
            "--mc",
            type=bool,
            default=False,
            help="MultiChannel input",
        )
        group.add_argument(
            "--write_collected_feats",
            type=str2bool,
            default=False,
@@ -463,6 +469,12 @@
            type=int,
            default=sys.maxsize,
            help="The maximum number update step to train",
        )
        parser.add_argument(
            "--batch_interval",
            type=int,
            default=-1,
            help="The batch interval for saving model.",
        )
        group.add_argument(
            "--patience",
@@ -541,6 +553,12 @@
            type=int,
            default=1,
            help="The number of gradient accumulation",
        )
        group.add_argument(
            "--bias_grad_times",
            type=float,
            default=1.0,
            help="To scale the gradient of contextual related params",
        )
        group.add_argument(
            "--no_forward_run",
@@ -629,8 +647,8 @@
        group.add_argument(
            "--init_param",
            type=str,
            action="append",
            default=[],
            nargs="*",
            help="Specify the file path used for initialization of parameters. "
                 "The format is '<file_path>:<src_key>:<dst_key>:<exclude_keys>', "
                 "where file_path is the model file path, "
@@ -656,7 +674,7 @@
            "--freeze_param",
            type=str,
            default=[],
            nargs="*",
            action="append",
            help="Freeze parameters",
        )
@@ -1147,10 +1165,10 @@
        elif args.distributed and args.simple_ddp:
            distributed_option.init_torch_distributed_pai(args)
            args.ngpu = dist.get_world_size()
            if args.dataset_type == "small":
            if args.dataset_type == "small" and args.ngpu > 0:
                if args.batch_size is not None:
                    args.batch_size = args.batch_size * args.ngpu
                if args.batch_bins is not None:
                if args.batch_bins is not None and args.ngpu > 0:
                    args.batch_bins = args.batch_bins * args.ngpu
        # filter samples if wav.scp and text are mismatch
@@ -1310,6 +1328,7 @@
                    data_path_and_name_and_type=args.train_data_path_and_name_and_type,
                    key_file=train_key_file,
                    batch_size=args.batch_size,
                    mc=args.mc,
                    dtype=args.train_dtype,
                    num_workers=args.num_workers,
                    allow_variable_data_keys=args.allow_variable_data_keys,
@@ -1321,6 +1340,7 @@
                    data_path_and_name_and_type=args.valid_data_path_and_name_and_type,
                    key_file=valid_key_file,
                    batch_size=args.valid_batch_size,
                    mc=args.mc,
                    dtype=args.train_dtype,
                    num_workers=args.num_workers,
                    allow_variable_data_keys=args.allow_variable_data_keys,
@@ -1355,15 +1375,15 @@
                from funasr.datasets.large_datasets.build_dataloader import ArkDataLoader
                train_iter_factory = ArkDataLoader(args.train_data_file, args.token_list, args.dataset_conf,
                                                   frontend_conf=args.frontend_conf if hasattr(args, "frontend_conf") else None,
                                                   seg_dict_file=args.seg_dict_file if hasattr(args,
                                                                                               "seg_dict_file") else None,
                                                   seg_dict_file=args.seg_dict_file if hasattr(args, "seg_dict_file") else None,
                                                   punc_dict_file=args.punc_list if hasattr(args, "punc_list") else None,
                                                   bpemodel_file=args.bpemodel if hasattr(args, "bpemodel") else None,
                                                   mode="train")
                valid_iter_factory = ArkDataLoader(args.valid_data_file, args.token_list, args.dataset_conf, 
                                                   frontend_conf=args.frontend_conf if hasattr(args, "frontend_conf") else None,
                                                   seg_dict_file=args.seg_dict_file if hasattr(args,
                                                                                               "seg_dict_file") else None,
                                                   seg_dict_file=args.seg_dict_file if hasattr(args, "seg_dict_file") else None,
                                                   punc_dict_file=args.punc_list if hasattr(args, "punc_list") else None,
                                                   bpemodel_file=args.bpemodel if hasattr(args, "bpemodel") else None,
                                                   mode="eval")
            elif args.dataset_type == "small":
                train_iter_factory = cls.build_iter_factory(
@@ -1576,13 +1596,21 @@
    ) -> AbsIterFactory:
        assert check_argument_types()
        if hasattr(args, "frontend_conf"):
            if args.frontend_conf is not None and "fs" in args.frontend_conf:
                dest_sample_rate = args.frontend_conf["fs"]
            else:
                dest_sample_rate = 16000
        else:
            dest_sample_rate = 16000
        dataset = ESPnetDataset(
            iter_options.data_path_and_name_and_type,
            float_dtype=args.train_dtype,
            preprocess=iter_options.preprocess_fn,
            max_cache_size=iter_options.max_cache_size,
            max_cache_fd=iter_options.max_cache_fd,
            dest_sample_rate=args.frontend_conf["fs"],
            dest_sample_rate=dest_sample_rate,
        )
        cls.check_task_requirements(
            dataset, args.allow_variable_data_keys, train=iter_options.train