ParallelWaveGAN/parallel_wavegan/bin/train.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Copyright 2019 Tomoki Hayashi
#  MIT License (https://opensource.org/licenses/MIT)

"""Train Parallel WaveGAN."""

import argparse
import logging
import os
import sys
from collections import defaultdict

import matplotlib
import numpy as np
import soundfile as sf
import torch
import yaml
from tensorboardX import SummaryWriter
from torch.utils.data import DataLoader
from tqdm import tqdm

import parallel_wavegan
import parallel_wavegan.models
import parallel_wavegan.optimizers
from parallel_wavegan.datasets import (
    AudioDataset,
    AudioMelDataset,
    AudioMelF0ExcitationDataset,
    AudioMelSCPDataset,
    AudioSCPDataset,
)
from parallel_wavegan.layers import PQMF
# from parallel_wavegan.layers import ISTFT
from parallel_wavegan.losses import (
    DiscriminatorAdversarialLoss,
    DurationPredictorLoss,
    FeatureMatchLoss,
    GeneratorAdversarialLoss,
    MelSpectrogramLoss,
    MultiResolutionSTFTLoss,
)
from parallel_wavegan.utils import read_hdf5

# set to avoid matplotlib error in CLI environment
matplotlib.use("Agg")


class Trainer(object):
    """Customized trainer module for Parallel WaveGAN training."""

    def __init__(
        self,
        steps,
        epochs,
        data_loader,
        sampler,
        model,
        criterion,
        optimizer,
        scheduler,
        config,
        device=torch.device("cpu"),
    ):
        """Initialize trainer.

        Args:
            steps (int): Initial global steps.
            epochs (int): Initial global epochs.
            data_loader (dict): Dict of data loaders. It must contrain "train" and "dev" loaders.
            model (dict): Dict of models. It must contrain "generator" and "discriminator" models.
            criterion (dict): Dict of criterions. It must contrain "stft" and "mse" criterions.
            optimizer (dict): Dict of optimizers. It must contrain "generator" and "discriminator" optimizers.
            scheduler (dict): Dict of schedulers. It must contrain "generator" and "discriminator" schedulers.
            config (dict): Config dict loaded from yaml format configuration file.
            device (torch.deive): Pytorch device instance.

        """
        self.steps = steps
        self.epochs = epochs
        self.data_loader = data_loader
        self.sampler = sampler
        self.model = model
        self.criterion = criterion
        self.optimizer = optimizer
        self.scheduler = scheduler
        self.config = config
        self.device = device
        self.writer = SummaryWriter(config["outdir"])
        self.finish_train = False
        self.total_train_loss = defaultdict(float)
        self.total_eval_loss = defaultdict(float)
        self.is_vq = "VQVAE" in config.get("generator_type", "ParallelWaveGANGenerator")
        self.use_duration_prediction = "Duration" in config.get(
            "generator_type", "ParallelWaveGANGenerator"
        )

    def run(self):
        """Run training."""
        self.tqdm = tqdm(
            initial=self.steps, total=self.config["train_max_steps"], desc="[train]"
        )
        while True:
            # train one epoch
            self._train_epoch()

            # check whether training is finished
            if self.finish_train:
                break

        self.tqdm.close()
        logging.info("Finished training.")

    def save_checkpoint(self, checkpoint_path):
        """Save checkpoint.

        Args:
            checkpoint_path (str): Checkpoint path to be saved.

        """
        state_dict = {
            "optimizer": {
                "generator": self.optimizer["generator"].state_dict(),
                "discriminator": self.optimizer["discriminator"].state_dict(),
            },
            "scheduler": {
                "generator": self.scheduler["generator"].state_dict(),
                "discriminator": self.scheduler["discriminator"].state_dict(),
            },
            "steps": self.steps,
            "epochs": self.epochs,
        }
        if self.config["distributed"]:
            state_dict["model"] = {
                "generator": self.model["generator"].module.state_dict(),
                "discriminator": self.model["discriminator"].module.state_dict(),
            }
        else:
            state_dict["model"] = {
                "generator": self.model["generator"].state_dict(),
                "discriminator": self.model["discriminator"].state_dict(),
            }

        if not os.path.exists(os.path.dirname(checkpoint_path)):
            os.makedirs(os.path.dirname(checkpoint_path))
        torch.save(state_dict, checkpoint_path)

    def load_checkpoint(self, checkpoint_path, load_only_params=False):
        """Load checkpoint.

        Args:
            checkpoint_path (str): Checkpoint path to be loaded.
            load_only_params (bool): Whether to load only model parameters.

        """
        state_dict = torch.load(checkpoint_path, map_location="cpu")
        if self.config["distributed"]:
            self.model["generator"].module.load_state_dict(
                state_dict["model"]["generator"],
            )
            self.model["discriminator"].module.load_state_dict(
                state_dict["model"]["discriminator"],
                strict=False,
            )
        else:
            self.model["generator"].load_state_dict(
                state_dict["model"]["generator"],
            )
            self.model["discriminator"].load_state_dict(
                state_dict["model"]["discriminator"],
                strict=False,
            )
        if not load_only_params:
            self.steps = state_dict["steps"]
            self.epochs = state_dict["epochs"]
            self.optimizer["generator"].load_state_dict(
                state_dict["optimizer"]["generator"]
            )
            self.optimizer["discriminator"].load_state_dict(
                state_dict["optimizer"]["discriminator"]
            )
            self.scheduler["generator"].load_state_dict(
                state_dict["scheduler"]["generator"]
            )
            self.scheduler["discriminator"].load_state_dict(
                state_dict["scheduler"]["discriminator"]
            )

    def _train_step(self, batch):
        """Train model one step."""
        # parse batch and send to device
        if self.use_duration_prediction:
            x, y, ds = self._parse_batch(batch)
        else:
            x, y = self._parse_batch(batch)

        #######################
        #      Generator      #
        #######################
        if self.steps > self.config.get("generator_train_start_steps", 0):
            # initialize
            gen_loss = 0.0

            if self.is_vq:
                # vq case
                if self.config["generator_params"]["in_channels"] == 1:
                    y_, z_e, z_q = self.model["generator"](y, *x)
                else:
                    y_mb = self.criterion["pqmf"].analysis(y)
                    y_, z_e, z_q = self.model["generator"](y_mb, *x)
                quantize_loss = self.criterion["mse"](z_q, z_e.detach())
                commit_loss = self.criterion["mse"](z_e, z_q.detach())
                self.total_train_loss["train/quantization_loss"] += quantize_loss.item()
                self.total_train_loss["train/commitment_loss"] += commit_loss.item()
                gen_loss += quantize_loss + self.config["lambda_commit"] * commit_loss
            elif self.use_duration_prediction:
                assert ds is not None
                y_, ds_ = self.model["generator"](x, ds)
                duration_loss = self.criterion["duration"](ds_, ds)
                self.total_train_loss["train/duration_loss"] += duration_loss.item()
                gen_loss += duration_loss
            else:
                y_ = self.model["generator"](*x)
                # print('generator output y_: (out_channels=4)', y_.shape)
            # reconstruct the signal from multi-band signal
            if self.config["generator_params"]["out_channels"] > 1:
                y_mb_ = y_
                y_ = self.criterion["pqmf"].synthesis(y_mb_)
                # print('y_mb_ (y_) before pqmf, and y_ after pqmf:', y_mb_.shape, y_.shape)
            
            # multi-resolution sfft loss
            if self.config["use_stft_loss"]:
                # print("use_stft_loss: y_.shape, y.shape:", y_.shape, y.shape)
                sc_loss, mag_loss = self.criterion["stft"](y_, y)
                gen_loss += sc_loss + mag_loss
                self.total_train_loss[
                    "train/spectral_convergence_loss"
                ] += sc_loss.item()
                self.total_train_loss[
                    "train/log_stft_magnitude_loss"
                ] += mag_loss.item()

            # subband multi-resolution stft loss
            if self.config["use_subband_stft_loss"]:
                
                gen_loss *= 0.5  # for balancing with subband stft loss
                if not self.is_vq:
                    y_mb = self.criterion["pqmf"].analysis(y)
                # print("use_subband_stft_loss: y_mb_.shape, y_mb.shape:", y_mb_.shape, y_mb.shape)
                sub_sc_loss, sub_mag_loss = self.criterion["sub_stft"](y_mb_, y_mb)
                gen_loss += 0.5 * (sub_sc_loss + sub_mag_loss)
                self.total_train_loss[
                    "train/sub_spectral_convergence_loss"
                ] += sub_sc_loss.item()
                self.total_train_loss[
                    "train/sub_log_stft_magnitude_loss"
                ] += sub_mag_loss.item()

            # mel spectrogram loss
            if self.config["use_mel_loss"]:
                mel_loss = self.criterion["mel"](y_, y)
                gen_loss += mel_loss
                self.total_train_loss["train/mel_loss"] += mel_loss.item()

            # weighting aux loss
            gen_loss *= self.config.get("lambda_aux", 1.0)

            # adversarial loss
            if self.steps > self.config["discriminator_train_start_steps"]:
                if self.config["discriminator_type"] == "MultibandParallelWaveGANDiscriminator":
                    p_, p_low_, p_mid_, p_high_ = self.model["discriminator"](y_)
                    adv_loss = self.criterion["gen_adv"](p_)
                    adv_loss += self.criterion["gen_adv"](p_low_)
                    adv_loss += self.criterion["gen_adv"](p_mid_)
                    adv_loss += self.criterion["gen_adv"](p_high_)

                elif self.config["discriminator_type"] == "LowbandParallelWaveGANDiscriminator":
                    p_, p_low_ = self.model["discriminator"](y_)
                    adv_loss = self.criterion["gen_adv"](p_)
                    adv_loss_low = self.criterion["gen_adv"](p_low_)
                    adv_loss += self.criterion["gen_adv"](p_low_)
                    self.total_train_loss["train/adversarial_loss_low"] += adv_loss_low.item()                     
                else:    
                    p_ = self.model["discriminator"](y_)
                    adv_loss = self.criterion["gen_adv"](p_)

                self.total_train_loss["train/adversarial_loss"] += adv_loss.item()

                # feature matching loss
                if self.config["use_feat_match_loss"]:
                    # no need to track gradients
                    with torch.no_grad():
                        p_, p_low_ = self.model["discriminator"](y_)
                        p, p_low = self.model["discriminator"](y)
                    fm_loss = self.criterion["feat_match"](p_, p)
                    fm_loss_low = self.criterion["feat_match"](p_low_, p_low)
                    self.total_train_loss[
                        "train/feature_matching_loss"
                    ] += fm_loss
                    self.total_train_loss[
                        "train/feature_matching_loss_low"
                    ] += fm_loss_low
                    adv_loss += self.config["lambda_feat_match"] * fm_loss + self.config["lambda_feat_match_low"] * fm_loss_low

                # add adversarial loss to generator loss
                gen_loss += self.config["lambda_adv"] * adv_loss

                # istft
                
            self.total_train_loss["train/generator_loss"] += gen_loss.item()

            # update generator
            self.optimizer["generator"].zero_grad()
            gen_loss.backward()
            if self.config["generator_grad_norm"] > 0:
                torch.nn.utils.clip_grad_norm_(
                    self.model["generator"].parameters(),
                    self.config["generator_grad_norm"],
                )
            self.optimizer["generator"].step()
            self.scheduler["generator"].step()

        #######################
        #    Discriminator    #
        #######################
        if self.steps > self.config["discriminator_train_start_steps"]:
            if self.config.get("update_prediction_after_generator_update", True):
                # re-compute y_ which leads better quality
                with torch.no_grad():
                    if self.is_vq:
                        if self.config["generator_params"]["in_channels"] == 1:
                            y_, _, _ = self.model["generator"](y, *x)
                        else:
                            y_, _, _ = self.model["generator"](y_mb, *x)
                    elif self.use_duration_prediction:
                        assert ds is not None
                        y_, _ = self.model["generator"](x, ds)
                    else:
                        y_ = self.model["generator"](*x)
                if self.config["generator_params"]["out_channels"] > 1:
                    y_ = self.criterion["pqmf"].synthesis(y_)

            # discriminator loss
            if self.config["discriminator_type"] == "MultibandParallelWaveGANDiscriminator":
                p_, p_low_, p_mid_, p_high_ = self.model["discriminator"](y_.detach())
                p, p_low, p_mid, p_high = self.model["discriminator"](y)
                real_loss, fake_loss = self.criterion["dis_adv"](p_, p)
                real_loss_low, fake_loss_low = self.criterion["dis_adv"](p_low_, p_low)
                real_loss_mid, fake_loss_mid = self.criterion["dis_adv"](p_mid_, p_mid)
                real_loss_high, fake_loss_high = self.criterion["dis_adv"](p_high_, p_high)
                real_loss += real_loss_low + real_loss_mid + real_loss_high
                fake_loss += fake_loss_low + fake_loss_mid + fake_loss_high
            elif self.config["discriminator_type"] == "LowbandParallelWaveGANDiscriminator":
                p_, p_low_ = self.model["discriminator"](y_.detach())
                p, p_low = self.model["discriminator"](y)
                real_loss, fake_loss = self.criterion["dis_adv"](p_, p)
                real_loss_low, fake_loss_low = self.criterion["dis_adv"](p_low_, p_low)
                real_loss += real_loss_low  
                fake_loss += fake_loss_low  
                self.total_train_loss["train/real_loss_low"] += real_loss_low.item()
                self.total_train_loss["train/fake_loss_low"] += fake_loss_low.item()
            else:
                p = self.model["discriminator"](y)
                p_ = self.model["discriminator"](y_.detach())
                real_loss, fake_loss = self.criterion["dis_adv"](p_, p)
            
            dis_loss = real_loss + fake_loss
            self.total_train_loss["train/real_loss"] += real_loss.item()
            self.total_train_loss["train/fake_loss"] += fake_loss.item()
            self.total_train_loss["train/discriminator_loss"] += dis_loss.item()

            # update discriminator
            self.optimizer["discriminator"].zero_grad()
            dis_loss.backward()
            if self.config["discriminator_grad_norm"] > 0:
                torch.nn.utils.clip_grad_norm_(
                    self.model["discriminator"].parameters(),
                    self.config["discriminator_grad_norm"],
                )
            self.optimizer["discriminator"].step()
            self.scheduler["discriminator"].step()

        # update counts
        self.steps += 1
        self.tqdm.update(1)
        self._check_train_finish()

    def _train_epoch(self):
        """Train model one epoch."""
        for train_steps_per_epoch, batch in enumerate(self.data_loader["train"], 1):
            # train one step
            self._train_step(batch)

            # check interval
            if self.config["rank"] == 0:
                self._check_log_interval()
                self._check_eval_interval()
                self._check_save_interval()

            # check whether training is finished
            if self.finish_train:
                return

        # update
        self.epochs += 1
        self.train_steps_per_epoch = train_steps_per_epoch
        logging.info(
            f"(Steps: {self.steps}) Finished {self.epochs} epoch training "
            f"({self.train_steps_per_epoch} steps per epoch)."
        )

        # needed for shuffle in distributed training
        if self.config["distributed"]:
            self.sampler["train"].set_epoch(self.epochs)

    @torch.no_grad()
    def _eval_step(self, batch):
        """Evaluate model one step."""
        # parse batch and send to device
        if self.use_duration_prediction:
            x, y, ds = self._parse_batch(batch)
        else:
            x, y = self._parse_batch(batch)

        #######################
        #      Generator      #
        #######################
        if self.is_vq:
            if self.config["generator_params"]["in_channels"] == 1:
                y_, z_e, z_q = self.model["generator"](y, *x)
            else:
                y_mb = self.criterion["pqmf"].analysis(y)
                y_, z_e, z_q = self.model["generator"](y_mb, *x)
            quantize_loss = self.criterion["mse"](z_q, z_e.detach())
            commit_loss = self.criterion["mse"](z_e, z_q.detach())
        elif self.use_duration_prediction:
            assert ds is not None
            y_, ds_ = self.model["generator"](x, ds)
            duration_loss = self.criterion["duration"](ds_, torch.log(ds))
        else:
            y_ = self.model["generator"](*x)
        if self.config["generator_params"]["out_channels"] > 1:
            y_mb_ = y_
            y_ = self.criterion["pqmf"].synthesis(y_mb_)

        # initialize
        aux_loss = 0.0

        # multi-resolution stft loss
        if self.config["use_stft_loss"]:
            sc_loss, mag_loss = self.criterion["stft"](y_, y)
            aux_loss += sc_loss + mag_loss
            self.total_eval_loss["eval/spectral_convergence_loss"] += sc_loss.item()
            self.total_eval_loss["eval/log_stft_magnitude_loss"] += mag_loss.item()

        # subband multi-resolution stft loss
        if self.config.get("use_subband_stft_loss", False):
            aux_loss *= 0.5  # for balancing with subband stft loss
            if not self.is_vq:
                y_mb = self.criterion["pqmf"].analysis(y)
            sub_sc_loss, sub_mag_loss = self.criterion["sub_stft"](y_mb_, y_mb)
            self.total_eval_loss[
                "eval/sub_spectral_convergence_loss"
            ] += sub_sc_loss.item()
            self.total_eval_loss[
                "eval/sub_log_stft_magnitude_loss"
            ] += sub_mag_loss.item()
            aux_loss += 0.5 * (sub_sc_loss + sub_mag_loss)

        # mel spectrogram loss
        if self.config["use_mel_loss"]:
            mel_loss = self.criterion["mel"](y_, y)
            aux_loss += mel_loss
            self.total_eval_loss["eval/mel_loss"] += mel_loss.item()

        # weighting stft loss
        aux_loss *= self.config.get("lambda_aux", 1.0)

        # adversarial loss
        p_ = self.model["discriminator"](y_)
        adv_loss = self.criterion["gen_adv"](p_)
        gen_loss = aux_loss + self.config["lambda_adv"] * adv_loss

        # feature matching loss
        if self.config["use_feat_match_loss"]:
            p = self.model["discriminator"](y)
            fm_loss = self.criterion["feat_match"](p_, p)
            self.total_eval_loss["eval/feature_matching_loss"] += fm_loss
            gen_loss += (
                self.config["lambda_adv"] * self.config["lambda_feat_match"] * fm_loss
            )

        #######################
        #    Discriminator    #
        #######################
        p = self.model["discriminator"](y)
        p_ = self.model["discriminator"](y_)

        # discriminator loss
        real_loss, fake_loss = self.criterion["dis_adv"](p_, p)
        dis_loss = real_loss + fake_loss

        # add to total eval loss
        self.total_eval_loss["eval/adversarial_loss"] += adv_loss.item()
        self.total_eval_loss["eval/generator_loss"] += gen_loss.item()
        self.total_eval_loss["eval/real_loss"] += real_loss.item()
        self.total_eval_loss["eval/fake_loss"] += fake_loss.item()
        self.total_eval_loss["eval/discriminator_loss"] += dis_loss.item()
        if self.is_vq:
            self.total_eval_loss["eval/quantization_loss"] += quantize_loss.item()
            self.total_eval_loss["eval/commitment_loss"] += commit_loss.item()
        if self.use_duration_prediction:
            self.total_eval_loss["eval/duration_loss"] += duration_loss.item()

    def _eval_epoch(self):
        """Evaluate model one epoch."""
        logging.info(f"(Steps: {self.steps}) Start evaluation.")
        # change mode
        for key in self.model.keys():
            self.model[key].eval()

        # calculate loss for each batch
        for eval_steps_per_epoch, batch in enumerate(
            tqdm(self.data_loader["dev"], desc="[eval]"), 1
        ):
            # eval one step
            self._eval_step(batch)

            # save intermediate result
            if eval_steps_per_epoch == 1:
                self._genearete_and_save_intermediate_result(batch)

        logging.info(
            f"(Steps: {self.steps}) Finished evaluation "
            f"({eval_steps_per_epoch} steps per epoch)."
        )

        # average loss
        for key in self.total_eval_loss.keys():
            self.total_eval_loss[key] /= eval_steps_per_epoch
            logging.info(
                f"(Steps: {self.steps}) {key} = {self.total_eval_loss[key]:.4f}."
            )

        # record
        self._write_to_tensorboard(self.total_eval_loss)

        # reset
        self.total_eval_loss = defaultdict(float)

        # restore mode
        for key in self.model.keys():
            self.model[key].train()

    @torch.no_grad()
    def _genearete_and_save_intermediate_result(self, batch):
        """Generate and save intermediate result."""
        # delayed import to avoid error related backend error
        import matplotlib.pyplot as plt

        # parse batch and send to device
        if self.use_duration_prediction:
            x_batch, y_batch, _ = self._parse_batch(batch)
        else:
            x_batch, y_batch = self._parse_batch(batch)

        # generate
        if self.is_vq:
            if self.config["generator_params"]["in_channels"] == 1:
                y_batch_, _, _ = self.model["generator"](y_batch, *x_batch)
            else:
                y_batch_, _, _ = self.model["generator"](
                    self.criterion["pqmf"].analysis(y_batch), *x_batch
                )
        elif self.use_duration_prediction:
            y_batch_, _ = self.model["generator"].synthesis(x_batch)
        else:
            y_batch_ = self.model["generator"](*x_batch)
        if self.config["generator_params"]["out_channels"] > 1:
            y_batch_ = self.criterion["pqmf"].synthesis(y_batch_)

        # check directory
        dirname = os.path.join(self.config["outdir"], f"predictions/{self.steps}steps")
        if not os.path.exists(dirname):
            os.makedirs(dirname)

        for idx, (y, y_) in enumerate(zip(y_batch, y_batch_), 1):
            # convert to ndarray
            y, y_ = y.view(-1).cpu().numpy(), y_.view(-1).cpu().numpy()

            # plot figure and save it
            figname = os.path.join(dirname, f"{idx}.png")
            plt.subplot(2, 1, 1)
            plt.plot(y)
            plt.title("groundtruth speech")
            plt.subplot(2, 1, 2)
            plt.plot(y_)
            plt.title(f"generated speech @ {self.steps} steps")
            plt.tight_layout()
            plt.savefig(figname)
            plt.close()

            # save as wavfile
            y = np.clip(y, -1, 1)
            y_ = np.clip(y_, -1, 1)
            sf.write(
                figname.replace(".png", "_ref.wav"),
                y,
                self.config["sampling_rate"],
                "PCM_16",
            )
            sf.write(
                figname.replace(".png", "_gen.wav"),
                y_,
                self.config["sampling_rate"],
                "PCM_16",
            )

            if idx >= self.config["num_save_intermediate_results"]:
                break

    def _parse_batch(self, batch):
        """Parse batch and send to the device."""
        # parse batch
        if self.use_duration_prediction:
            inputs, targets, durations = batch
        else:
            inputs, targets = batch

        # send inputs to device
        if isinstance(inputs, torch.Tensor):
            x = inputs.to(self.device)
        elif isinstance(inputs, (tuple, list)):
            x = [None if x is None else x.to(self.device) for x in inputs]
        else:
            raise ValueError(f"Not supported type ({type(inputs)}).")

        # send targets to device
        if isinstance(targets, torch.Tensor):
            y = targets.to(self.device)
        elif isinstance(targets, (tuple, list)):
            y = [None if y is None else y.to(self.device) for y in targets]
        else:
            raise ValueError(f"Not supported type ({type(targets)}).")

        if self.use_duration_prediction:
            # send durations to device (for model with duration prediction only)
            if isinstance(durations, torch.Tensor):
                ds = durations.to(self.device)
            elif isinstance(durations, (tuple, list)):
                ds = [None if d is None else d.to(self.device) for d in durations]
            else:
                raise ValueError(f"Not supported type ({type(durations)}).")

            return x, y, ds

        return x, y

    def _write_to_tensorboard(self, loss):
        """Write to tensorboard."""
        for key, value in loss.items():
            self.writer.add_scalar(key, value, self.steps)

    def _check_save_interval(self):
        if self.steps % self.config["save_interval_steps"] == 0:
            self.save_checkpoint(
                os.path.join(self.config["outdir"], f"checkpoint-{self.steps}steps.pkl")
            )
            logging.info(f"Successfully saved checkpoint @ {self.steps} steps.")

    def _check_eval_interval(self):
        if self.steps % self.config["eval_interval_steps"] == 0:
            self._eval_epoch()

    def _check_log_interval(self):
        if self.steps % self.config["log_interval_steps"] == 0:
            for key in self.total_train_loss.keys():
                self.total_train_loss[key] /= self.config["log_interval_steps"]
                logging.info(
                    f"(Steps: {self.steps}) {key} = {self.total_train_loss[key]:.4f}."
                )
            self._write_to_tensorboard(self.total_train_loss)

            # reset
            self.total_train_loss = defaultdict(float)

    def _check_train_finish(self):
        if self.steps >= self.config["train_max_steps"]:
            self.finish_train = True


