Architectures/ToucanTTS/toucantts_train_loop.py

import os
import time

import torch
import torch.multiprocessing
import wandb
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data.dataloader import DataLoader
from tqdm import tqdm

from Preprocessing.AudioPreprocessor import AudioPreprocessor
from Preprocessing.EnCodecAudioPreprocessor import CodecAudioPreprocessor
from Utility.WarmupScheduler import ToucanWarmupScheduler as WarmupScheduler
from Utility.utils import delete_old_checkpoints
from Utility.utils import get_most_recent_checkpoint
from Utility.utils import plot_progress_spec_toucantts
from run_weight_averaging import average_checkpoints
from run_weight_averaging import get_n_recent_checkpoints_paths
from run_weight_averaging import load_net_toucan
from run_weight_averaging import save_model_for_use


def collate_and_pad(batch):
    # text, text_len, speech, speech_len, durations, energy, pitch, utterance condition, language_id, speaker embedding
    return (pad_sequence([datapoint[0] for datapoint in batch], batch_first=True).float(),
            torch.stack([datapoint[1] for datapoint in batch]).squeeze(1),
            [datapoint[2] for datapoint in batch],
            torch.stack([datapoint[3] for datapoint in batch]).squeeze(1),
            pad_sequence([datapoint[4] for datapoint in batch], batch_first=True),
            pad_sequence([datapoint[5] for datapoint in batch], batch_first=True),
            pad_sequence([datapoint[6] for datapoint in batch], batch_first=True),
            None,
            torch.stack([datapoint[8] for datapoint in batch]),
            torch.stack([datapoint[9] for datapoint in batch]))


