python/FunASR-XL.git

			@@ -43,12 +43,15 @@
			from funasr.iterators.chunk_iter_factory import ChunkIterFactory
			from funasr.iterators.multiple_iter_factory import MultipleIterFactory
			from funasr.iterators.sequence_iter_factory import SequenceIterFactory
			from funasr.main_funcs.collect_stats import collect_stats
			from funasr.optimizers.sgd import SGD
			from funasr.optimizers.fairseq_adam import FairseqAdam
			from funasr.samplers.build_batch_sampler import BATCH_TYPES
			from funasr.samplers.build_batch_sampler import build_batch_sampler
			from funasr.samplers.unsorted_batch_sampler import UnsortedBatchSampler
			from funasr.schedulers.noam_lr import NoamLR
			from funasr.schedulers.warmup_lr import WarmupLR
			from funasr.schedulers.tri_stage_scheduler import TriStageLR
			from funasr.torch_utils.load_pretrained_model import load_pretrained_model
			from funasr.torch_utils.model_summary import model_summary
			from funasr.torch_utils.pytorch_version import pytorch_cudnn_version
			@@ -83,6 +86,7 @@

			optim_classes = dict(
			adam=torch.optim.Adam,
			fairseq_adam=FairseqAdam,
			adamw=torch.optim.AdamW,
			sgd=SGD,
			adadelta=torch.optim.Adadelta,
			@@ -149,6 +153,7 @@
			CosineAnnealingLR=torch.optim.lr_scheduler.CosineAnnealingLR,
			noamlr=NoamLR,
			warmuplr=WarmupLR,
			tri_stage=TriStageLR,
			cycliclr=torch.optim.lr_scheduler.CyclicLR,
			onecyclelr=torch.optim.lr_scheduler.OneCycleLR,
			CosineAnnealingWarmRestarts=torch.optim.lr_scheduler.CosineAnnealingWarmRestarts,
			@@ -459,6 +464,12 @@
			default=sys.maxsize,
			help="The maximum number update step to train",
			)
			parser.add_argument(
			"--batch_interval",
			type=int,
			default=10000,
			help="The batch interval for saving model.",
			)
			group.add_argument(
			"--patience",
			type=int_or_none,
			@@ -634,12 +645,12 @@
			"and exclude_keys excludes keys of model states for the initialization."
			"e.g.\n"
			" # Load all parameters"
			" --init_param some/where/model.pth\n"
			" --init_param some/where/model.pb\n"
			" # Load only decoder parameters"
			" --init_param some/where/model.pth:decoder:decoder\n"
			" --init_param some/where/model.pb:decoder:decoder\n"
			" # Load only decoder parameters excluding decoder.embed"
			" --init_param some/where/model.pth:decoder:decoder:decoder.embed\n"
			" --init_param some/where/model.pth:decoder:decoder:decoder.embed\n",
			" --init_param some/where/model.pb:decoder:decoder:decoder.embed\n"
			" --init_param some/where/model.pb:decoder:decoder:decoder.embed\n",
			)
			group.add_argument(
			"--ignore_init_mismatch",
			@@ -1188,12 +1199,18 @@
			# logging.basicConfig() is invoked in main_worker() instead of main()
			# because it can be invoked only once in a process.
			# FIXME(kamo): Should we use logging.getLogger()?
			# BUGFIX: Remove previous handlers and reset log level
			for handler in logging.root.handlers[:]:
			logging.root.removeHandler(handler)
			logging.basicConfig(
			level=args.log_level,
			format=f"[{os.uname()[1].split('.')[0]}]"
			f" %(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
			)
			else:
			# BUGFIX: Remove previous handlers and reset log level
			for handler in logging.root.handlers[:]:
			logging.root.removeHandler(handler)
			# Suppress logging if RANK != 0
			logging.basicConfig(
			level="ERROR",
			@@ -1276,6 +1293,52 @@

			if args.dry_run:
			pass
			elif args.collect_stats:
			# Perform on collect_stats mode. This mode has two roles
			# - Derive the length and dimension of all input data
			# - Accumulate feats, square values, and the length for whitening

			if args.valid_batch_size is None:
			args.valid_batch_size = args.batch_size

			if len(args.train_shape_file) != 0:
			train_key_file = args.train_shape_file[0]
			else:
			train_key_file = None
			if len(args.valid_shape_file) != 0:
			valid_key_file = args.valid_shape_file[0]
			else:
			valid_key_file = None

			collect_stats(
			model=model,
			train_iter=cls.build_streaming_iterator(
			data_path_and_name_and_type=args.train_data_path_and_name_and_type,
			key_file=train_key_file,
			batch_size=args.batch_size,
			dtype=args.train_dtype,
			num_workers=args.num_workers,
			allow_variable_data_keys=args.allow_variable_data_keys,
			ngpu=args.ngpu,
			preprocess_fn=cls.build_preprocess_fn(args, train=False),
			collate_fn=cls.build_collate_fn(args, train=False),
			),
			valid_iter=cls.build_streaming_iterator(
			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,
			dtype=args.train_dtype,
			num_workers=args.num_workers,
			allow_variable_data_keys=args.allow_variable_data_keys,
			ngpu=args.ngpu,
			preprocess_fn=cls.build_preprocess_fn(args, train=False),
			collate_fn=cls.build_collate_fn(args, train=False),
			),
			output_dir=output_dir,
			ngpu=args.ngpu,
			log_interval=args.log_interval,
			write_collected_feats=args.write_collected_feats,
			)
			else:
			logging.info("Training args: {}".format(args))
			# 6. Loads pre-trained model
			@@ -1297,12 +1360,16 @@
			if args.dataset_type == "large":
			from funasr.datasets.large_datasets.build_dataloader import ArkDataLoader
			train_iter_factory = ArkDataLoader(args.train_data_file, args.token_list, args.dataset_conf,
			seg_dict_file=args.seg_dict_file if hasattr(args,
			"seg_dict_file") else None,
			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,
			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,
			seg_dict_file=args.seg_dict_file if hasattr(args,
			"seg_dict_file") else None,
			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,
			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(
			@@ -1515,12 +1582,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=dest_sample_rate,
			)
			cls.check_task_requirements(
			dataset, args.allow_variable_data_keys, train=iter_options.train
			@@ -1792,6 +1868,7 @@
			key_file: str = None,
			batch_size: int = 1,
			fs: dict = None,
			mc: bool = False,
			dtype: str = np.float32,
			num_workers: int = 1,
			allow_variable_data_keys: bool = False,
			@@ -1810,6 +1887,7 @@
			data_path_and_name_and_type,
			float_dtype=dtype,
			fs=fs,
			mc=mc,
			preprocess=preprocess_fn,
			key_file=key_file,
			)