class Collater(object):
    """Customized collater for Pytorch DataLoader in training."""

    def __init__(
        self,
        batch_max_steps=20480,
        hop_size=256,
        aux_context_window=2,
        use_noise_input=False,
        use_f0_and_excitation=False,
        use_aux_input=True,
        use_duration=False,
        use_global_condition=False,
        use_local_condition=False,
        pad_value=0,
    ):
        """Initialize customized collater for PyTorch DataLoader.

        Args:
            batch_max_steps (int): The maximum length of input signal in batch.
            hop_size (int): Hop size of auxiliary features.
            aux_context_window (int): Context window size for auxiliary feature conv.
            use_noise_input (bool): Whether to use noise input.
            use_f0_and_excitation (bool): Whether to use f0 and ext. input.
            use_aux_input (bool): Whether to use auxiliary input.
            use_duration (bool): Whether to use duration for duration prediction.
            use_global_condition (bool): Whether to use global conditioning.
            use_local_condition (bool): Whether to use local conditioning.

        """
        if hop_size is not None:
            if batch_max_steps % hop_size != 0:
                batch_max_steps += -(batch_max_steps % hop_size)
            assert batch_max_steps % hop_size == 0
            self.hop_size = hop_size
            self.batch_max_frames = batch_max_steps // hop_size
        self.batch_max_steps = batch_max_steps
        self.aux_context_window = aux_context_window
        self.use_noise_input = use_noise_input
        self.use_f0_and_excitation = use_f0_and_excitation
        self.use_aux_input = use_aux_input
        self.use_duration = use_duration
        self.use_global_condition = use_global_condition
        self.use_local_condition = use_local_condition
        self.pad_value = pad_value
        if not self.use_aux_input:
            assert not self.use_noise_input, "Not supported."
            assert not self.use_duration, "Not supported."
        if self.use_noise_input:
            assert not self.use_duration, "Not supported."
        if self.use_local_condition:
            assert not self.use_aux_input and not self.use_duration, "Not supported."
        if self.use_global_condition:
            assert not self.use_aux_input and not self.use_duration, "Not supported."

        # set useful values in random cutting
        if self.use_aux_input or self.use_local_condition:
            self.start_offset = aux_context_window
            self.end_offset = -(self.batch_max_frames + aux_context_window)
            self.mel_threshold = self.batch_max_frames + 2 * aux_context_window
        else:
            self.start_offset = 0
            self.end_offset = -self.batch_max_steps
            self.audio_threshold = self.batch_max_steps

    def __call__(self, batch):
        """Convert into batch tensors.

        Args:
            batch (list): list of tuple of the pair of audio and features.

        Returns:
            Tuple: Tuple of Gaussian noise batch (B, 1, T) and auxiliary feature
                batch (B, C, T'), where T = (T' - 2 * aux_context_window) * hop_size.
                If use_noise_input = False, Gaussian noise batch is not included.
                If use_aux_input = False, auxiliary feature batch is not included.
                If both use_noise_input and use_aux_input to False, this tuple is
                not returned.
            Tensor: Target signal batch (B, 1, T).

        """
        if self.use_aux_input:
            #################################
            #          MEL2WAV CASE         #
            #################################
            # check length
            batch = [
                self._adjust_length(*b) for b in batch if len(b[1]) > self.mel_threshold
            ]
            xs, cs = [b[0] for b in batch], [b[1] for b in batch]
            if self.use_f0_and_excitation:
                fs, es = [b[2] for b in batch], [b[3] for b in batch]

            # make batch with random cut
            c_lengths = [len(c) for c in cs]
            start_frames = np.array(
                [
                    np.random.randint(self.start_offset, cl + self.end_offset)
                    for cl in c_lengths
                ]
            )
            x_starts = start_frames * self.hop_size
            x_ends = x_starts + self.batch_max_steps
            c_starts = start_frames - self.aux_context_window
            c_ends = start_frames + self.batch_max_frames + self.aux_context_window
            y_batch = [x[start:end] for x, start, end in zip(xs, x_starts, x_ends)]
            c_batch = [c[start:end] for c, start, end in zip(cs, c_starts, c_ends)]

            # convert each batch to tensor, asuume that each item in batch has the same length
            y_batch, c_batch = np.array(y_batch), np.array(c_batch)
            y_batch = torch.tensor(y_batch, dtype=torch.float).unsqueeze(1)  # (B, 1, T)

            if self.use_f0_and_excitation:
                f_batch = [f[start:end] for f, start, end in zip(fs, c_starts, c_ends)]
                e_batch = [e[start:end] for e, start, end in zip(es, c_starts, c_ends)]
                f_batch, e_batch = np.array(f_batch), np.array(e_batch)
                f_batch = torch.tensor(f_batch, dtype=torch.float).unsqueeze(
                    1
                )  # (B, 1, T')
                e_batch = torch.tensor(e_batch, dtype=torch.float)  # (B, 1, T', C')
                e_batch = e_batch.reshape(e_batch.shape[0], 1, -1)  # (B, 1, T' * C')

            # duration calculation and return with duration information
            if self.use_duration:
                updated_c_batch, d_batch = [], []
                for c in c_batch:
                    # NOTE(jiatong): assume 0 is the discrete symbol
                    # (refer to cvss_c/local/preprocess_hubert.py)
                    code, d = torch.unique_consecutive(
                        torch.tensor(c, dtype=torch.long), return_counts=True, dim=0
                    )
                    updated_c_batch.append(code)
                    d_batch.append(d)
                c_batch = self._pad_list(updated_c_batch, self.pad_value).transpose(
                    2, 1
                )  # (B, C, T')
                d_batch = self._pad_list(d_batch, 0)
                return c_batch, y_batch, d_batch

            # process data without duration prediction
            c_batch = torch.tensor(c_batch, dtype=torch.float).transpose(
                2, 1
            )  # (B, C, T')

            input_items = (c_batch,)
            if self.use_noise_input:
                # make input noise signal batch tensor
                z_batch = torch.randn(y_batch.size())  # (B, 1, T)
                input_items = (z_batch,) + input_items
            if self.use_f0_and_excitation:
                input_items = input_items + (f_batch, e_batch)

            return input_items, y_batch
        else:
            #################################
            #        VQ-WAV2WAV CASE        #
            #################################
            if self.use_local_condition:
                # check length
                batch_idx = [
                    idx
                    for idx, b in enumerate(batch)
                    if len(b[1]) >= self.mel_threshold
                ]

                # fix length
                batch_ = [
                    self._adjust_length(batch[idx][0], batch[idx][1])
                    for idx in batch_idx
                ]

                # decide random index
                l_lengths = [len(b[1]) for b in batch_]
                l_starts = np.array(
                    [
                        np.random.randint(self.start_offset, ll + self.end_offset)
                        for ll in l_lengths
                    ]
                )
                l_ends = l_starts + self.batch_max_frames
                y_starts = l_starts * self.hop_size
                y_ends = y_starts + self.batch_max_steps

                # make random batch
                y_batch = [
                    b[0][start:end] for b, start, end in zip(batch_, y_starts, y_ends)
                ]
                l_batch = [
                    b[1][start:end] for b, start, end in zip(batch_, l_starts, l_ends)
                ]
                if self.use_global_condition:
                    g_batch = [batch[idx][2].reshape(1) for idx in batch_idx]
            else:
                # check length
                if self.use_global_condition:
                    batch = [b for b in batch if len(b[0]) >= self.audio_threshold]
                else:
                    batch = [(b,) for b in batch if len(b) >= self.audio_threshold]

                # decide random index
                y_lengths = [len(b[0]) for b in batch]
                y_starts = np.array(
                    [
                        np.random.randint(self.start_offset, yl + self.end_offset)
                        for yl in y_lengths
                    ]
                )
                y_ends = y_starts + self.batch_max_steps

                # make random batch
                y_batch = [
                    b[0][start:end] for b, start, end in zip(batch, y_starts, y_ends)
                ]
                if self.use_global_condition:
                    g_batch = [b[1].reshape(1) for b in batch]

            # convert each batch to tensor, asuume that each item in batch has the same length
            y_batch = torch.tensor(y_batch, dtype=torch.float).unsqueeze(1)  # (B, 1, T)
            if self.use_local_condition:
                l_batch = torch.tensor(l_batch, dtype=torch.float).transpose(
                    2, 1
                )  # (B, C' T')
            else:
                l_batch = None
            if self.use_global_condition:
                g_batch = torch.tensor(g_batch, dtype=torch.long).view(-1)  # (B,)
            else:
                g_batch = None

            # NOTE(kan-bayashi): Always return "l" and "g" since VQ-VAE can accept None
            return (l_batch, g_batch), y_batch

    def _adjust_length(self, x, c, f0=None, excitation=None):
        """Adjust the audio and feature lengths.

        Note:
            Basically we assume that the length of x and c are adjusted
            through preprocessing stage, but if we use other library processed
            features, this process will be needed.

        """
        if len(x) < len(c) * self.hop_size:
            x = np.pad(x, (0, len(c) * self.hop_size - len(x)), mode="edge")

        # check the legnth is valid
        assert len(x) == len(c) * self.hop_size

        if f0 is not None and excitation is not None:
            return x, c, f0, excitation
        else:
            return x, c

    def _pad_list(self, xs, pad_value):
        """Perform padding for the list of tensors.

        Args:
            xs (List): List of Tensors [(T_1, `*`), (T_2, `*`), ..., (T_B, `*`)].
            pad_value (float): Value for padding.

        Returns:
            Tensor: Padded tensor (B, Tmax, `*`).

        Examples:
            >>> x = [torch.ones(4), torch.ones(2), torch.ones(1)]
            >>> x
            [tensor([1., 1., 1., 1.]), tensor([1., 1.]), tensor([1.])]
            >>> pad_list(x, 0)
            tensor([[1., 1., 1., 1.],
                    [1., 1., 0., 0.],
                    [1., 0., 0., 0.]])

        """
        n_batch = len(xs)
        max_len = max(x.size(0) for x in xs)
        pad = xs[0].new(n_batch, max_len, *xs[0].size()[1:]).fill_(pad_value)

        for i in range(n_batch):
            pad[i, : xs[i].size(0)] = xs[i]

        return pad


