<reponame>rhasspy/tacotron2-train
import logging
import time
import typing
from pathlib import Path

import numpy as np
import torch
from torch.cuda.amp import GradScaler, autocast
from torch.utils.data import DataLoader

from .checkpoint import Checkpoint, save_checkpoint
from .config import TrainingConfig
from .loss_function import LossType
from .models import ModelType, OptimizerType, setup_model
from .utils import to_gpu

_LOGGER = logging.getLogger("tacotron2_train")

# -----------------------------------------------------------------------------


def train(
    train_loader: DataLoader,
    config: TrainingConfig,
    model_dir: Path,
    model: typing.Optional[ModelType] = None,
    optimizer: typing.Optional[OptimizerType] = None,
    global_step: int = 1,
    checkpoint_epochs: int = 1,
    learning_rate: float = 1e-3,
    rank: int = 0,
):
    """Run training for the specified number of epochs"""
    torch.manual_seed(config.seed)

    model, optimizer = setup_model(config, model=model, optimizer=optimizer)

    assert model is not None
    assert optimizer is not None

    criterion = LossType(
        ga_alpha=config.model.guided_attention_alpha,
        ga_sigma=config.model.guided_attention_sigma,
    )
    criterion.cuda()

    # Gradient scaler
    scaler = GradScaler() if config.fp16_run else None

    # Begin training
    for epoch in range(1, config.epochs + 1):
        _LOGGER.debug(
            "Begin epoch %s/%s (global step=%s)", epoch, config.epochs, global_step
        )
        epoch_start_time = time.perf_counter()
        global_step, learning_rate = train_step(
            global_step=global_step,
            epoch=epoch,
            learning_rate=learning_rate,
            config=config,
            model=model,
            optimizer=optimizer,
            criterion=criterion,
            train_loader=train_loader,
            fp16_run=config.fp16_run,
            scaler=scaler,
        )

        if ((epoch % checkpoint_epochs) == 0) and (rank == 0):
            # Save checkpoint
            checkpoint_path = model_dir / f"checkpoint_{global_step}.pth"
            _LOGGER.debug("Saving checkpoint to %s", checkpoint_path)
            save_checkpoint(
                Checkpoint(
                    model=model,
                    optimizer=optimizer,
                    learning_rate=learning_rate,
                    global_step=global_step,
                    version=config.version,
                ),
                checkpoint_path,
            )

            # Save checkpoint config
            config_path = model_dir / f"config_{global_step}.json"
            with open(config_path, "w") as config_file:
                config.save(config_file)

            _LOGGER.info("Saved checkpoint to %s", checkpoint_path)

        epoch_end_time = time.perf_counter()
        _LOGGER.debug(
            "Epoch %s complete in %s second(s) (global step=%s)",
            epoch,
            epoch_end_time - epoch_start_time,
            global_step,
        )


