Amphion/models/vc/Noro/noro_trainer.py

# Copyright (c) 2023 Amphion.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import os
import shutil
import time
import json5
import torch
import numpy as np
from tqdm import tqdm
from utils.util import ValueWindow
from torch.utils.data import DataLoader
from models.vc.Noro.noro_base_trainer import Noro_base_Trainer
from torch.nn import functional as F
from models.base.base_sampler import VariableSampler

from diffusers import get_scheduler
import accelerate
from accelerate.logging import get_logger
from accelerate.utils import ProjectConfiguration
from models.vc.Noro.noro_model import Noro_VCmodel
from models.vc.Noro.noro_dataset import VCCollator, VCDataset, batch_by_size
from processors.content_extractor import HubertExtractor
from models.vc.Noro.noro_loss import diff_loss, ConstractiveSpeakerLoss
from utils.mel import mel_spectrogram_torch
from utils.f0 import get_f0_features_using_dio, interpolate
from torch.nn.utils.rnn import pad_sequence


class NoroTrainer(Noro_base_Trainer):
    def __init__(self, args, cfg):
        self.args = args
        self.cfg = cfg
        cfg.exp_name = args.exp_name
        self.content_extractor = "mhubert"

        # Initialize accelerator and ensure all processes are ready
        self._init_accelerator()
        self.accelerator.wait_for_everyone()

        # Initialize logger on the main process
        if self.accelerator.is_main_process:
            self.logger = get_logger(args.exp_name, log_level="INFO")

        # Configure noise and speaker usage
        self.use_ref_noise = self.cfg.trans_exp.use_ref_noise

        # Log configuration on the main process
        if self.accelerator.is_main_process:
            self.logger.info(f"use_ref_noise: {self.use_ref_noise}")

        # Initialize a time window for monitoring metrics
        self.time_window = ValueWindow(50)

        # Log the start of training
        if self.accelerator.is_main_process:
            self.logger.info("=" * 56)
            self.logger.info("||\t\tNew training process started.\t\t||")
            self.logger.info("=" * 56)
            self.logger.info("\n")
            self.logger.debug(f"Using {args.log_level.upper()} logging level.")
            self.logger.info(f"Experiment name: {args.exp_name}")
            self.logger.info(f"Experiment directory: {self.exp_dir}")

        # Initialize checkpoint directory
        self.checkpoint_dir = os.path.join(self.exp_dir, "checkpoint")
        if self.accelerator.is_main_process:
            os.makedirs(self.checkpoint_dir, exist_ok=True)
            self.logger.debug(f"Checkpoint directory: {self.checkpoint_dir}")

        # Initialize training counters
        self.batch_count: int = 0
        self.step: int = 0
        self.epoch: int = 0
        self.max_epoch = (
            self.cfg.train.max_epoch if self.cfg.train.max_epoch > 0 else float("inf")
        )
        if self.accelerator.is_main_process:
            self.logger.info(
                f"Max epoch: {self.max_epoch if self.max_epoch < float('inf') else 'Unlimited'}"
            )

        # Check basic configuration
        if self.accelerator.is_main_process:
            self._check_basic_configs()
            self.save_checkpoint_stride = self.cfg.train.save_checkpoint_stride
            self.keep_last = [
                i if i > 0 else float("inf") for i in self.cfg.train.keep_last
            ]
            self.run_eval = self.cfg.train.run_eval

        # Set random seed
        with self.accelerator.main_process_first():
            self._set_random_seed(self.cfg.train.random_seed)

        # Setup data loader
        with self.accelerator.main_process_first():
            if self.accelerator.is_main_process:
                self.logger.info("Building dataset...")
            self.train_dataloader = self._build_dataloader()
            self.speaker_num = len(self.train_dataloader.dataset.speaker2id)
            if self.accelerator.is_main_process:
                self.logger.info("Speaker num: {}".format(self.speaker_num))

        # Build model
        with self.accelerator.main_process_first():
            if self.accelerator.is_main_process:
                self.logger.info("Building model...")
            self.model, self.w2v = self._build_model()

        # Resume training if specified
        with self.accelerator.main_process_first():
            if self.accelerator.is_main_process:
                self.logger.info("Resume training: {}".format(args.resume))
            if args.resume:
                if self.accelerator.is_main_process:
                    self.logger.info("Resuming from checkpoint...")
                ckpt_path = self._load_model(
                    self.checkpoint_dir,
                    args.checkpoint_path,
                    resume_type=args.resume_type,
                )
            self.checkpoint_dir = os.path.join(self.exp_dir, "checkpoint")
            if self.accelerator.is_main_process:
                os.makedirs(self.checkpoint_dir, exist_ok=True)
            if self.accelerator.is_main_process:
                self.logger.debug(f"Checkpoint directory: {self.checkpoint_dir}")

        # Initialize optimizer & scheduler
        with self.accelerator.main_process_first():
            if self.accelerator.is_main_process:
                self.logger.info("Building optimizer and scheduler...")
            self.optimizer = self._build_optimizer()
            self.scheduler = self._build_scheduler()

        # Prepare model, w2v, optimizer, and scheduler for accelerator
        self.model = self._prepare_for_accelerator(self.model)
        self.w2v = self._prepare_for_accelerator(self.w2v)
        self.optimizer = self._prepare_for_accelerator(self.optimizer)
        self.scheduler = self._prepare_for_accelerator(self.scheduler)

        # Build criterion
        with self.accelerator.main_process_first():
            if self.accelerator.is_main_process:
                self.logger.info("Building criterion...")
            self.criterion = self._build_criterion()

        self.config_save_path = os.path.join(self.exp_dir, "args.json")
        self.task_type = "VC"
        self.contrastive_speaker_loss = ConstractiveSpeakerLoss()

        if self.accelerator.is_main_process:
            self.logger.info("Task type: {}".format(self.task_type))

    def _init_accelerator(self):
        self.exp_dir = os.path.join(
            os.path.abspath(self.cfg.log_dir), self.args.exp_name
        )
        project_config = ProjectConfiguration(
            project_dir=self.exp_dir,
            logging_dir=os.path.join(self.exp_dir, "log"),
        )
        self.accelerator = accelerate.Accelerator(
            gradient_accumulation_steps=self.cfg.train.gradient_accumulation_step,
            log_with=self.cfg.train.tracker,
            project_config=project_config,
        )
        if self.accelerator.is_main_process:
            os.makedirs(project_config.project_dir, exist_ok=True)
            os.makedirs(project_config.logging_dir, exist_ok=True)
        self.accelerator.wait_for_everyone()
        with self.accelerator.main_process_first():
            self.accelerator.init_trackers(self.args.exp_name)

    def _build_model(self):
        w2v = HubertExtractor(self.cfg)
        model = Noro_VCmodel(cfg=self.cfg.model, use_ref_noise=self.use_ref_noise)
        return model, w2v

    def _build_dataloader(self):
        np.random.seed(int(time.time()))
        if self.accelerator.is_main_process:
            self.logger.info("Use Dynamic Batchsize...")
        train_dataset = VCDataset(self.cfg.trans_exp)
        train_collate = VCCollator(self.cfg)
        batch_sampler = batch_by_size(
            train_dataset.num_frame_indices,
            train_dataset.get_num_frames,
            max_tokens=self.cfg.train.max_tokens * self.accelerator.num_processes,
            max_sentences=self.cfg.train.max_sentences * self.accelerator.num_processes,
            required_batch_size_multiple=self.accelerator.num_processes,
        )
        np.random.shuffle(batch_sampler)
        batches = [
            x[self.accelerator.local_process_index :: self.accelerator.num_processes]
            for x in batch_sampler
            if len(x) % self.accelerator.num_processes == 0
        ]
        train_loader = DataLoader(
            train_dataset,
            collate_fn=train_collate,
            num_workers=self.cfg.train.dataloader.num_worker,
            batch_sampler=VariableSampler(
                batches, drop_last=False, use_random_sampler=True
            ),
            pin_memory=self.cfg.train.dataloader.pin_memory,
        )
        self.accelerator.wait_for_everyone()
        return train_loader

    def _build_optimizer(self):
        optimizer = torch.optim.AdamW(
            filter(lambda p: p.requires_grad, self.model.parameters()),
            **self.cfg.train.adam,
        )
        return optimizer

    def _build_scheduler(self):
        lr_scheduler = get_scheduler(
            self.cfg.train.lr_scheduler,
            optimizer=self.optimizer,
            num_warmup_steps=self.cfg.train.lr_warmup_steps,
            num_training_steps=self.cfg.train.num_train_steps,
        )
        return lr_scheduler

    def _build_criterion(self):
        criterion = torch.nn.L1Loss(reduction="mean")
        return criterion

    def _dump_cfg(self, path):
        os.makedirs(os.path.dirname(path), exist_ok=True)
        json5.dump(
            self.cfg,
            open(path, "w"),
            indent=4,
            sort_keys=True,
            ensure_ascii=False,
            quote_keys=True,
        )

    def load_model(self, checkpoint):
        self.step = checkpoint["step"]
        self.epoch = checkpoint["epoch"]
        self.model.load_state_dict(checkpoint["model"])
        self.optimizer.load_state_dict(checkpoint["optimizer"])
        self.scheduler.load_state_dict(checkpoint["scheduler"])

    def _prepare_for_accelerator(self, component):
        if isinstance(component, dict):
            for key in component.keys():
                component[key] = self.accelerator.prepare(component[key])
        else:
            component = self.accelerator.prepare(component)
        return component

    def _train_step(self, batch):
        total_loss = 0.0
        train_losses = {}
        device = self.accelerator.device

        # Move all Tensor data to the specified device
        batch = {
            k: v.to(device) if isinstance(v, torch.Tensor) else v
            for k, v in batch.items()
        }

        speech = batch["speech"]
        ref_speech = batch["ref_speech"]

        with torch.set_grad_enabled(False):
            # Extract features and spectrograms
            mel = mel_spectrogram_torch(speech, self.cfg).transpose(1, 2)
            ref_mel = mel_spectrogram_torch(ref_speech, self.cfg).transpose(1, 2)
            mask = batch["mask"]
            ref_mask = batch["ref_mask"]

            # Extract pitch and content features
            audio = speech.cpu().numpy()
            f0s = []
            for i in range(audio.shape[0]):
                wav = audio[i]
                f0 = get_f0_features_using_dio(wav, self.cfg.preprocess)
                f0, _ = interpolate(f0)
                frame_num = len(wav) // self.cfg.preprocess.hop_size
                f0 = torch.from_numpy(f0[:frame_num]).to(speech.device)
                f0s.append(f0)

            pitch = pad_sequence(f0s, batch_first=True, padding_value=0).float()
            pitch = (pitch - pitch.mean(dim=1, keepdim=True)) / (
                pitch.std(dim=1, keepdim=True) + 1e-6
            )  # Normalize pitch (B,T)
            _, content_feature = self.w2v.extract_content_features(
                speech
            )  # semantic (B, T, 768)

            if self.use_ref_noise:
                noisy_ref_mel = mel_spectrogram_torch(
                    batch["noisy_ref_speech"], self.cfg
                ).transpose(1, 2)

        if self.use_ref_noise:
            diff_out, (ref_emb, noisy_ref_emb), (cond_emb, _) = self.model(
                x=mel,
                content_feature=content_feature,
                pitch=pitch,
                x_ref=ref_mel,
                x_mask=mask,
                x_ref_mask=ref_mask,
                noisy_x_ref=noisy_ref_mel,
            )
        else:
            diff_out, (ref_emb, _), (cond_emb, _) = self.model(
                x=mel,
                content_feature=content_feature,
                pitch=pitch,
                x_ref=ref_mel,
                x_mask=mask,
                x_ref_mask=ref_mask,
            )

        if self.use_ref_noise:
            # B x N_query x D
            ref_emb = torch.mean(ref_emb, dim=1)  # B x D
            noisy_ref_emb = torch.mean(noisy_ref_emb, dim=1)  # B x D
            all_ref_emb = torch.cat([ref_emb, noisy_ref_emb], dim=0)  # 2B x D
            all_speaker_ids = torch.cat(
                [batch["speaker_id"], batch["speaker_id"]], dim=0
            )  # 2B
            cs_loss = self.contrastive_speaker_loss(all_ref_emb, all_speaker_ids) * 0.25
            total_loss += cs_loss
            train_losses["ref_loss"] = cs_loss

        diff_loss_x0 = diff_loss(diff_out["x0_pred"], mel, mask=mask)
        total_loss += diff_loss_x0
        train_losses["diff_loss_x0"] = diff_loss_x0

        diff_loss_noise = diff_loss(
            diff_out["noise_pred"], diff_out["noise"], mask=mask
        )
        total_loss += diff_loss_noise
        train_losses["diff_loss_noise"] = diff_loss_noise
        train_losses["total_loss"] = total_loss

        self.optimizer.zero_grad()
        self.accelerator.backward(total_loss)
        if self.accelerator.sync_gradients:
            self.accelerator.clip_grad_norm_(
                filter(lambda p: p.requires_grad, self.model.parameters()), 0.5
            )
        self.optimizer.step()
        self.scheduler.step()

        for item in train_losses:
            train_losses[item] = train_losses[item].item()

        train_losses["learning_rate"] = f"{self.optimizer.param_groups[0]['lr']:.1e}"
        train_losses["batch_size"] = batch["speaker_id"].shape[0]

        return (train_losses["total_loss"], train_losses, None)

    def _train_epoch(self):
        r"""Training epoch. Should return average loss of a batch (sample) over

        one epoch. See ``train_loop`` for usage.

        """
        if isinstance(self.model, dict):
            for key in self.model.keys():
                self.model[key].train()
        else:
            self.model.train()
        if isinstance(self.w2v, dict):
            for key in self.w2v.keys():
                self.w2v[key].eval()
        else:
            self.w2v.eval()

        epoch_sum_loss: float = 0.0  # total loss
        # Put the data to cuda device
        device = self.accelerator.device
        with device:
            torch.cuda.empty_cache()
        self.model = self.model.to(device)
        self.w2v = self.w2v.to(device)

        for batch in tqdm(
            self.train_dataloader,
            desc=f"Training Epoch {self.epoch}",
            unit="batch",
            colour="GREEN",
            leave=False,
            dynamic_ncols=True,
            smoothing=0.04,
            disable=not self.accelerator.is_main_process,
        ):
            speech = batch["speech"].cpu().numpy()
            speech = speech[0]
            self.batch_count += 1
            self.step += 1
            if len(speech) >= 16000 * 25:
                continue
            with self.accelerator.accumulate(self.model):
                total_loss, train_losses, _ = self._train_step(batch)

            if self.batch_count % self.cfg.train.gradient_accumulation_step == 0:
                epoch_sum_loss += total_loss
                self.current_loss = total_loss
                if isinstance(train_losses, dict):
                    for key, loss in train_losses.items():
                        self.accelerator.log(
                            {"Epoch/Train {} Loss".format(key): loss},
                            step=self.step,
                        )
                if self.accelerator.is_main_process and self.batch_count % 10 == 0:
                    self.echo_log(train_losses, mode="Training")

                self.save_checkpoint()
        self.accelerator.wait_for_everyone()

        return epoch_sum_loss, None

    def train_loop(self):
        r"""Training loop. The public entry of training process."""
        # Wait everyone to prepare before we move on
        self.accelerator.wait_for_everyone()
        # Dump config file
        if self.accelerator.is_main_process:
            self._dump_cfg(self.config_save_path)

        # Wait to ensure good to go
        self.accelerator.wait_for_everyone()
        # Stop when meeting max epoch or self.cfg.train.num_train_steps
        while (
            self.epoch < self.max_epoch and self.step < self.cfg.train.num_train_steps
        ):
            if self.accelerator.is_main_process:
                self.logger.info("\n")
                self.logger.info("-" * 32)
                self.logger.info("Epoch {}: ".format(self.epoch))
                self.logger.info("Start training...")

            train_total_loss, _ = self._train_epoch()

            self.epoch += 1
            self.accelerator.wait_for_everyone()
            if isinstance(self.scheduler, dict):
                for key in self.scheduler.keys():
                    self.scheduler[key].step()
            else:
                self.scheduler.step()

        # Finish training and save final checkpoint
        self.accelerator.wait_for_everyone()
        if self.accelerator.is_main_process:
            self.accelerator.save_state(
                os.path.join(
                    self.checkpoint_dir,
                    "final_epoch-{:04d}_step-{:07d}_loss-{:.6f}".format(
                        self.epoch, self.step, train_total_loss
                    ),
                )
            )
        self.accelerator.end_training()
        if self.accelerator.is_main_process:
            self.logger.info("Training finished...")

    def save_checkpoint(self):
        self.accelerator.wait_for_everyone()
        # Main process only
        if self.accelerator.is_main_process:
            if self.batch_count % self.save_checkpoint_stride[0] == 0:
                keep_last = self.keep_last[0]
                # Read all folders in self.checkpoint_dir
                all_ckpts = os.listdir(self.checkpoint_dir)
                # Exclude non-folders
                all_ckpts = [
                    ckpt
                    for ckpt in all_ckpts
                    if os.path.isdir(os.path.join(self.checkpoint_dir, ckpt))
                ]
                if len(all_ckpts) > keep_last:
                    # Keep only the last keep_last folders in self.checkpoint_dir, sorted by step "epoch-{:04d}_step-{:07d}_loss-{:.6f}"
                    all_ckpts = sorted(
                        all_ckpts, key=lambda x: int(x.split("_")[1].split("-")[1])
                    )
                    for ckpt in all_ckpts[:-keep_last]:
                        shutil.rmtree(os.path.join(self.checkpoint_dir, ckpt))
                checkpoint_filename = "epoch-{:04d}_step-{:07d}_loss-{:.6f}".format(
                    self.epoch, self.step, self.current_loss
                )
                path = os.path.join(self.checkpoint_dir, checkpoint_filename)
                self.logger.info("Saving state to {}...".format(path))
                self.accelerator.save_state(path)
                self.logger.info("Finished saving state.")
        self.accelerator.wait_for_everyone()

    def echo_log(self, losses, mode="Training"):
        message = [
            "{} - Epoch {} Step {}: [{:.3f} s/step]".format(
                mode, self.epoch + 1, self.step, self.time_window.average
            )
        ]

        for key in sorted(losses.keys()):
            if isinstance(losses[key], dict):
                for k, v in losses[key].items():
                    message.append(
                        str(k).split("/")[-1] + "=" + str(round(float(v), 5))
                    )
            else:
                message.append(
                    str(key).split("/")[-1] + "=" + str(round(float(losses[key]), 5))
                )
        self.logger.info(", ".join(message))