def main():
    """Run training process."""
    parser = argparse.ArgumentParser(
        description=(
            "Train Parallel WaveGAN (See detail in parallel_wavegan/bin/train.py)."
        )
    )
    parser.add_argument(
        "--train-wav-scp",
        default=None,
        type=str,
        help=(
            "kaldi-style wav.scp file for training. "
            "you need to specify either train-*-scp or train-dumpdir."
        ),
    )
    parser.add_argument(
        "--train-feats-scp",
        default=None,
        type=str,
        help=(
            "kaldi-style feats.scp file for training. "
            "you need to specify either train-*-scp or train-dumpdir."
        ),
    )
    parser.add_argument(
        "--train-segments",
        default=None,
        type=str,
        help="kaldi-style segments file for training.",
    )
    parser.add_argument(
        "--train-dumpdir",
        default=None,
        type=str,
        help=(
            "directory including training data. "
            "you need to specify either train-*-scp or train-dumpdir."
        ),
    )
    parser.add_argument(
        "--dev-wav-scp",
        default=None,
        type=str,
        help=(
            "kaldi-style wav.scp file for validation. "
            "you need to specify either dev-*-scp or dev-dumpdir."
        ),
    )
    parser.add_argument(
        "--dev-feats-scp",
        default=None,
        type=str,
        help=(
            "kaldi-style feats.scp file for vaidation. "
            "you need to specify either dev-*-scp or dev-dumpdir."
        ),
    )
    parser.add_argument(
        "--dev-segments",
        default=None,
        type=str,
        help="kaldi-style segments file for validation.",
    )
    parser.add_argument(
        "--dev-dumpdir",
        default=None,
        type=str,
        help=(
            "directory including development data. "
            "you need to specify either dev-*-scp or dev-dumpdir."
        ),
    )
    parser.add_argument(
        "--outdir",
        type=str,
        required=True,
        help="directory to save checkpoints.",
    )
    parser.add_argument(
        "--config",
        type=str,
        required=True,
        help="yaml format configuration file.",
    )
    parser.add_argument(
        "--pretrain",
        default="",
        type=str,
        nargs="?",
        help='checkpoint file path to load pretrained params. (default="")',
    )
    parser.add_argument(
        "--resume",
        default="",
        type=str,
        nargs="?",
        help='checkpoint file path to resume training. (default="")',
    )
    parser.add_argument(
        "--verbose",
        type=int,
        default=1,
        help="logging level. higher is more logging. (default=1)",
    )
    parser.add_argument(
        "--rank",
        "--local_rank",
        default=0,
        type=int,
        help="rank for distributed training. no need to explictly specify.",
    )
    args = parser.parse_args()

    args.distributed = False
    if not torch.cuda.is_available():
        device = torch.device("cpu")
    else:
        device = torch.device("cuda")
        # effective when using fixed size inputs
        # see https://discuss.pytorch.org/t/what-does-torch-backends-cudnn-benchmark-do/5936
        torch.backends.cudnn.benchmark = True
        torch.cuda.set_device(args.rank)
        # setup for distributed training
        # see example: https://github.com/NVIDIA/apex/tree/master/examples/simple/distributed
        if "WORLD_SIZE" in os.environ:
            args.world_size = int(os.environ["WORLD_SIZE"])
            args.distributed = args.world_size > 1
        if args.distributed:
            torch.distributed.init_process_group(backend="nccl", init_method="env://")

    # suppress logging for distributed training
    if args.rank != 0:
        sys.stdout = open(os.devnull, "w")

    # set logger
    if args.verbose > 1:
        logging.basicConfig(
            level=logging.DEBUG,
            stream=sys.stdout,
            format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
        )
    elif args.verbose > 0:
        logging.basicConfig(
            level=logging.INFO,
            stream=sys.stdout,
            format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
        )
    else:
        logging.basicConfig(
            level=logging.WARN,
            stream=sys.stdout,
            format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
        )
        logging.warning("Skip DEBUG/INFO messages")

    # check directory existence
    if not os.path.exists(args.outdir):
        os.makedirs(args.outdir)

    # check arguments
    if (args.train_feats_scp is not None and args.train_dumpdir is not None) or (
        args.train_feats_scp is None and args.train_dumpdir is None
    ):
        raise ValueError("Please specify either --train-dumpdir or --train-*-scp.")
    if (args.dev_feats_scp is not None and args.dev_dumpdir is not None) or (
        args.dev_feats_scp is None and args.dev_dumpdir is None
    ):
        raise ValueError("Please specify either --dev-dumpdir or --dev-*-scp.")

    # load and save config
    with open(args.config) as f:
        config = yaml.load(f, Loader=yaml.Loader)
    config.update(vars(args))
    config["version"] = parallel_wavegan.__version__  # add version info
    with open(os.path.join(args.outdir, "config.yml"), "w") as f:
        yaml.dump(config, f, Dumper=yaml.Dumper)
    for key, value in config.items():
        logging.info(f"{key} = {value}")

    # get configuration
    generator_type = config.get("generator_type", "ParallelWaveGANGenerator")
    use_aux_input = "VQVAE" not in generator_type
    use_noise_input = (
        "ParallelWaveGAN" in generator_type and "VQVAE" not in generator_type
    )
    use_duration = "Duration" in generator_type
    use_local_condition = config.get("use_local_condition", False)
    use_global_condition = config.get("use_global_condition", False)
    use_f0_and_excitation = generator_type == "UHiFiGANGenerator"

    # setup query and load function
    if args.train_wav_scp is None or args.dev_wav_scp is None:
        local_query = None
        local_load_fn = None
        global_query = None
        global_load_fn = None
        if config["format"] == "hdf5":
            audio_query, mel_query = "*.h5", "*.h5"
            audio_load_fn = lambda x: read_hdf5(x, "wave")  # NOQA
            mel_load_fn = lambda x: read_hdf5(x, "feats")  # NOQA
            if use_f0_and_excitation:
                f0_query, excitation_query = "*.h5", "*.h5"
                f0_load_fn = lambda x: read_hdf5(x, "f0")  # NOQA
                excitation_load_fn = lambda x: read_hdf5(x, "excitation")  # NOQA
            if use_local_condition:
                local_query = "*.h5"
                local_load_fn = lambda x: read_hdf5(x, "local")  # NOQA
            if use_global_condition:
                global_query = "*.h5"
                global_load_fn = lambda x: read_hdf5(x, "global")  # NOQA
        elif config["format"] == "npy":
            audio_query, mel_query = "*-wave.npy", "*-feats.npy"
            audio_load_fn = np.load
            mel_load_fn = np.load
            if use_f0_and_excitation:
                f0_query, excitation_query = "*-f0.npy", "*-excitation.npy"
                f0_load_fn = np.load
                excitation_load_fn = np.load
            if use_local_condition:
                local_query = "*-local.npy"
                local_load_fn = np.load
            if use_global_condition:
                global_query = "*-global.npy"
                global_load_fn = np.load
        else:
            raise ValueError("support only hdf5 or npy format.")

    # setup length threshold
    if config["remove_short_samples"]:
        audio_length_threshold = config["batch_max_steps"]
        mel_length_threshold = config["batch_max_steps"] // config[
            "hop_size"
        ] + 2 * config["generator_params"].get("aux_context_window", 0)
    else:
        mel_length_threshold = None
        audio_length_threshold = None

    # define dataset for training data
    if args.train_dumpdir is not None:
        if not use_f0_and_excitation:
            if use_aux_input:
                train_dataset = AudioMelDataset(
                    root_dir=args.train_dumpdir,
                    audio_query=audio_query,
                    audio_load_fn=audio_load_fn,
                    mel_query=mel_query,
                    mel_load_fn=mel_load_fn,
                    local_query=local_query,
                    local_load_fn=local_load_fn,
                    global_query=global_query,
                    global_load_fn=global_load_fn,
                    mel_length_threshold=mel_length_threshold,
                    allow_cache=config.get("allow_cache", False),  # keep compatibility
                )
            else:
                train_dataset = AudioDataset(
                    root_dir=args.train_dumpdir,
                    audio_query=audio_query,
                    audio_load_fn=audio_load_fn,
                    local_query=local_query,
                    local_load_fn=local_load_fn,
                    global_query=global_query,
                    global_load_fn=global_load_fn,
                    audio_length_threshold=audio_length_threshold,
                    allow_cache=config.get("allow_cache", False),  # keep compatibility
                )
        else:
            train_dataset = AudioMelF0ExcitationDataset(
                root_dir=args.train_dumpdir,
                audio_query=audio_query,
                mel_query=mel_query,
                f0_query=f0_query,
                excitation_query=excitation_query,
                audio_load_fn=audio_load_fn,
                mel_load_fn=mel_load_fn,
                f0_load_fn=f0_load_fn,
                excitation_load_fn=excitation_load_fn,
                mel_length_threshold=mel_length_threshold,
                allow_cache=config.get("allow_cache", False),  # keep compatibility
            )
    else:
        if use_f0_and_excitation:
            raise NotImplementedError(
                "SCP format is not supported for f0 and excitation."
            )
        if use_local_condition:
            raise NotImplementedError("Not supported.")
        if use_global_condition:
            raise NotImplementedError("Not supported.")
        if use_aux_input:
            train_dataset = AudioMelSCPDataset(
                wav_scp=args.train_wav_scp,
                feats_scp=args.train_feats_scp,
                segments=args.train_segments,
                mel_length_threshold=mel_length_threshold,
                allow_cache=config.get("allow_cache", False),  # keep compatibility
            )
        else:
            train_dataset = AudioSCPDataset(
                wav_scp=args.train_wav_scp,
                segments=args.train_segments,
                audio_length_threshold=audio_length_threshold,
                allow_cache=config.get("allow_cache", False),  # keep compatibility
            )

    # define dataset for validation
    if args.dev_dumpdir is not None:
        if not use_f0_and_excitation:
            if use_aux_input:
                dev_dataset = AudioMelDataset(
                    root_dir=args.dev_dumpdir,
                    audio_query=audio_query,
                    audio_load_fn=audio_load_fn,
                    mel_query=mel_query,
                    mel_load_fn=mel_load_fn,
                    local_query=local_query,
                    local_load_fn=local_load_fn,
                    global_query=global_query,
                    global_load_fn=global_load_fn,
                    mel_length_threshold=mel_length_threshold,
                    allow_cache=config.get("allow_cache", False),  # keep compatibility
                )
            else:
                dev_dataset = AudioDataset(
                    root_dir=args.dev_dumpdir,
                    audio_query=audio_query,
                    audio_load_fn=audio_load_fn,
                    local_query=local_query,
                    local_load_fn=local_load_fn,
                    global_query=global_query,
                    global_load_fn=global_load_fn,
                    audio_length_threshold=audio_length_threshold,
                    allow_cache=config.get("allow_cache", False),  # keep compatibility
                )
        else:
            dev_dataset = AudioMelF0ExcitationDataset(
                root_dir=args.dev_dumpdir,
                audio_query=audio_query,
                mel_query=mel_query,
                f0_query=f0_query,
                excitation_query=excitation_query,
                audio_load_fn=audio_load_fn,
                mel_load_fn=mel_load_fn,
                f0_load_fn=f0_load_fn,
                excitation_load_fn=excitation_load_fn,
                mel_length_threshold=mel_length_threshold,
                allow_cache=config.get("allow_cache", False),  # keep compatibility
            )
    else:
        if use_f0_and_excitation:
            raise NotImplementedError(
                "SCP format is not supported for f0 and excitation."
            )
        if use_local_condition:
            raise NotImplementedError("Not supported.")
        if use_global_condition:
            raise NotImplementedError("Not supported.")
        if use_aux_input:
            dev_dataset = AudioMelSCPDataset(
                wav_scp=args.dev_wav_scp,
                feats_scp=args.dev_feats_scp,
                segments=args.dev_segments,
                mel_length_threshold=mel_length_threshold,
                allow_cache=config.get("allow_cache", False),  # keep compatibility
            )
        else:
            dev_dataset = AudioSCPDataset(
                wav_scp=args.dev_wav_scp,
                segments=args.dev_segments,
                audio_length_threshold=audio_length_threshold,
                allow_cache=config.get("allow_cache", False),  # keep compatibility
            )

    # store into dataset dict
    dataset = {
        "train": train_dataset,
        "dev": dev_dataset,
    }
    logging.info(f"The number of training files = {len(train_dataset)}.")
    logging.info(f"The number of development files = {len(dev_dataset)}.")

    # get data loader
    collater = Collater(
        batch_max_steps=config["batch_max_steps"],
        hop_size=config.get("hop_size", None),
        aux_context_window=config["generator_params"].get("aux_context_window", 0),
        use_f0_and_excitation=use_f0_and_excitation,
        use_noise_input=use_noise_input,
        use_aux_input=use_aux_input,
        use_duration=use_duration,
        use_global_condition=use_global_condition,
        use_local_condition=use_local_condition,
        pad_value=config["generator_params"].get(
            "num_embs", 0
        ),  # assume 0-based discrete symbol
    )
    sampler = {"train": None, "dev": None}
    if args.distributed:
        # setup sampler for distributed training
        from torch.utils.data.distributed import DistributedSampler

        sampler["train"] = DistributedSampler(
            dataset=dataset["train"],
            num_replicas=args.world_size,
            rank=args.rank,
            shuffle=True,
        )
        sampler["dev"] = DistributedSampler(
            dataset=dataset["dev"],
            num_replicas=args.world_size,
            rank=args.rank,
            shuffle=False,
        )
    data_loader = {
        "train": DataLoader(
            dataset=dataset["train"],
            shuffle=False if args.distributed else True,
            collate_fn=collater,
            batch_size=config["batch_size"],
            num_workers=config["num_workers"],
            sampler=sampler["train"],
            pin_memory=config["pin_memory"],
        ),
        "dev": DataLoader(
            dataset=dataset["dev"],
            shuffle=False if args.distributed else True,
            collate_fn=collater,
            batch_size=config["batch_size"],
            num_workers=config["num_workers"],
            sampler=sampler["dev"],
            pin_memory=config["pin_memory"],
        ),
    }

    # define models
    generator_class = getattr(
        parallel_wavegan.models,
        # keep compatibility
        config.get("generator_type", "ParallelWaveGANGenerator"),
    )
    discriminator_class = getattr(
        parallel_wavegan.models,
        # keep compatibility
        config.get("discriminator_type", "ParallelWaveGANDiscriminator"),
    )
    model = {
        "generator": generator_class(
            **config["generator_params"],
        ).to(device),
        "discriminator": discriminator_class(
            **config["discriminator_params"],
        ).to(device),
    }
    
    generator_params = sum(p.numel() for p in model["generator"].parameters())
    discriminator_params = sum(p.numel() for p in model["discriminator"].parameters())
    logging.info("number of generator parameters:", generator_params)
    logging.info("number of discriminator parameters:", discriminator_params)
    
    # define criterions
    criterion = {
        "gen_adv": GeneratorAdversarialLoss(
            # keep compatibility
            **config.get("generator_adv_loss_params", {})
        ).to(device),
        "dis_adv": DiscriminatorAdversarialLoss(
            # keep compatibility
            **config.get("discriminator_adv_loss_params", {})
        ).to(device),
        "mse": torch.nn.MSELoss().to(device),
    }
    if config.get("use_stft_loss", True):  # keep compatibility
        config["use_stft_loss"] = True
        criterion["stft"] = MultiResolutionSTFTLoss(
            **config["stft_loss_params"],
        ).to(device)
    if config.get("use_subband_stft_loss", False):  # keep compatibility
        assert config["generator_params"]["out_channels"] > 1
        criterion["sub_stft"] = MultiResolutionSTFTLoss(
            **config["subband_stft_loss_params"],
        ).to(device)
    else:
        config["use_subband_stft_loss"] = False
    if config.get("use_feat_match_loss", False):  # keep compatibility
        criterion["feat_match"] = FeatureMatchLoss(
            # keep compatibility
            **config.get("feat_match_loss_params", {}),
        ).to(device)
    else:
        config["use_feat_match_loss"] = False
    if config.get("use_mel_loss", False):  # keep compatibility
        if config.get("mel_loss_params", None) is None:
            criterion["mel"] = MelSpectrogramLoss(
                fs=config["sampling_rate"],
                fft_size=config["fft_size"],
                hop_size=config["hop_size"],
                win_length=config["win_length"],
                window=config["window"],
                num_mels=config["num_mels"],
                fmin=config["fmin"],
                fmax=config["fmax"],
            ).to(device)
        else:
            criterion["mel"] = MelSpectrogramLoss(
                **config["mel_loss_params"],
            ).to(device)
    else:
        config["use_mel_loss"] = False
    if config.get("use_duration_loss", False):  # keep compatibility
        if config.get("duration_loss_params", None) is None:
            criterion["duration"] = DurationPredictorLoss(
                offset=config["offset"],
                reduction=config["reduction"],
            ).to(device)
        else:
            criterion["duration"] = DurationPredictorLoss(
                **config["duration_loss_params"],
            ).to(device)
    else:
        config["use_duration_loss"] = False

    # define special module for subband processing
    if config["generator_params"]["out_channels"] > 1:
        criterion["pqmf"] = PQMF(
            subbands=config["generator_params"]["out_channels"],
            # keep compatibility
            **config.get("pqmf_params", {}),
        ).to(device)

    # if config["use_istft"]:
    #     criterion["istft"] = ISTFT(
    #         config["fft_size"], config["hop_size"], config["win_size"], device=device)
        
    # define optimizers and schedulers
    generator_optimizer_class = getattr(
        parallel_wavegan.optimizers,
        # keep compatibility
        config.get("generator_optimizer_type", "RAdam"),
    )
    discriminator_optimizer_class = getattr(
        parallel_wavegan.optimizers,
        # keep compatibility
        config.get("discriminator_optimizer_type", "RAdam"),
    )
    optimizer = {
        "generator": generator_optimizer_class(
            model["generator"].parameters(),
            **config["generator_optimizer_params"],
        ),
        "discriminator": discriminator_optimizer_class(
            model["discriminator"].parameters(),
            **config["discriminator_optimizer_params"],
        ),
    }
    generator_scheduler_class = getattr(
        torch.optim.lr_scheduler,
        # keep compatibility
        config.get("generator_scheduler_type", "StepLR"),
    )
    discriminator_scheduler_class = getattr(
        torch.optim.lr_scheduler,
        # keep compatibility
        config.get("discriminator_scheduler_type", "StepLR"),
    )
    scheduler = {
        "generator": generator_scheduler_class(
            optimizer=optimizer["generator"],
            **config["generator_scheduler_params"],
        ),
        "discriminator": discriminator_scheduler_class(
            optimizer=optimizer["discriminator"],
            **config["discriminator_scheduler_params"],
        ),
    }
    if args.distributed:
        # wrap model for distributed training
        # try:
        #     from apex.parallel import DistributedDataParallel
        # except ImportError:
        #     raise ImportError(
        #         "apex is not installed. please check https://github.com/NVIDIA/apex."
        #     )
        from torch.nn.parallel import DistributedDataParallel
        model["generator"] = DistributedDataParallel(model["generator"],find_unused_parameters=True)
        model["discriminator"] = DistributedDataParallel(model["discriminator"],find_unused_parameters=True)

    # show settings
    logging.info(model["generator"])
    logging.info(model["discriminator"])
    logging.info(optimizer["generator"])
    logging.info(optimizer["discriminator"])
    logging.info(scheduler["generator"])
    logging.info(scheduler["discriminator"])
    for criterion_ in criterion.values():
        logging.info(criterion_)

    # define trainer
    trainer = Trainer(
        steps=0,
        epochs=0,
        data_loader=data_loader,
        sampler=sampler,
        model=model,
        criterion=criterion,
        optimizer=optimizer,
        scheduler=scheduler,
        config=config,
        device=device,
    )

    # load pretrained parameters from checkpoint
    if len(args.pretrain) != 0:
        trainer.load_checkpoint(args.pretrain, load_only_params=True)
        logging.info(f"Successfully load parameters from {args.pretrain}.")

    # resume from checkpoint
    if len(args.resume) != 0:
        trainer.load_checkpoint(args.resume)
        logging.info(f"Successfully resumed from {args.resume}.")

    # run training loop
    try:
        trainer.run()
    finally:
        trainer.save_checkpoint(
            os.path.join(config["outdir"], f"checkpoint-{trainer.steps}steps.pkl")
        )
        logging.info(f"Successfully saved checkpoint @ {trainer.steps}steps.")


if __name__ == "__main__":
    main()