def train_step(
    global_step: int,
    epoch: int,
    learning_rate: float,
    config: TrainingConfig,
    model: ModelType,
    optimizer: OptimizerType,
    criterion: LossType,
    train_loader: DataLoader,
    fp16_run: bool,
    scaler: typing.Optional[GradScaler] = None,
) -> typing.Tuple[int, float]:
    steps_per_epoch = len(train_loader)

    model.train()
    for batch_idx, batch in enumerate(train_loader):
        learning_rate = adjust_learning_rate(global_step, epoch, optimizer, config)

        model.zero_grad()
        (
            text_padded,
            input_lengths,
            mel_padded,
            gate_padded,
            output_lengths,
            _len_x,
        ) = batch

        # Put on GPU
        text_padded = to_gpu(text_padded).long()
        input_lengths = to_gpu(input_lengths).long()
        mel_padded = to_gpu(mel_padded).float()
        gate_padded = to_gpu(gate_padded).float()
        output_lengths = to_gpu(output_lengths).long()

        with autocast(enabled=fp16_run):
            mel_outputs, mel_outputs_postnet, gate_outputs, alignments = model(
                text_padded, input_lengths, mel_padded, output_lengths
            )

            # Adjust for case where n_frames_per_step > 1
            remainder = output_lengths.max() % config.model.n_frames_per_step
            if remainder != 0:
                alignment_lengths = (
                    output_lengths + (config.model.n_frames_per_step - remainder)
                ) // config.model.n_frames_per_step
            else:
                alignment_lengths = output_lengths // config.model.n_frames_per_step

            loss = criterion(
                mel_outputs,
                mel_outputs_postnet,
                gate_outputs,
                mel_padded,
                gate_padded,
                alignments,
                input_lengths,
                alignment_lengths,
            )

        reduced_loss = loss.item()

        if np.isnan(reduced_loss):
            raise RuntimeError(f"loss is NaN at global step {global_step}")

        if fp16_run:
            # Float16
            assert scaler is not None
            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(
                model.parameters(), config.grad_clip_threshold
            )

            scaler.step(optimizer)
            scaler.update()
        else:
            # Float32
            loss.backward()
            torch.nn.utils.clip_grad_norm_(
                model.parameters(), config.grad_clip_threshold
            )
            optimizer.step()

        _LOGGER.debug(
            "Loss: %s (step=%s/%s)", reduced_loss, batch_idx + 1, steps_per_epoch
        )
        global_step += 1

    return global_step, learning_rate


# -----------------------------------------------------------------------------


