def evaluate(
self,
test_set,
max_key=None,
min_key=None,
progressbar=None,
test_loader_kwargs={},
):
"""Iterate test_set and evaluate brain performance. By default, loads
the best-performing checkpoint (as recorded using the checkpointer).
Arguments
---------
test_set : Dataset, DataLoader
If a DataLoader is given, it is iterated directly. Otherwise passed
to self.make_dataloader()
max_key : str
Key to use for finding best checkpoint, passed to on_evaluate_start
min_key : str
Key to use for finding best checkpoint, passed to on_evaluate_start
progressbar : bool
Whether to display the progress in a progressbar.
test_loader_kwargs : dict
Kwargs passed to `make_dataloader()` if test_set is a Dataset, not
DataLoader. NOTE: loader_kwargs["ckpt_prefix"] gets automatically
overwritten to None (so that the test DataLoader is not added to
the checkpointer).
Returns
-------
average test loss
"""
if progressbar is None:
progressbar = self.progressbar
if not isinstance(test_set, DataLoader):
test_loader_kwargs["ckpt_prefix"] = None
test_set = self.make_dataloader(test_set, Stage.TEST,
**test_loader_kwargs)
self.on_evaluate_start(max_key=max_key, min_key=min_key)
self.on_stage_start(Stage.TEST, epoch=None)
self.modules.eval()
avg_test_loss = 0.0
with torch.no_grad():
for batch in tqdm(test_set,
dynamic_ncols=True,
disable=not progressbar):
self.step += 1
loss = self.evaluate_batch(batch, stage=Stage.TEST)
avg_test_loss = self.update_average(loss, avg_test_loss)
# Debug mode only runs a few batches
if self.debug and self.step == self.debug_batches:
break
# Only run evaluation "on_stage_end" on main process
run_on_main(self.on_stage_end,
args=[Stage.TEST, avg_test_loss, None])
self.step = 0
def from_hparams(
cls,
source,
hparams_file="hyperparams.yaml",
overrides={},
savedir=None,
use_auth_token=False,
**kwargs,
):
"""Fetch and load based from outside source based on HyperPyYAML file
The source can be a location on the filesystem or online/huggingface
The hyperparams file should contain a "modules" key, which is a
dictionary of torch modules used for computation.
The hyperparams file should contain a "pretrainer" key, which is a
speechbrain.utils.parameter_transfer.Pretrainer
Arguments
---------
source : str
The location to use for finding the model. See
``speechbrain.pretrained.fetching.fetch`` for details.
hparams_file : str
The name of the hyperparameters file to use for constructing
the modules necessary for inference. Must contain two keys:
"modules" and "pretrainer", as described.
overrides : dict
Any changes to make to the hparams file when it is loaded.
savedir : str or Path
Where to put the pretraining material. If not given, will use
./pretrained_models/<class-name>-hash(source).
use_auth_token : bool (default: False)
If true Hugginface's auth_token will be used to load private models from the HuggingFace Hub,
default is False because majority of models are public.
"""
if savedir is None:
clsname = cls.__name__
savedir = f"./pretrained_models/{clsname}-{hash(source)}"
hparams_local_path = fetch(hparams_file, source, savedir,
use_auth_token)
# Load the modules:
with open(hparams_local_path) as fin:
hparams = load_hyperpyyaml(fin, overrides)
# Pretraining:
pretrainer = hparams["pretrainer"]
pretrainer.set_collect_in(savedir)
# For distributed setups, have this here:
run_on_main(pretrainer.collect_files,
kwargs={"default_source": source})
# Load on the CPU. Later the params can be moved elsewhere by specifying
# run_opts={"device": ...}
pretrainer.load_collected(device="cpu")
# Now return the system
return cls(hparams["modules"], hparams, **kwargs)
"forward",
"follow",
"learn",
"visual",
]
else:
raise ValueError("number_of_commands must be 12 or 35")
# Data preparation
run_on_main(
prepare_GSC,
kwargs={
"data_folder": hparams["data_folder"],
"save_folder": hparams["output_folder"],
"validation_percentage": hparams["validation_percentage"],
"testing_percentage": hparams["testing_percentage"],
"percentage_unknown": hparams["percentage_unknown"],
"percentage_silence": hparams["percentage_silence"],
"words_wanted": words_wanted,
"skip_prep": hparams["skip_prep"],
},
)
# Dataset IO prep: creating Dataset objects and proper encodings for phones
train_data, valid_data, test_data, label_encoder = dataio_prep(hparams)
# Brain class initialization
speaker_brain = SpeakerBrain(
modules=hparams["modules"],
opt_class=hparams["opt_class"],
hparams=hparams,
hparams["base_folder_dm"]):
print(
"Please, specify a valid base_folder_dm folder when using dynamic mixing"
)
sys.exit(1)
# Data preparation for LibriMix
from recipes.LibriMix.prepare_data import prepare_librimix as prepare_libri
# create the csv files
run_on_main(
prepare_libri,
kwargs={
"datapath": hparams["data_folder"],
"savepath": hparams["save_folder"],
"n_spks": hparams["num_spks"],
"skip_prep": hparams["skip_prep"],
"fs": hparams["sample_rate"],
"librimix_addnoise": hparams["use_wham_noise"],
},
)
# Data preparation for WHAMR
from recipes.WHAMandWHAMR.prepare_data import create_wham_whamr_csv
from recipes.WHAMandWHAMR.separation.train import (
dataio_prep as dataio_prep_whamr, )
create_wham_whamr_csv(
datapath=hparams["whamr_data_folder"],
savepath=hparams["save_folder"],
fs=hparams["sample_rate"],
# Create experiment directory
sb.create_experiment_directory(
experiment_directory=hparams["output_folder"],
hyperparams_to_save=hparams_file,
overrides=overrides,
)
# Due to DDP, we do the preparation ONLY on the main python process
run_on_main(
prepare_common_voice,
kwargs={
"data_folder": hparams["data_folder"],
"save_folder": hparams["save_folder"],
"train_tsv_file": hparams["train_tsv_file"],
"dev_tsv_file": hparams["dev_tsv_file"],
"test_tsv_file": hparams["test_tsv_file"],
"accented_letters": hparams["accented_letters"],
"language": hparams["language"],
"skip_prep": hparams["skip_prep"],
},
)
# Create the datasets objects as well as tokenization and encoding :-D
train_data, valid_data, test_data, tokenizer = dataio_prepare(hparams)
# Trainer initialization
asr_brain = ASR(
modules=hparams["modules"],
hparams=hparams,
run_opts=run_opts,
)
# Check if wsj0_tr is set with dynamic mixing
if hparams["dynamic_mixing"] and not os.path.exists(hparams["wsj0_tr"]):
print(
"Please, specify a valid wsj0_tr folder when using dynamic mixing")
sys.exit(1)
# Data preparation
from recipes.WSJ0Mix.prepare_data import prepare_wsjmix # noqa
run_on_main(
prepare_wsjmix,
kwargs={
"datapath": hparams["data_folder"],
"savepath": hparams["save_folder"],
"n_spks": hparams["num_spks"],
"skip_prep": hparams["skip_prep"],
},
)
# Create dataset objects
if hparams["dynamic_mixing"]:
if hparams["num_spks"] == 2:
from dynamic_mixing import dynamic_mix_data_prep # noqa
train_data = dynamic_mix_data_prep(hparams)
elif hparams["num_spks"] == 3:
from dynamic_mixing import dynamic_mix_data_prep_3mix # noqa
# Load hyperparameters file with command-line overrides
hparams_file, run_opts, overrides = sb.parse_arguments(sys.argv[1:])
with open(hparams_file) as fin:
hparams = load_hyperpyyaml(fin, overrides)
# Initialize ddp (useful only for multi-GPU DDP training)
sb.utils.distributed.ddp_init_group(run_opts)
# Data preparation
from voicebank_prepare import prepare_voicebank # noqa
run_on_main(
prepare_voicebank,
kwargs={
"data_folder": hparams["data_folder"],
"save_folder": hparams["data_folder"],
"skip_prep": hparams["skip_prep"],
},
)
# Create dataset objects
datasets = dataio_prep(hparams)
# Create experiment directory
sb.create_experiment_directory(
experiment_directory=hparams["output_folder"],
hyperparams_to_save=hparams_file,
overrides=overrides,
)
if hparams["use_tensorboard"]:
def fit(
self,
epoch_counter,
train_set,
valid_set=None,
progressbar=None,
train_loader_kwargs={},
valid_loader_kwargs={},
):
"""Iterate epochs and datasets to improve objective.
Relies on the existence of multiple functions that can (or should) be
overridden. The following methods are used and expected to have a
certain behavior:
* ``fit_batch()``
* ``evaluate_batch()``
* ``update_average()``
If the initialization was done with distributed_count > 0 and the
distributed_backend is ddp, this will generally handle multiprocess
logic, like splitting the training data into subsets for each device and
only saving a checkpoint on the main process.
Arguments
---------
epoch_counter : iterable
Each call should return an integer indicating the epoch count.
train_set : Dataset, DataLoader
A set of data to use for training. If a Dataset is given, a
DataLoader is automatically created. If a DataLoader is given, it is
used directly.
valid_set : Dataset, DataLoader
A set of data to use for validation. If a Dataset is given, a
DataLoader is automatically created. If a DataLoader is given, it is
used directly.
train_loader_kwargs : dict
Kwargs passed to `make_dataloader()` for making the train_loader
(if train_set is a Dataset, not DataLoader).
E.G. batch_size, num_workers.
DataLoader kwargs are all valid.
valid_loader_kwargs : dict
Kwargs passed to `make_dataloader()` for making the valid_loader
(if valid_set is a Dataset, not DataLoader).
E.g., batch_size, num_workers.
DataLoader kwargs are all valid.
progressbar : bool
Whether to display the progress of each epoch in a progressbar.
"""
if not isinstance(train_set, DataLoader):
train_set = self.make_dataloader(
train_set, stage=sb.Stage.TRAIN, **train_loader_kwargs
)
if valid_set is not None and not isinstance(valid_set, DataLoader):
valid_set = self.make_dataloader(
valid_set,
stage=sb.Stage.VALID,
ckpt_prefix=None,
**valid_loader_kwargs,
)
self.on_fit_start()
if progressbar is None:
progressbar = not self.noprogressbar
# Iterate epochs
for epoch in epoch_counter:
# Training stage
self.on_stage_start(Stage.TRAIN, epoch)
self.modules.train()
# Reset nonfinite count to 0 each epoch
self.nonfinite_count = 0
if self.train_sampler is not None and hasattr(
self.train_sampler, "set_epoch"
):
self.train_sampler.set_epoch(epoch)
# Time since last intra-epoch checkpoint
last_ckpt_time = time.time()
# Only show progressbar if requested and main_process
enable = progressbar and sb.utils.distributed.if_main_process()
with tqdm(
train_set,
initial=self.step,
dynamic_ncols=True,
disable=not enable,
) as t:
for batch in t:
self.step += 1
loss = self.fit_batch(batch)
self.avg_train_loss = self.update_average(
loss, self.avg_train_loss
)
t.set_postfix(train_loss=self.avg_train_loss)
# Debug mode only runs a few batches
if self.debug and self.step == self.debug_batches:
break
if (
self.checkpointer is not None
and self.ckpt_interval_minutes > 0
and time.time() - last_ckpt_time
>= self.ckpt_interval_minutes * 60.0
):
run_on_main(self._save_intra_epoch_ckpt)
last_ckpt_time = time.time()
# Run train "on_stage_end" on all processes
self.on_stage_end(Stage.TRAIN, self.avg_train_loss, epoch)
self.avg_train_loss = 0.0
self.step = 0
# Validation stage
if valid_set is not None:
self.on_stage_start(Stage.VALID, epoch)
self.modules.eval()
avg_valid_loss = 0.0
with torch.no_grad():
for batch in tqdm(
valid_set, dynamic_ncols=True, disable=not enable
):
self.step += 1
loss = self.evaluate_batch(batch, stage=Stage.VALID)
avg_valid_loss = self.update_average(
loss, avg_valid_loss
)
# Debug mode only runs a few batches
if self.debug and self.step == self.debug_batches:
break
# Only run validation "on_stage_end" on main process
self.step = 0
run_on_main(
self.on_stage_end,
args=[Stage.VALID, avg_valid_loss, epoch],
)
# Debug mode only runs a few epochs
if self.debug and epoch == self.debug_epochs:
break
# Check if wsj0_tr is set with dynamic mixing
if hparams["dynamic_mixing"] and not os.path.exists(
hparams["base_folder_dm"]):
print(
"Please, specify a valid base_folder_dm folder when using dynamic mixing"
)
sys.exit(1)
# Data preparation
from recipes.WHAMandWHAMR.prepare_data import prepare_wham_whamr_csv
run_on_main(
prepare_wham_whamr_csv,
kwargs={
"datapath": hparams["data_folder"],
"savepath": hparams["save_folder"],
"skip_prep": hparams["skip_prep"],
"fs": hparams["sample_rate"],
},
)
# if whamr, and we do speedaugment we need to prepare the csv file
if "whamr" in hparams["data_folder"] and hparams["use_speedperturb"]:
from recipes.WHAMandWHAMR.prepare_data import create_whamr_rir_csv
from recipes.WHAMandWHAMR.meta.create_whamr_rirs import create_rirs
# If the Room Impulse Responses do not exist, we create them
if not os.path.exists(hparams["rir_path"]):
print("ing Room Impulse Responses...")
run_on_main(
create_rirs,
# Load hyperparameters file with command-line overrides
params_file, run_opts, overrides = sb.core.parse_arguments(sys.argv[1:])
with open(params_file) as fin:
params = load_hyperpyyaml(fin, overrides)
# Dataset prep (parsing VoxCeleb and annotation into csv files)
from ami_prepare import prepare_ami # noqa
run_on_main(
prepare_ami,
kwargs={
"data_folder": params["data_folder"],
"save_folder": params["save_folder"],
"manual_annot_folder": params["manual_annot_folder"],
"split_type": params["split_type"],
"skip_TNO": params["skip_TNO"],
"mic_type": params["mic_type"],
"vad_type": params["vad_type"],
"max_subseg_dur": params["max_subseg_dur"],
"overlap": params["overlap"],
},
)
# Create experiment directory
sb.core.create_experiment_directory(
experiment_directory=params["output_folder"],
hyperparams_to_save=params_file,
overrides=overrides,
)
# Few more experiment directories (to have cleaner structure)
hparams_file, run_opts, overrides = sb.parse_arguments(sys.argv[1:])
with open(hparams_file) as fin:
hparams = load_hyperpyyaml(fin, overrides)
# If distributed_launch=True then
# create ddp_group with the right communication protocol
sb.utils.distributed.ddp_init_group(run_opts)
# Create experiment directory
sb.create_experiment_directory(
experiment_directory=hparams["output_folder"],
hyperparams_to_save=hparams_file,
overrides=overrides,
)
# 1. # Dataset prep
from prepare import prepare_FSC # noqa
# multi-gpu (ddp) save data preparation
run_on_main(
prepare_FSC,
kwargs={
"data_folder": hparams["data_folder"],
"save_folder": hparams["output_folder"],
"skip_prep": hparams["skip_prep"],
},
)
# Train tokenizer
hparams["tokenizer"]()
hparams = load_hyperpyyaml(fin, overrides)
# Create experiment directory
sb.create_experiment_directory(
experiment_directory=hparams["output_folder"],
hyperparams_to_save=hparams_file,
overrides=overrides,
)
# Prepare data
from voicebank_prepare import prepare_voicebank # noqa E402
run_on_main(
prepare_voicebank,
kwargs={
"data_folder": hparams["data_folder"],
"save_folder": hparams["data_folder"],
},
)
datasets, tokenizer = dataio_prep(hparams)
# Load pretrained models
if "pretrained_path" in hparams:
for model_name, path in hparams["pretrained_path"].items():
hparams["modules"][model_name].load_state_dict(torch.load(path))
if "pretrain_checkpointer" in hparams:
hparams["pretrain_checkpointer"].recover_if_possible()
# Initialize trainer
# Create experiment directory
sb.create_experiment_directory(
experiment_directory=hparams["output_folder"],
hyperparams_to_save=hparams_file,
overrides=overrides,
)
# Dataset prep
from prepare import prepare_FSC # noqa
# multi-gpu (ddp) save data preparation
run_on_main(
prepare_FSC,
kwargs={
"data_folder": hparams["data_folder"],
"save_folder": hparams["output_folder"],
"skip_prep": hparams["skip_prep"],
},
)
# here we create the datasets objects as well as tokenization and encoding
(
train_set,
valid_set,
test_set,
tokenizer,
) = dataio_prepare(hparams)
# We download and pretrain the tokenizer
run_on_main(hparams["pretrainer"].collect_files)
hparams["pretrainer"].load_collected(device=run_opts["device"])
hparams = load_hyperpyyaml(fin, overrides)
# Create experiment directory
sb.create_experiment_directory(
experiment_directory=hparams["output_folder"],
hyperparams_to_save=hparams_file,
overrides=overrides,
)
# Prepare data
from voicebank_prepare import prepare_voicebank # noqa E402
run_on_main(
prepare_voicebank,
kwargs={
"data_folder": hparams["data_folder"],
"save_folder": hparams["data_folder"],
"skip_prep": hparams["skip_prep"],
},
)
# Load pretrained models
for model in ["asr", "enhance", "perceptual"]:
pretrained = f"{model}_pretrained"
if pretrained in hparams:
# We download the model from HuggingFace (by default).
run_on_main(hparams[pretrained].collect_files)
hparams[pretrained].load_collected()
# Switch encoder based on task
if "asr_pretrained" in hparams:
token_encoder = hparams["tokenizer"]
hparams = load_hyperpyyaml(fin, overrides)
show_results_every = 100 # plots results every N iterations
# If distributed_launch=True then
# create ddp_group with the right communication protocol
sb.utils.distributed.ddp_init_group(run_opts)
# 1. # Dataset prep (parsing Librispeech)
from prepare import prepare_TAS # noqa
# multi-gpu (ddp) save data preparation
run_on_main(
prepare_TAS,
kwargs={
"data_folder": hparams["data_folder"],
"train_splits": hparams["train_splits"],
"type": "direct",
},
)
# Create experiment directory
sb.create_experiment_directory(
experiment_directory=hparams["output_folder"],
hyperparams_to_save=hparams_file,
overrides=overrides,
)
# here we create the datasets objects as well as tokenization and encoding
(
train_set,
valid_set,
)
# here we create the datasets objects as well as tokenization and encoding
train_dataloader, enrol_dataloader, test_dataloader = dataio_prep(params)
# Initialize PLDA vars
modelset, segset = [], []
embeddings = numpy.empty(shape=[0, params["emb_dim"]], dtype=numpy.float64)
# Embedding file for train data
xv_file = os.path.join(params["save_folder"],
"VoxCeleb1_train_embeddings_stat_obj.pkl")
# We download the pretrained LM from HuggingFace (or elsewhere depending on
# the path given in the YAML file). The tokenizer is loaded at the same time.
run_on_main(params["pretrainer"].collect_files)
params["pretrainer"].load_collected()
params["embedding_model"].eval()
params["embedding_model"].to(params["device"])
# Computing training embeddings (skip it of if already extracted)
if not os.path.exists(xv_file):
logger.info("Extracting embeddings from Training set..")
with tqdm(train_dataloader, dynamic_ncols=True) as t:
for batch in t:
snt_id = batch.id
wav, lens = batch.sig
spk_ids = batch.spk_id
# Flattening speaker ids
def __init__(
self,
model_dir,
vocab_size,
annotation_train=None,
annotation_read=None,
model_type="unigram",
char_format_input=False,
character_coverage=1.0,
user_defined_symbols=None,
max_sentencepiece_length=10,
bos_id=-1,
eos_id=-1,
pad_id=-1,
unk_id=0,
split_by_whitespace=True,
num_sequences=None,
annotation_list_to_check=None,
annotation_format="csv",
):
if model_type not in ["unigram", "bpe", "char"]:
raise ValueError(
"model_type must be one of : [unigram, bpe, char]")
if not os.path.isdir(model_dir):
os.makedirs(model_dir)
if not isinstance(vocab_size, int):
raise ValueError("vocab_size must be integer.")
self.annotation_train = annotation_train
self.annotation_read = annotation_read
self.annotation_format = annotation_format
if self.annotation_train is not None:
ext = os.path.splitext(self.annotation_train)[1]
self.text_file = self.annotation_train.replace(ext, ".txt")
self.prefix_model_file = os.path.join(
model_dir,
str(vocab_size) + "_" + model_type)
self.vocab_size = str(vocab_size)
self.model_type = model_type
self.char_format_input = char_format_input
self.character_coverage = str(character_coverage)
self.max_sentencepiece_length = str(max_sentencepiece_length)
self.bos_id = str(bos_id)
self.eos_id = str(eos_id)
self.pad_id = str(pad_id)
self.unk_id = str(unk_id)
self.num_sequences = num_sequences
self.split_by_whitespace = split_by_whitespace
self.user_defined_symbols = user_defined_symbols
if not os.path.isfile(self.prefix_model_file + ".model"):
logger.info("Train tokenizer with type:" + self.model_type)
if not os.path.isfile(self.text_file):
if annotation_format == "csv":
run_on_main(self._csv2text)
elif annotation_format == "json":
run_on_main(self._json2text)
else:
raise ValueError(
"Annotation format not supported. Supported formats are csv and json. Got "
+ annotation_format)
run_on_main(self._train_BPE)
else:
logger.info("Tokenizer is already trained.")
logger.info("==== Loading Tokenizer ===")
logger.info("Tokenizer path: " + self.prefix_model_file + ".model")
logger.info("Tokenizer vocab_size: " + str(self.vocab_size))
logger.info("Tokenizer type: " + self.model_type)
self.sp = spm.SentencePieceProcessor()
self.sp.load(self.prefix_model_file + ".model")
if annotation_list_to_check is not None:
run_on_main(
self._check_coverage_from_bpe,
kwargs={"list_annotation_files": annotation_list_to_check},
)
# Create experiment directory
sb.create_experiment_directory(
experiment_directory=hparams["output_folder"],
hyperparams_to_save=hparams_file,
overrides=overrides,
)
# Dataset prep (parsing TAS)
from prepare import prepare_TAS # noqa
# multi-gpu (ddp) save data preparation
run_on_main(
prepare_TAS,
kwargs={
"data_folder": hparams["data_folder"],
"train_splits": hparams["train_splits"],
"type": "multistage",
"skip_prep": hparams["skip_prep"],
},
)
# here we create the datasets objects as well as tokenization and encoding
(
train_set,
valid_set,
test_real_set,
test_synth_set,
tokenizer,
) = dataio_prepare(hparams)
# Brain class initialization
if hparams["dynamic_mixing"] and not os.path.exists(
hparams["base_folder_dm"]):
print(
"Please, specify a valid base_folder_dm folder when using dynamic mixing"
)
sys.exit(1)
# Data preparation
from recipes.LibriMix.prepare_data import prepare_librimix
run_on_main(
prepare_librimix,
kwargs={
"datapath": hparams["data_folder"],
"savepath": hparams["save_folder"],
"n_spks": hparams["num_spks"],
"skip_prep": hparams["skip_prep"],
"librimix_addnoise": hparams["use_wham_noise"],
"fs": hparams["sample_rate"],
},
)
# Create dataset objects
if hparams["dynamic_mixing"]:
from dynamic_mixing import (
dynamic_mix_data_prep_librimix as dynamic_mix_data_prep, )
# if the base_folder for dm is not processed, preprocess them
if "processed" not in hparams["base_folder_dm"]:
# if the processed folder already exists we just use it otherwise we do the preprocessing
if not os.path.exists(
if __name__ == "__main__":
# CLI:
hparams_file, run_opts, overrides = sb.parse_arguments(sys.argv[1:])
with open(hparams_file) as fin:
hparams = load_hyperpyyaml(fin, overrides)
# If distributed_launch=True then
# create ddp_group with the right communication protocol
sb.utils.distributed.ddp_init_group(run_opts)
# 1. # Dataset prep
from prepare import prepare_FSC # noqa
# multi-gpu (ddp) save data preparation
run_on_main(
prepare_FSC,
kwargs={"data_folder": hparams["data_folder"]},
)
# Create experiment directory
sb.create_experiment_directory(
experiment_directory=hparams["output_folder"],
hyperparams_to_save=hparams_file,
overrides=overrides,
)
# Train tokenizer
hparams["tokenizer"]()
# Create experiment directory
sb.create_experiment_directory(
experiment_directory=hparams["output_folder"],
hyperparams_to_save=hparams_file,
overrides=overrides,
)
# multi-gpu (ddp) save data preparation
run_on_main(
prepare_librispeech,
kwargs={
"data_folder": hparams["data_folder"],
"tr_splits": hparams["train_splits"],
"dev_splits": hparams["dev_splits"],
"te_splits": hparams["test_splits"],
"save_folder": hparams["data_folder"],
"merge_lst": hparams["train_splits"],
"merge_name": hparams["train_csv"],
"skip_prep": hparams["skip_prep"],
},
)
# here we create the datasets objects as well as tokenization and encoding
train_data, valid_data, test_datasets, tokenizer = dataio_prepare(hparams)
# We download the pretrained LM from HuggingFace (or elsewhere depending on
# the path given in the YAML file). The tokenizer is loaded at the same time.
run_on_main(hparams["pretrainer"].collect_files)
hparams["pretrainer"].load_collected(device=run_opts["device"])
# If distributed_launch=True then
# create ddp_group with the right communication protocol
sb.utils.distributed.ddp_init_group(run_opts)
# 1. # Dataset prep (parsing Librispeech)
from librispeech_prepare import prepare_librispeech # noqa
# multi-gpu (ddp) save data preparation
run_on_main(
prepare_librispeech,
kwargs={
"data_folder": hparams["data_folder"],
"tr_splits": hparams["train_splits"],
"dev_splits": hparams["dev_splits"],
"te_splits": hparams["test_splits"],
"save_folder": hparams["data_folder"],
"merge_lst": hparams["train_splits"],
"merge_name": hparams["train_csv"],
},
)
# Create experiment directory
sb.create_experiment_directory(
experiment_directory=hparams["output_folder"],
hyperparams_to_save=hparams_file,
overrides=overrides,
)
# Train tokenizer
hparams["tokenizer"]()
# Data preparation
from dns_prepare import prepare_dns # noq
# Create experiment directory
sb.create_experiment_directory(
experiment_directory=hparams["output_folder"],
hyperparams_to_save=hparams_file,
overrides=overrides,
)
run_on_main(
prepare_dns,
kwargs={
"data_folder": hparams["data_folder"],
"save_folder": hparams["save_folder"],
"valid_folder": hparams["valid_folder"],
"seg_size": 10.0,
"skip_prep": hparams["skip_prep"],
},
)
if hparams["use_tensorboard"]:
from speechbrain.utils.train_logger import TensorboardLogger
hparams["hparams"]["tensorboard_train_logger"] = TensorboardLogger(
hparams["tensorboard_logs"])
# Create the folder to save enhanced files (+ support for DDP)
try:
# all writing command must be done with the main_process
if sb.utils.distributed.if_main_process():
from librispeech_prepare import prepare_librispeech # noqa
# Create experiment directory
sb.create_experiment_directory(
experiment_directory=hparams["output_folder"],
hyperparams_to_save=hparams_file,
overrides=overrides,
)
# multi-gpu (ddp) save data preparation
run_on_main(
prepare_librispeech,
kwargs={
"data_folder": hparams["data_folder"],
"save_folder": hparams["save_folder"],
"create_lexicon": True,
"skip_prep": hparams["skip_prep"],
},
)
# Dataset IO prep: creating Dataset objects and proper encodings for phones
train_data, valid_data, test_data, phoneme_encoder = dataio_prep(hparams)
# Trainer initialization
asr_brain = ASR(
modules=hparams["modules"],
opt_class=hparams["opt_class"],
hparams=hparams,
run_opts=run_opts,
checkpointer=hparams["checkpointer"],
# Create experiment directory
sb.create_experiment_directory(
experiment_directory=hparams["output_folder"],
hyperparams_to_save=hparams_file,
overrides=overrides,
)
if hparams["use_tensorboard"]:
from speechbrain.utils.train_logger import TensorboardLogger
hparams["tensorboard_train_logger"] = TensorboardLogger(
hparams["tensorboard_logs"]
)
# Create the folder to save enhanced files (+ support for DDP)
run_on_main(create_folder, kwargs={"folder": hparams["enhanced_folder"]})
se_brain = MetricGanBrain(
modules=hparams["modules"],
hparams=hparams,
run_opts=run_opts,
checkpointer=hparams["checkpointer"],
)
se_brain.train_set = datasets["train"]
se_brain.historical_set = {}
se_brain.noisy_scores = {}
se_brain.batch_size = hparams["dataloader_options"]["batch_size"]
se_brain.valid_batch_size = hparams["valid_dataloader_options"]["batch_size"]
se_brain.sub_stage = SubStage.GENERATOR
# shutil.rmtree(hparams["MetricGAN_folder"])
experiment_directory=hparams["output_folder"],
hyperparams_to_save=hparams_file,
overrides=overrides,
)
# Initialize ddp (useful only for multi-GPU DDP training)
sb.utils.distributed.ddp_init_group(run_opts)
# Dataset prep (parsing VoxCeleb and annotation into csv files)
from voxceleb_prepare import prepare_voxceleb # noqa
run_on_main(
prepare_voxceleb,
kwargs={
"data_folder": hparams["data_folder"],
"save_folder": hparams["data_folder"],
"splits": ["train", "dev", "test"],
"split_ratio": [90, 10],
"seg_dur": int(hparams["sentence_len"]) * 100,
},
)
(
train_dataloader,
valid_dataloader,
enrol_dataloader,
test_dataloader,
) = dataio_prep(hparams)
# Dictionary to store the last waveform read for each speaker
wav_stored = {}
# Initialize ddp (useful only for multi-GPU DDP training)
sb.utils.distributed.ddp_init_group(run_opts)
# Create experiment directory
sb.create_experiment_directory(
experiment_directory=hparams["output_folder"],
hyperparams_to_save=hparams_file,
overrides=overrides,
)
# multi-gpu (ddp) save data preparation
run_on_main(
prepare_timit,
kwargs={
"data_folder": hparams["data_folder"],
"splits": ["train", "dev", "test"],
"save_folder": hparams["data_folder"],
"phn_set": hparams["phn_set"],
"skip_prep": hparams["skip_prep"],
},
)
# Dataset IO prep: creating Dataset objects and proper encodings for phones
train_data, valid_data, test_data, label_encoder = dataio_prep(hparams)
# Trainer initialization
align_brain = AlignBrain(
modules=hparams["modules"],
opt_class=hparams["opt_class"],
hparams=hparams,
run_opts=run_opts,
checkpointer=hparams["checkpointer"],
# Create experiment directory
sb.create_experiment_directory(
experiment_directory=hparams["output_folder"],
hyperparams_to_save=hparams_file,
overrides=overrides,
)
# Dataset prep (parsing TAS)
from prepare import prepare_TAS # noqa
# multi-gpu (ddp) save data preparation
run_on_main(
prepare_TAS,
kwargs={
"data_folder": hparams["data_folder"],
"save_folder": hparams["output_folder"],
"train_splits": hparams["train_splits"],
"type": "multistage",
"skip_prep": hparams["skip_prep"],
},
)
# here we create the datasets objects as well as tokenization and encoding
(
train_set,
valid_set,
test_real_set,
test_synth_set,
all_real_set,
tokenizer,
) = dataio_prepare(hparams)
# Create experiment directory
sb.create_experiment_directory(
experiment_directory=hparams["output_folder"],
hyperparams_to_save=hparams_file,
overrides=overrides,
)
# Dataset prep (parsing SLURP)
from prepare import prepare_SLURP # noqa
# multi-gpu (ddp) save data preparation
run_on_main(
prepare_SLURP,
kwargs={
"data_folder": hparams["data_folder"],
"save_folder": hparams["output_folder"],
"train_splits": hparams["train_splits"],
"slu_type": "direct",
"skip_prep": hparams["skip_prep"],
},
)
# here we create the datasets objects as well as tokenization and encoding
(
train_set,
valid_set,
test_set,
tokenizer,
) = dataio_prepare(hparams)
# We download and pretrain the tokenizer
run_on_main(hparams["pretrainer"].collect_files)
# Initialize ddp (useful only for multi-GPU DDP training)
sb.utils.distributed.ddp_init_group(run_opts)
# Create experiment directory
sb.create_experiment_directory(
experiment_directory=hparams["output_folder"],
hyperparams_to_save=hparams_file,
overrides=overrides,
)
# multi-gpu (ddp) save data preparation
run_on_main(
prepare_timit,
kwargs={
"data_folder": hparams["data_folder"],
"save_json_train": hparams["train_annotation"],
"save_json_valid": hparams["valid_annotation"],
"save_json_test": hparams["test_annotation"],
"skip_prep": hparams["skip_prep"],
},
)
# Dataset IO prep: creating Dataset objects and proper encodings for phones
train_data, valid_data, test_data, label_encoder = dataio_prep(hparams)
# Trainer initialization
asr_brain = ASR_Brain(
modules=hparams["modules"],
opt_class=hparams["opt_class"],
hparams=hparams,
run_opts=run_opts,
checkpointer=hparams["checkpointer"],