def train_loop(net,
               train_dataset,
               device,
               save_directory,
               batch_size,
               lang,
               lr,
               warmup_steps,
               path_to_checkpoint,
               fine_tune,
               resume,
               steps,
               use_wandb,
               train_sampler,
               gpu_count,
               steps_per_checkpoint
               ):
    """
    see train loop arbiter for explanations of the arguments
    """
    net = net.to(device)
    if gpu_count > 1:
        rank = int(os.environ["LOCAL_RANK"])
    else:
        rank = 0
    if steps_per_checkpoint is None:
        steps_per_checkpoint = len(train_dataset) // batch_size

    if steps < warmup_steps * 5:
        print(f"too much warmup given the amount of steps, reducing warmup to {warmup_steps} steps")
        warmup_steps = steps // 5

    torch.multiprocessing.set_sharing_strategy('file_system')
    batch_sampler_train = torch.utils.data.BatchSampler(train_sampler, batch_size, drop_last=True)
    train_loader = DataLoader(dataset=train_dataset,
                              batch_sampler=batch_sampler_train,
                              num_workers=0,  # has to be 0, otherwise copies of the dataset are created, which is not feasible for large scale trainings. This is not optimal for small trainings, but necessary for scalability.
                              pin_memory=True,
                              prefetch_factor=None,
                              collate_fn=collate_and_pad)
    ap = CodecAudioPreprocessor(input_sr=-1, device=device)
    spec_extractor = AudioPreprocessor(input_sr=16000, output_sr=16000, device=device)

    step_counter = 0

    if isinstance(net, torch.nn.parallel.DistributedDataParallel):
        model = net.module
    else:
        model = net
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)

    scheduler = WarmupScheduler(optimizer, peak_lr=lr, warmup_steps=warmup_steps, max_steps=steps)

    epoch = 0
    if resume:
        path_to_checkpoint = get_most_recent_checkpoint(checkpoint_dir=save_directory)
    if path_to_checkpoint is not None:
        check_dict = torch.load(path_to_checkpoint, map_location=device)
        model.load_state_dict(check_dict["model"])
        if not fine_tune:
            optimizer.load_state_dict(check_dict["optimizer"])
            scheduler.load_state_dict(check_dict["scheduler"])
            step_counter = check_dict["step_counter"]
    start_time = time.time()
    regression_losses_total = list()
    stochastic_losses_total = list()
    duration_losses_total = list()
    pitch_losses_total = list()
    energy_losses_total = list()
    while True:
        net.train()
        epoch += 1
        for batch in tqdm(train_loader):

            text_tensors = batch[0].to(device)
            text_lengths = batch[1].squeeze().to(device)
            speech_indexes = batch[2]
            speech_lengths = batch[3].squeeze().to(device)
            gold_durations = batch[4].to(device)
            gold_pitch = batch[6].to(device)  # mind the switched order
            gold_energy = batch[5].to(device)  # mind the switched order
            lang_ids = batch[8].squeeze(1).to(device)

            speech_batch = list()  # I wish this could be done in the collate function or in the getitem, but using DL models in multiprocessing on very large datasets causes just way too many issues.
            for speech_sample in speech_indexes:
                with torch.inference_mode():
                    wave = ap.indexes_to_audio(speech_sample.int().to(device)).detach()
                    mel = spec_extractor.audio_to_mel_spec_tensor(wave, explicit_sampling_rate=16000).transpose(0, 1).detach().cpu()
                gold_speech_sample = mel.clone()
                speech_batch.append(gold_speech_sample)
            gold_speech = pad_sequence(speech_batch, batch_first=True).to(device)

            run_stochastic = (step_counter > warmup_steps) or fine_tune

            train_loss = 0.0
            utterance_embedding = batch[9].to(device)
            regression_loss, stochastic_loss, duration_loss, pitch_loss, energy_loss = net(
                text_tensors=text_tensors,
                text_lengths=text_lengths,
                gold_speech=gold_speech,
                speech_lengths=speech_lengths,
                gold_durations=gold_durations,
                gold_pitch=gold_pitch,
                gold_energy=gold_energy,
                utterance_embedding=utterance_embedding,
                lang_ids=lang_ids,
                return_feats=False,
                run_stochastic=run_stochastic
            )

            if torch.isnan(regression_loss) or torch.isnan(duration_loss) or torch.isnan(pitch_loss) or torch.isnan(energy_loss):
                print("One of the losses turned to NaN! Skipping this batch ...")
                continue

            train_loss = train_loss + duration_loss
            train_loss = train_loss + pitch_loss
            train_loss = train_loss + energy_loss
            train_loss = train_loss + regression_loss

            regression_losses_total.append(regression_loss.item())
            duration_losses_total.append(duration_loss.item())
            pitch_losses_total.append(pitch_loss.item())
            energy_losses_total.append(energy_loss.item())

            if stochastic_loss is not None:

                if torch.isnan(stochastic_loss):
                    print("Flow loss turned to NaN! Skipping this batch ...")
                    continue

                stochastic_losses_total.append(stochastic_loss.item())
                train_loss = train_loss + stochastic_loss
            else:
                stochastic_losses_total.append(0)

            optimizer.zero_grad()
            if type(train_loss) is float:
                print("There is no loss for this step! Skipping ...")
                continue
            if gpu_count > 1:
                torch.distributed.barrier()
            train_loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0, error_if_nonfinite=False)
            optimizer.step()
            scheduler.step()
            step_counter += 1
            if step_counter % steps_per_checkpoint == 0:
                # evaluation interval is happening
                if rank == 0:
                    net.eval()
                    default_embedding = train_dataset[0][9].to(device)
                    torch.save({
                        "model"       : model.state_dict(),
                        "optimizer"   : optimizer.state_dict(),
                        "step_counter": step_counter,
                        "scheduler"   : scheduler.state_dict(),
                        "default_emb" : default_embedding,
                        "config"      : model.config
                    }, os.path.join(save_directory, "checkpoint_{}.pt".format(step_counter)))

                    delete_old_checkpoints(save_directory, keep=5)

                    print(f"\nEpoch:                  {epoch}")
                    print(f"Time elapsed:           {round((time.time() - start_time) / 60)} Minutes")
                    print("Reconstruction Loss:    {}".format(round(sum(regression_losses_total) / len(regression_losses_total), 3)))
                    print(f"Steps:                  {step_counter}\n")

                    if use_wandb:
                        wandb.log({
                            "regression_loss": round(sum(regression_losses_total) / len(regression_losses_total), 5),
                            "stochastic_loss": round(sum(stochastic_losses_total) / len(stochastic_losses_total), 5),
                            "duration_loss"  : round(sum(duration_losses_total) / len(duration_losses_total), 5),
                            "pitch_loss"     : round(sum(pitch_losses_total) / len(pitch_losses_total), 5),
                            "energy_loss"    : round(sum(energy_losses_total) / len(energy_losses_total), 5),
                            "learning_rate"  : optimizer.param_groups[0]['lr']
                        }, step=step_counter)
                    regression_losses_total = list()
                    stochastic_losses_total = list()
                    duration_losses_total = list()
                    pitch_losses_total = list()
                    energy_losses_total = list()

                    path_to_most_recent_plot = plot_progress_spec_toucantts(model,
                                                                            device,
                                                                            save_dir=save_directory,
                                                                            step=step_counter,
                                                                            lang=lang,
                                                                            default_emb=default_embedding,
                                                                            run_stochastic=run_stochastic)
                    if use_wandb:
                        wandb.log({
                            "progress_plot": wandb.Image(path_to_most_recent_plot)
                        }, step=step_counter)

                    checkpoint_paths = get_n_recent_checkpoints_paths(checkpoint_dir=save_directory, n=1)
                    averaged_model, default_embed = average_checkpoints(checkpoint_paths, load_func=load_net_toucan)
                    save_model_for_use(model=averaged_model, default_embed=default_embed, name=os.path.join(save_directory, "best.pt"))

                    if step_counter > steps:
                        return  # DONE

                    net.train()

        print("\n\n\nEPOCH COMPLETE\n\n\n")