Support specifying number of workers for on-the-fly feature extraction.

This commit is contained in:
Fangjun Kuang 2022-04-14 10:26:45 +08:00
parent bbf074a36b
commit 0cc13bc702

View File

@ -101,6 +101,13 @@ class AsrDataModule:
"collect the batches.",
)
group.add_argument(
"--on-the-fly-num-workers",
type=int,
default=0,
help="The number of workers for on-the-fly feature extraction",
)
group.add_argument(
"--enable-spec-aug",
type=str2bool,
@ -212,7 +219,10 @@ class AsrDataModule:
train = K2SpeechRecognitionDataset(
cut_transforms=transforms,
input_strategy=(
OnTheFlyFeatures(Fbank(FbankConfig(num_mel_bins=80)))
OnTheFlyFeatures(
extractor=Fbank(FbankConfig(num_mel_bins=80)),
num_workers=self.args.on_the_fly_num_workers,
)
if on_the_fly_feats
else PrecomputedFeatures()
),