|
|
import sys |
|
|
|
|
|
import torch |
|
|
import speechbrain as sb |
|
|
from speechbrain.dataio import dataset |
|
|
from speechbrain.utils.distributed import run_on_main |
|
|
from hyperpyyaml import load_hyperpyyaml |
|
|
|
|
|
|
|
|
class LM(sb.core.Brain): |
|
|
def compute_forward(self, batch, stage): |
|
|
batch = batch.to(self.device) |
|
|
tokens_bos, _ = batch.tokens_bos |
|
|
logits = self.hparams.model(tokens_bos) |
|
|
pred = self.hparams.log_softmax(logits) |
|
|
return pred |
|
|
|
|
|
def compute_objectives(self, predictions, batch, stage): |
|
|
batch = batch.to(self.device) |
|
|
tokens_eos, tokens_len = batch.tokens_eos |
|
|
loss = self.hparams.compute_cost( |
|
|
predictions, tokens_eos, length=tokens_len |
|
|
) |
|
|
return loss |
|
|
|
|
|
def fit_batch(self, batch): |
|
|
predictions = self.compute_forward(batch, sb.Stage.TRAIN) |
|
|
loss = self.compute_objectives(predictions, batch, sb.Stage.TRAIN) |
|
|
|
|
|
(loss / self.hparams.accumulation_steps).backward() |
|
|
|
|
|
if self.step % self.hparams.accumulation_steps == 0: |
|
|
self.check_gradients(loss) |
|
|
|
|
|
self.optimizer.step() |
|
|
self.optimizer.zero_grad() |
|
|
|
|
|
if isinstance( |
|
|
self.hparams.lr_annealing, sb.nnet.schedulers.NoamScheduler |
|
|
) or isinstance( |
|
|
self.hparams.lr_annealing, |
|
|
sb.nnet.schedulers.CyclicCosineScheduler, |
|
|
): |
|
|
self.hparams.lr_annealing(self.optimizer) |
|
|
|
|
|
return loss |
|
|
|
|
|
def on_stage_end(self, stage, stage_loss, epoch): |
|
|
stage_stats = {"loss": stage_loss} |
|
|
if stage == sb.Stage.TRAIN: |
|
|
self.train_stats = stage_stats |
|
|
|
|
|
if stage == sb.Stage.VALID and sb.utils.distributed.if_main_process(): |
|
|
if not ( |
|
|
isinstance( |
|
|
self.hparams.lr_annealing, sb.nnet.schedulers.NoamScheduler |
|
|
) |
|
|
or isinstance( |
|
|
self.hparams.lr_annealing, |
|
|
sb.nnet.schedulers.CyclicCosineScheduler, |
|
|
) |
|
|
): |
|
|
old_lr, new_lr = self.hparams.lr_annealing(stage_loss) |
|
|
sb.nnet.schedulers.update_learning_rate(self.optimizer, new_lr) |
|
|
else: |
|
|
old_lr = self.hparams.lr_annealing.current_lr |
|
|
|
|
|
self.hparams.train_logger.log_stats( |
|
|
stats_meta={"epoch": epoch, "lr": old_lr}, |
|
|
train_stats=self.train_stats, |
|
|
valid_stats=stage_stats, |
|
|
) |
|
|
self.checkpointer.save_and_keep_only( |
|
|
meta=stage_stats, min_keys=["loss"], |
|
|
) |
|
|
|
|
|
if stage == sb.Stage.TEST and sb.utils.distributed.if_main_process(): |
|
|
self.hparams.train_logger.log_stats( |
|
|
stats_meta={"Epoch loaded": self.hparams.epoch_counter.current}, |
|
|
test_stats=stage_stats, |
|
|
) |
|
|
|
|
|
|
|
|
def dataio_prepare(hparams): |
|
|
@sb.utils.data_pipeline.takes("transcription") |
|
|
@sb.utils.data_pipeline.provides( |
|
|
"transcription", "tokens_bos", "tokens_eos" |
|
|
) |
|
|
def transcription_pipline(transcription): |
|
|
yield transcription |
|
|
tokens_list = hparams["tokenizer"].encode_as_ids(transcription) |
|
|
tokens_bos = torch.LongTensor([hparams["bos_index"]] + (tokens_list)) |
|
|
yield tokens_bos |
|
|
tokens_eos = torch.LongTensor(tokens_list + [hparams["eos_index"]]) |
|
|
yield tokens_eos |
|
|
|
|
|
data_folder = hparams["data_folder"] |
|
|
datasets = {} |
|
|
for dataset_name in ["train", "dev", "test"]: |
|
|
json_path = f"{data_folder}/{dataset_name}.json" |
|
|
datasets[dataset_name] = dataset.DynamicItemDataset.from_json( |
|
|
json_path=json_path, |
|
|
replacements={"data_root": data_folder}, |
|
|
dynamic_items=[transcription_pipline], |
|
|
output_keys=["transcription", "tokens_bos", "tokens_eos"], |
|
|
) |
|
|
|
|
|
return datasets |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
hparams_file_path, run_opts, overrides = sb.parse_arguments(sys.argv[1:]) |
|
|
with open(hparams_file_path) as hparams_file: |
|
|
hparams = load_hyperpyyaml(hparams_file, overrides) |
|
|
|
|
|
sb.utils.distributed.ddp_init_group(run_opts) |
|
|
|
|
|
sb.create_experiment_directory( |
|
|
experiment_directory=hparams["output_folder"], |
|
|
hyperparams_to_save=hparams_file_path, |
|
|
overrides=overrides, |
|
|
) |
|
|
|
|
|
run_on_main(hparams["pretrainer"].collect_files) |
|
|
hparams["pretrainer"].load_collected(device=run_opts["device"]) |
|
|
|
|
|
datasets = dataio_prepare(hparams) |
|
|
|
|
|
lm_brain = LM( |
|
|
modules=hparams["modules"], |
|
|
opt_class=hparams["optimizer"], |
|
|
hparams=hparams, |
|
|
run_opts=run_opts, |
|
|
checkpointer=hparams["checkpointer"], |
|
|
) |
|
|
|
|
|
lm_brain.fit( |
|
|
lm_brain.hparams.epoch_counter, |
|
|
datasets["train"], |
|
|
datasets["dev"], |
|
|
train_loader_kwargs=hparams["train_dataloader_opts"], |
|
|
valid_loader_kwargs=hparams["valid_dataloader_opts"], |
|
|
) |
|
|
|
|
|
|
|
|
lm_brain.evaluate( |
|
|
datasets["test"], |
|
|
min_key="loss", |
|
|
test_loader_kwargs=hparams["test_dataloader_opts"], |
|
|
) |
|
|
|