def adjust_learning_rate(
    iteration: int, epoch: int, optimizer: OptimizerType, config: TrainingConfig
) -> float:

    p = 0
    if config.anneal_steps is not None:
        for a_step in config.anneal_steps:
            if epoch >= int(a_step):
                p = p + 1

    if config.anneal_factor == 0.3:
        lr = config.learning_rate * ((0.1 ** (p // 2)) * (1.0 if p % 2 == 0 else 0.3))
    else:
        lr = config.learning_rate * (config.anneal_factor ** p)

    for param_group in optimizer.param_groups:
        param_group["lr"] = lr

    return lr

import logging
import time
import typing
from pathlib import Path

import numpy as np
import torch
from torch.cuda.amp import GradScaler, autocast
from torch.utils.data import DataLoader

from .checkpoint import Checkpoint, save_checkpoint
from .config import TrainingConfig
from .loss_function import LossType
from .models import ModelType, OptimizerType, setup_model
from .utils import to_gpu

_LOGGER = logging.getLogger("tacotron2_train")

# -----------------------------------------------------------------------------


def train(
    train_loader: DataLoader,
    config: TrainingConfig,
    model_dir: Path,
    model: typing.Optional[ModelType] = None,
    optimizer: typing.Optional[OptimizerType] = None,
    global_step: int = 1,
    checkpoint_epochs: int = 1,
    learning_rate: float = 1e-3,
    rank: int = 0,
):
    """Run training for the specified number of epochs"""
    torch.manual_seed(config.seed)

    model, optimizer = setup_model(config, model=model, optimizer=optimizer)

    assert model is not None
    assert optimizer is not None

    criterion = LossType(
        ga_alpha=config.model.guided_attention_alpha,
        ga_sigma=config.model.guided_attention_sigma,
    )
    criterion.cuda()

    # Gradient scaler
    scaler = GradScaler() if config.fp16_run else None

    # Begin training
    for epoch in range(1, config.epochs + 1):
        _LOGGER.debug(
            "Begin epoch %s/%s (global step=%s)", epoch, config.epochs, global_step
        )
        epoch_start_time = time.perf_counter()
        global_step, learning_rate = train_step(
            global_step=global_step,
            epoch=epoch,
            learning_rate=learning_rate,
            config=config,
            model=model,
            optimizer=optimizer,
            criterion=criterion,
            train_loader=train_loader,
            fp16_run=config.fp16_run,
            scaler=scaler,
        )

        if ((epoch % checkpoint_epochs) == 0) and (rank == 0):
            # Save checkpoint
            checkpoint_path = model_dir / f"checkpoint_{global_step}.pth"
            _LOGGER.debug("Saving checkpoint to %s", checkpoint_path)
            save_checkpoint(
                Checkpoint(
                    model=model,
                    optimizer=optimizer,
                    learning_rate=learning_rate,
                    global_step=global_step,
                    version=config.version,
                ),
                checkpoint_path,
            )

            # Save checkpoint config
            config_path = model_dir / f"config_{global_step}.json"
            with open(config_path, "w") as config_file:
                config.save(config_file)

            _LOGGER.info("Saved checkpoint to %s", checkpoint_path)

        epoch_end_time = time.perf_counter()
        _LOGGER.debug(
            "Epoch %s complete in %s second(s) (global step=%s)",
            epoch,
            epoch_end_time - epoch_start_time,
            global_step,
        )


def train_step(
    global_step: int,
    epoch: int,
    learning_rate: float,
    config: TrainingConfig,
    model: ModelType,
    optimizer: OptimizerType,
    criterion: LossType,
    train_loader: DataLoader,
    fp16_run: bool,
    scaler: typing.Optional[GradScaler] = None,
) -> typing.Tuple[int, float]:
    steps_per_epoch = len(train_loader)

    model.train()
    for batch_idx, batch in enumerate(train_loader):
        learning_rate = adjust_learning_rate(global_step, epoch, optimizer, config)

        model.zero_grad()
        (
            text_padded,
            input_lengths,
            mel_padded,
            gate_padded,
            output_lengths,
            _len_x,
        ) = batch

        # Put on GPU
        text_padded = to_gpu(text_padded).long()
        input_lengths = to_gpu(input_lengths).long()
        mel_padded = to_gpu(mel_padded).float()
        gate_padded = to_gpu(gate_padded).float()
        output_lengths = to_gpu(output_lengths).long()

        with autocast(enabled=fp16_run):
            mel_outputs, mel_outputs_postnet, gate_outputs, alignments = model(
                text_padded, input_lengths, mel_padded, output_lengths
            )

            # Adjust for case where n_frames_per_step > 1
            remainder = output_lengths.max() % config.model.n_frames_per_step
            if remainder != 0:
                alignment_lengths = (
                    output_lengths + (config.model.n_frames_per_step - remainder)
                ) // config.model.n_frames_per_step
            else:
                alignment_lengths = output_lengths // config.model.n_frames_per_step

            loss = criterion(
                mel_outputs,
                mel_outputs_postnet,
                gate_outputs,
                mel_padded,
                gate_padded,
                alignments,
                input_lengths,
                alignment_lengths,
            )

        reduced_loss = loss.item()

        if np.isnan(reduced_loss):
            raise RuntimeError(f"loss is NaN at global step {global_step}")

        if fp16_run:
            # Float16
            assert scaler is not None
            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(
                model.parameters(), config.grad_clip_threshold
            )

            scaler.step(optimizer)
            scaler.update()
        else:
            # Float32
            loss.backward()
            torch.nn.utils.clip_grad_norm_(
                model.parameters(), config.grad_clip_threshold
            )
            optimizer.step()

        _LOGGER.debug(
            "Loss: %s (step=%s/%s)", reduced_loss, batch_idx + 1, steps_per_epoch
        )
        global_step += 1

    return global_step, learning_rate


# -----------------------------------------------------------------------------


def adjust_learning_rate(
    iteration: int, epoch: int, optimizer: OptimizerType, config: TrainingConfig
) -> float:

    p = 0
    if config.anneal_steps is not None:
        for a_step in config.anneal_steps:
            if epoch >= int(a_step):
                p = p + 1

    if config.anneal_factor == 0.3:
        lr = config.learning_rate * ((0.1 ** (p // 2)) * (1.0 if p % 2 == 0 else 0.3))
    else:
        lr = config.learning_rate * (config.anneal_factor ** p)

    for param_group in optimizer.param_groups:
        param_group["lr"] = lr

    return lr
# ----------------------------------------------------------------------------- Run training for the specified number of epochs # Gradient scaler # Begin training # Save checkpoint # Save checkpoint config # Put on GPU # Adjust for case where n_frames_per_step > 1 # Float16 # Float32 # -----------------------------------------------------------------------------