From 062483df1a0df850f4a5b511e8f004b93dc273f9 Mon Sep 17 00:00:00 2001 From: bmalaski <5288648+bmalaski@users.noreply.github.com> Date: Thu, 20 Jun 2024 13:56:31 -0400 Subject: [PATCH 1/7] start merge --- .gitignore | 173 ++- metrics_logger.py | 161 +++ trainer_alltalk/trainer.py | 2197 ++++++++++++++++++++++++++++++++++++ 3 files changed, 2527 insertions(+), 4 deletions(-) create mode 100644 metrics_logger.py create mode 100644 trainer_alltalk/trainer.py diff --git a/.gitignore b/.gitignore index c1d928d..563c227 100644 --- a/.gitignore +++ b/.gitignore @@ -1,8 +1,173 @@ -voices/*.* -models/*.* -outputs/*.* -finetune/*.* +voices/* +models/* +outputs/* +finetune/* +alltalk_environment/* system/config/at_github_sha.json system/tts_engines/tts_engines.json confignew.json diagnostics.log + +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/latest/usage/project/#working-with-version-control +.pdm.toml +.pdm-python +.pdm-build/ + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. + +.idea/* \ No newline at end of file diff --git a/metrics_logger.py b/metrics_logger.py new file mode 100644 index 0000000..342ce8d --- /dev/null +++ b/metrics_logger.py @@ -0,0 +1,161 @@ +import datetime +import io +import logging + +from matplotlib import pyplot as plt +from trainer import ConsoleLogger +from trainer.logging.console_logger import tcolors +from trainer.utils.distributed import rank_zero_only + +logger = logging.getLogger("trainer") + +class MetricsLogger(ConsoleLogger): + def __init__(self, ): + super().__init__() + self.cur_epoch_start = None + self.cur_epoch_end = None + self.cur_epoch = None + self.epoch_durations = [] + self.max_epoch = None + self.batch_steps = None + self.global_steps = None + self.loss = [] + self.avg_loss = [] + self.loss_mel_ce = [] + self.avg_loss_mel_ce = [] + self.loss_text_ce = [] + self.avg_loss_text_ce = [] + self.learning_rates = [] + self.total_duration = None + self.estimated_duration = None + self.start = None + + @staticmethod + def get_now(): + return datetime.datetime.now() + @staticmethod + def format_duration(seconds): + if seconds: + # Calculate hours, minutes, and seconds + hours = seconds // 3600 + minutes = (seconds % 3600) // 60 + seconds = seconds % 60 + return f"{int(hours)}:{int(minutes)}:{int(seconds)}" + else: + return None + + @rank_zero_only + def print_epoch_start(self, epoch, max_epoch, output_path=None): + self.cur_epoch = epoch + self.cur_epoch_start = self.get_now() + self.max_epoch = max_epoch + super().print_epoch_start(epoch, max_epoch, output_path) + + @rank_zero_only + def print_train_start(self): + super().print_train_start() + + @rank_zero_only + def print_train_step(self, batch_steps, step, global_step, loss_dict, avg_loss_dict): + self.batch_steps = batch_steps + self.global_steps = global_step + + for key, value in loss_dict.items(): + if key == 'loss_mel_ce': + self.loss_mel_ce.append((global_step, value)) + elif key == 'loss': + self.loss.append((global_step, value)) + elif key == 'loss_text_ce': + self.loss_text_ce.append((global_step, value)) + elif key == 'current_lr': + self.learning_rates.append((global_step, value)) + super().print_train_step(batch_steps, step, global_step, loss_dict, avg_loss_dict) + + # pylint: disable=unused-argument + @rank_zero_only + def print_train_epoch_end(self, global_step, epoch, epoch_time, print_dict): + super().print_epoch_end(global_step, epoch, epoch_time, print_dict) + if self.estimated_duration: + self.log_with_flush( + "\n{}{} >> ETA: {} {}".format(tcolors.UNDERLINE, tcolors.BOLD, self.format_duration(self.estimated_duration), tcolors.ENDC), + ) + + @rank_zero_only + def print_eval_start(self): + super().print_eval_start() + + @rank_zero_only + def print_eval_step(self, step, loss_dict, avg_loss_dict): + super().print_eval_step(step, loss_dict, avg_loss_dict) + + @rank_zero_only + def print_epoch_end(self, epoch, avg_loss_dict): + self.cur_epoch_end = self.get_now() + self.epoch_durations.append((self.cur_epoch_end - self.cur_epoch_start).total_seconds()) + for key, value in avg_loss_dict.items(): + if key == 'avg_loss_mel_ce': + self.avg_loss_mel_ce.append((epoch, value)) + elif key == 'avg_loss': + self.avg_loss.append((epoch, value)) + elif key == 'avg_loss_text_ce': + self.avg_loss_text_ce.append((epoch, value)) + + if len(self.epoch_durations) > 1 and self.max_epoch: + self.total_duration = sum(self.epoch_durations) + additional_data_points_needed = self.max_epoch - len(self.epoch_durations) + if additional_data_points_needed > 0: + weighted_avg_duration = sum((i + 1) * duration for i, duration in enumerate(self.epoch_durations)) / sum(range(1, len(self.epoch_durations) + 1)) + self.estimated_duration = (weighted_avg_duration * additional_data_points_needed) + + super().print_epoch_end(epoch, avg_loss_dict) + + def plot_metrics(self, show_loss_mel_ce=False, show_loss_text_ce=False, show_loss=True): + fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(24, 18)) + + ax2.set_title('Steps') + ax2.set_xlabel('Steps') + ax2.set_ylabel('Value') + + ax1.set_title('Epochs') + ax1.set_xlabel('Epoch') + ax1.set_ylabel('Value') + + ax3.set_title('Learning Rates') + ax3.set_xlabel('Steps') + ax3.set_ylabel('Value') + + if self.loss and self.loss_text_ce and self.loss_mel_ce: + if show_loss: + ax2.plot(*zip(*self.loss), color='green', linestyle='-', label='Loss') + ax2.fill_between(*zip(*self.loss), color='lightgreen', alpha=0.1) + if show_loss_text_ce: + ax2.plot(*zip(*self.loss_text_ce), color='red', linestyle='-', label='Loss Text') + ax2.fill_between(*zip(*self.loss_text_ce), color='lightcoral', alpha=0.1) + if show_loss_mel_ce: + ax2.plot(*zip(*self.loss_mel_ce), color='blue', linestyle='-', label='Loss MEL') + ax2.fill_between(*zip(*self.loss_mel_ce), color='lightblue', alpha=0.1) + ax2.legend() + + if self.learning_rates: + ax3.plot(*zip(*self.learning_rates), linestyle='-', color='blue', label='Learning Rate') + ax3.fill_between(*zip(*self.learning_rates), color='lightblue', alpha=0.1) + ax3.legend() + + if self.avg_loss and self.avg_loss_text_ce and self.avg_loss_mel_ce: + if show_loss: + ax1.plot(*zip(*self.avg_loss), color='green', linestyle='-', label='Avg Loss') + ax1.fill_between(*zip(*self.avg_loss), color='lightgreen', alpha=0.1) + if show_loss_text_ce: + ax1.plot(*zip(*self.avg_loss_text_ce), color='red', linestyle='-', label='Avg Loss Text') + ax1.fill_between(*zip(*self.avg_loss_text_ce), color='lightcoral', alpha=0.1) + if show_loss_mel_ce: + ax1.plot(*zip(*self.avg_loss_mel_ce), color='blue', linestyle='-', label='Avg Loss MEL') + ax1.fill_between(*zip(*self.avg_loss_mel_ce), color='lightblue', alpha=0.1) + ax1.set_xlim(left=0) + ax1.legend() + + plt.style.use('dark_background') + plt.tight_layout() + plt.savefig("training_metrics.png", format='png', bbox_inches='tight', pad_inches=0.1) + plt.close() + return "training_metrics.png" \ No newline at end of file diff --git a/trainer_alltalk/trainer.py b/trainer_alltalk/trainer.py new file mode 100644 index 0000000..e786a2e --- /dev/null +++ b/trainer_alltalk/trainer.py @@ -0,0 +1,2197 @@ +# -*- coding: utf-8 -*- +import gc +import importlib +import logging +import os +import platform +import shutil +import sys +import time +import traceback +from contextlib import nullcontext +from dataclasses import dataclass, field +from inspect import signature +from typing import Callable, Dict, List, Tuple, Union + +import torch +import torch.distributed as dist +from TTS.tts.layers.xtts.trainer.gpt_trainer import GPTTrainerConfig, GPTArgs, XttsAudioConfig +from coqpit import Coqpit +import json +from torch import nn +from torch.nn.parallel import DistributedDataParallel as DDP_th +from torch.optim.lr_scheduler import ReduceLROnPlateau +from torch.utils.data import DataLoader + +from trainer.analytics import ping_training_run +from trainer.callbacks import TrainerCallback +from trainer.generic_utils import ( + KeepAverage, + count_parameters, + get_experiment_folder_path, + get_git_branch, + isimplemented, + remove_experiment_folder, + set_partial_state_dict, + to_cuda, +) +from trainer.io import ( + copy_model_files, + get_last_checkpoint, + load_fsspec, + save_best_model, + save_checkpoint, +) +from trainer.logging import ConsoleLogger, DummyLogger, logger_factory +from trainer.trainer_utils import ( + get_optimizer, + get_scheduler, + is_apex_available, + print_training_env, + setup_torch_training_env, +) +from trainer.utils.cuda_memory import cuda_meminfo, should_reduce_batch_size +from trainer.utils.distributed import ( + get_rank, + init_distributed, + rank_zero_logger_info, + rank_zero_only, +) + +logger = logging.getLogger("trainer") + +if is_apex_available(): + from apex import amp # pylint: disable=import-error + + +@dataclass +class TrainerConfig(Coqpit): + """Config fields tweaking the Trainer for a model. + A ````ModelConfig```, by inheriting ```TrainerConfig``` must be defined for using πŸ‘Ÿ. + Inherit this by a new model config and override the fields as needed. + All the fields can be overridden from comman-line as ```--coqpit.arg_name=value```. + + Example:: + + Run the training code by overriding the ```lr``` and ```plot_step``` fields. + + >>> python train.py --coqpit.plot_step=22 --coqpit.lr=0.001 + + Defining a model using ```TrainerConfig```. + + >>> from trainer import TrainerConfig + >>> class MyModelConfig(TrainerConfig): + ... optimizer: str = "Adam" + ... lr: float = 0.001 + ... epochs: int = 1 + ... ... + >>> class MyModel(nn.module): + ... def __init__(self, config): + ... ... + >>> model = MyModel(MyModelConfig()) + + """ + + # Fields for the run + output_path: str = field(default="output") + logger_uri: str = field( + default=None, + metadata={ + "help": "URI to save training artifacts by the logger. If not set, logs will be saved in the output_path. Defaults to None" + }, + ) + run_name: str = field(default="run", metadata={"help": "Name of the run. Defaults to 'run'"}) + project_name: str = field(default=None, metadata={"help": "Name of the project. Defaults to None"}) + run_description: str = field( + default="🐸Coqui trainer run.", + metadata={"help": "Notes and description about the run. Defaults to '🐸Coqui trainer run.'"}, + ) + # Fields for logging + print_step: int = field( + default=25, metadata={"help": "Print training stats on the terminal every print_step steps. Defaults to 25"} + ) + plot_step: int = field( + default=100, metadata={"help": "Plot training stats on the logger every plot_step steps. Defaults to 100"} + ) + model_param_stats: bool = field( + default=False, metadata={"help": "Log model parameters stats on the logger dashboard. Defaults to False"} + ) + wandb_entity: str = field(default=None, metadata={"help": "Wandb entity to log the run. Defaults to None"}) + dashboard_logger: str = field( + default="tensorboard", metadata={"help": "Logger to use for the tracking dashboard. Defaults to 'tensorboard'"} + ) + # Fields for checkpointing + save_on_interrupt: bool = field( + default=True, metadata={"help": "Save checkpoint on interrupt (Ctrl+C). Defaults to True"} + ) + log_model_step: int = field( + default=None, + metadata={ + "help": "Save checkpoint to the logger every log_model_step steps. If not defined `save_step == log_model_step`." + }, + ) + save_step: int = field( + default=10000, metadata={"help": "Save local checkpoint every save_step steps. Defaults to 10000"} + ) + save_n_checkpoints: int = field(default=5, metadata={"help": "Keep n local checkpoints. Defaults to 5"}) + save_checkpoints: bool = field(default=True, metadata={"help": "Save checkpoints locally. Defaults to True"}) + save_all_best: bool = field( + default=False, metadata={"help": "Save all best checkpoints and keep the older ones. Defaults to False"} + ) + save_best_after: int = field( + default=0, metadata={"help": "Wait N steps to save best checkpoints. Defaults to 0"} + ) + target_loss: str = field( + default=None, metadata={"help": "Target loss name to select the best model. Defaults to None"} + ) + # Fields for eval and test run + print_eval: bool = field(default=False, metadata={"help": "Print eval steps on the terminal. Defaults to False"}) + test_delay_epochs: int = field(default=0, metadata={"help": "Wait N epochs before running the test. Defaults to 0"}) + run_eval: bool = field( + default=True, metadata={"help": "Run evalulation epoch after training epoch. Defaults to True"} + ) + run_eval_steps: int = field( + default=None, + metadata={ + "help": "Run evalulation epoch after N steps. If None, waits until training epoch is completed. Defaults to None" + }, + ) + # Fields for distributed training + distributed_backend: str = field( + default="nccl", metadata={"help": "Distributed backend to use. Defaults to 'nccl'"} + ) + distributed_url: str = field( + default="tcp://localhost:54321", + metadata={"help": "Distributed url to use. Defaults to 'tcp://localhost:54321'"}, + ) + # Fields for training specs + mixed_precision: bool = field(default=False, metadata={"help": "Use mixed precision training. Defaults to False"}) + precision: str = field( + default="fp16", + metadata={ + "help": "Precision to use in mixed precision training. `fp16` for float16 and `bf16` for bfloat16. Defaults to 'f16'" + }, + ) + epochs: int = field(default=1000, metadata={"help": "Number of epochs to train. Defaults to 1000"}) + batch_size: int = field(default=32, metadata={"help": "Batch size to use. Defaults to 32"}) + eval_batch_size: int = field(default=16, metadata={"help": "Batch size to use for eval. Defaults to 16"}) + grad_clip: float = field( + default=0.0, metadata={"help": "Gradient clipping value. Disabled if <= 0. Defaults to 0.0"} + ) + scheduler_after_epoch: bool = field( + default=True, + metadata={"help": "Step the scheduler after each epoch else step after each iteration. Defaults to True"}, + ) + # Fields for optimzation + lr: Union[float, List[float]] = field( + default=0.001, metadata={"help": "Learning rate for each optimizer. Defaults to 0.001"} + ) + optimizer: Union[str, List[str]] = field(default=None, metadata={"help": "Optimizer(s) to use. Defaults to None"}) + optimizer_params: Union[Dict, List[Dict]] = field( + default_factory=dict, metadata={"help": "Optimizer(s) arguments. Defaults to {}"} + ) + lr_scheduler: Union[str, List[str]] = field( + default=None, metadata={"help": "Learning rate scheduler(s) to use. Defaults to None"} + ) + lr_scheduler_params: Dict = field( + default_factory=dict, metadata={"help": "Learning rate scheduler(s) arguments. Defaults to {}"} + ) + use_grad_scaler: bool = field( + default=False, + metadata={ + "help": "Enable/disable gradient scaler explicitly. It is enabled by default with AMP training. Defaults to False" + }, + ) + allow_tf32: bool = field( + default=False, + metadata={ + "help": "A bool that controls whether TensorFloat-32 tensor cores may be used in matrix multiplications on Ampere or newer GPUs. Default to False." + }, + ) + cudnn_enable: bool = field(default=True, metadata={"help": "Enable/disable cudnn explicitly. Defaults to True"}) + cudnn_deterministic: bool = field( + default=False, + metadata={ + "help": "Enable/disable deterministic cudnn operations. Set this True for reproducibility but it slows down training significantly. Defaults to False." + }, + ) + cudnn_benchmark: bool = field( + default=False, + metadata={ + "help": "Enable/disable cudnn benchmark explicitly. Set this False if your input size change constantly. Defaults to False" + }, + ) + training_seed: int = field( + default=54321, + metadata={"help": "Global seed for torch, random and numpy random number generator. Defaults to 54321"}, + ) + + +@dataclass +class TrainerArgs(Coqpit): + """Trainer arguments that can be accessed from the command line. + + Examples:: + >>> python train.py --restore_path /path/to/checkpoint.pth + """ + + continue_path: str = field( + default="", + metadata={ + "help": "Path to a training folder to continue training. Restore the model from the last checkpoint and continue training under the same folder." + }, + ) + restore_path: str = field( + default="", + metadata={ + "help": "Path to a model checkpoit. Restore the model with the given checkpoint and start a new training." + }, + ) + best_path: str = field( + default="", + metadata={ + "help": "Best model file to be used for extracting the best loss. If not specified, the latest best model in continue path is used" + }, + ) + use_ddp: bool = field( + default=False, + metadata={"help": "Use DDP in distributed training. It is to set in `distribute.py`. Do not set manually."}, + ) + use_accelerate: bool = field(default=False, metadata={"help": "Use HF Accelerate as the back end for training."}) + grad_accum_steps: int = field( + default=1, + metadata={ + "help": "Number of gradient accumulation steps. It is used to accumulate gradients over multiple batches." + }, + ) + overfit_batch: bool = field(default=False, metadata={"help": "Overfit a single batch for debugging."}) + skip_train_epoch: bool = field( + default=False, + metadata={"help": "Skip training and only run evaluation and test."}, + ) + start_with_eval: bool = field( + default=False, + metadata={"help": "Start with evaluation and test."}, + ) + small_run: int = field( + default=None, + metadata={ + "help": "Only use a subset of the samples for debugging. Set the number of samples to use. Defaults to None. " + }, + ) + gpu: int = field( + default=None, metadata={"help": "GPU ID to use if ```CUDA_VISIBLE_DEVICES``` is not set. Defaults to None."} + ) + # only for DDP + rank: int = field(default=0, metadata={"help": "Process rank in a distributed training. Don't set manually."}) + group_id: str = field( + default="", metadata={"help": "Process group id in a distributed training. Don't set manually."} + ) + + +class Trainer: + def __init__( # pylint: disable=dangerous-default-value + self, + args: TrainerArgs, + config: Coqpit, + output_path: str, + c_logger: ConsoleLogger = None, + dashboard_logger: "Logger" = None, + model: nn.Module = None, + get_model: Callable = None, + get_data_samples: Callable = None, + train_samples: List = None, + eval_samples: List = None, + test_samples: List = None, + train_loader: DataLoader = None, + eval_loader: DataLoader = None, + training_assets: Dict = {}, + parse_command_line_args: bool = True, + callbacks: Dict[str, Callable] = {}, + gpu: int = None, + ) -> None: + """Simple yet powerful πŸΈπŸ’¬ TTS trainer for PyTorch. It can train all the available `tts` and `vocoder` models + or easily be customized. + + Notes: + + Supports Automatic Mixed Precision training. If `Apex` is availabe, it automatically picks that, else + it uses PyTorch's native `amp` module. `Apex` may provide more stable training in some cases. + + Args: + + args (Union[Coqpit, Namespace]): Training arguments parsed either from console by `argparse` or `TrainerArgs` + config object. + + config (Coqpit): Model config object. It includes all the values necessary for initializing, training, evaluating + and testing the model. + + output_path (str): Path to the output training folder. All the files are saved under thi path. + + c_logger (ConsoleLogger, optional): Console logger for printing training status. If not provided, the default + console logger is used. Defaults to None. + + dashboard_logger Union[TensorboardLogger, WandbLogger]: Dashboard logger. If not provided, the tensorboard logger is used. + Defaults to None. + + model (nn.Module, optional): Initialized and ready-to-train model. If it is not defined, `Trainer` + initializes a model from the provided config. Defaults to None. + + get_model (Callable): + A function that returns a model. It is used to initialize the model when `model` is not provided. + It either takes the config as the only argument or does not take any argument. + Defaults to None + + get_data_samples (Callable): + A function that returns a list of training and evaluation samples. Used if `train_samples` and + `eval_samples` are None. Defaults to None. + + train_samples (List): + A list of training samples used by the model's `get_train_data_loader` to init the `dataset` and the + `data_loader`. Defaults to None. + + eval_samples (List): + A list of evaluation samples used by the model's `get_eval_data_loader` to init the `dataset` and the + `data_loader`. Defaults to None. + + train_loader (DataLoader): + A pytorch data loader object for training epochs. Leave as None if you want it to be made during training. Defaults to None. + + eval_loader (DataLoader): + A pytorch data loader object for evaluation epochs. Leave as None to be generated during training. Defaults to None. + + test_samples (List): + A list of test samples used by the model's `get_test_data_loader` to init the `dataset` and the + `data_loader`. If None, the ```model.test_run()``` is expected to load the data. Defaults to None. + + training_assets (Dict): + A dictionary of assets to be used at training and passed to the model's ```train_log(), eval_log(), get_data_loader()``` + during training. It can include `AudioProcessor` or/and `Tokenizer`. Defaults to {}. + + parse_command_line_args (bool): + If true, parse command-line arguments and update `TrainerArgs` and model `config` values. Set it + to false if you parse the arguments yourself. Defaults to True. + + callbacks (Dict[str, Callable]): + A dictionary of callbacks to be used during training. The keys are the callback names and the values + + gpu (int): + GPU ID to use for training If "CUDA_VISIBLE_DEVICES" is not set. Defaults to None. + + Example:: + + Running trainer with a model. + + >>> args = TrainerArgs(...) + >>> config = ModelConfig(...) + >>> model = Model(config) + >>> trainer = Trainer(args, config, output_path, model=model) + >>> trainer.fit() + + TODO: + - Wrap model for not calling .module in DDP. + - Deepspeed integration + - Profiler integration. + - Overfitting to a batch. + - TPU training + """ + if parse_command_line_args: + # parse command-line arguments to override TrainerArgs() + args, coqpit_overrides = self.parse_argv(args) + + # get ready for training and parse command-line arguments to override the model config + config, new_fields = self.init_training(args, coqpit_overrides, config) + elif args.continue_path or args.restore_path: + config, new_fields = self.init_training(args, {}, config) + else: + new_fields = {} + + # set the output path + if args.continue_path: + # use the same path as the continuing run + output_path = args.continue_path + else: + # override the output path if it is provided + output_path = config.output_path if output_path is None else output_path + # create a new output folder name + output_path = get_experiment_folder_path(config.output_path, config.run_name) + os.makedirs(output_path, exist_ok=True) + + # copy training assets to the output folder + copy_model_files(config, output_path, new_fields) + + # init class members + self.args = args + self.config = config + self.output_path = output_path + self.training_assets = training_assets + self.grad_accum_steps = args.grad_accum_steps + self.overfit_batch = args.overfit_batch + self.skip_train_epoch = args.skip_train_epoch + self.start_with_eval = args.start_with_eval + + assert self.grad_accum_steps > 0, " [!] grad_accum_steps must be greater than 0." + + # setup logging + log_file = os.path.join(self.output_path, f"trainer_{args.rank}_log.txt") + self._setup_logger_config(log_file) + + # setup training environment + self.use_cuda, self.num_gpus = self.setup_training_environment(args=args, config=config, gpu=gpu) + + # init loggers + self.dashboard_logger, self.c_logger = self.init_loggers(self.config, output_path, dashboard_logger, c_logger) + # self.c_logger.logger = logger + + if not self.config.log_model_step: + self.config.log_model_step = self.config.save_step + + # make sure that start_with_eval is disabled if eval is disabled + if not self.config.run_eval and self.start_with_eval: + self.start_with_eval = False + + self.total_steps_done = 0 + self.epochs_done = 0 + self.restore_step = 0 + self.restore_epoch = 0 + self.best_loss = {"train_loss": float("inf"), "eval_loss": float("inf") if self.config.run_eval else None} + self.train_loader = None + self.test_loader = None + self.eval_loader = None + + self.keep_avg_train = None + self.keep_avg_eval = None + + self.use_amp_scaler = ( + self.use_cuda + if self.config.mixed_precision and self.config.precision == "fp16" + else self.config.use_grad_scaler + ) + + if train_samples is not None: + # use the provided samples + self.train_samples = train_samples + self.eval_samples = eval_samples + self.test_samples = test_samples + elif get_data_samples is not None: + # run `get_data_samples` to init the data samples + ( # pylint: disable=unbalanced-tuple-unpacking + self.train_samples, + self.eval_samples, + self.test_samples, + ) = self.run_get_data_samples(config, get_data_samples) + else: + # expecting to load the samples in `model.get_data_loader()` + self.train_samples = None + self.eval_samples = None + self.test_samples = None + + # define custom train and eval loader + self.train_loader = train_loader + self.eval_loader = eval_loader + + # only use a subset of the samples if small_run is set + self.setup_small_run(args.small_run) + + # init the model + if model is None and get_model is None: + raise ValueError("[!] `model` and `get_model` cannot both be None.") + if model is not None: + self.model = model + else: + self.run_get_model(self.config, get_model) + + # init model's training assets + if isimplemented(self.model, "init_for_training"): + self.model.init_for_training() + + # setup criterion + self.criterion = self.get_criterion(self.model) + + # DISTRUBUTED + if self.use_pt_ddp: + rank_zero_logger_info(" > Using PyTorch DDP", logger) + init_distributed( + args.rank, + self.num_gpus, + args.group_id, + self.config.distributed_backend, + self.config.distributed_url, + ) + + if self.use_cuda: + self.model.cuda() + if isinstance(self.criterion, list): + for criterion in self.criterion: + if isinstance(criterion, torch.nn.Module): + criterion.cuda() + else: + if isinstance(self.criterion, torch.nn.Module): + self.criterion.cuda() + + # setup optimizer + self.optimizer = self.get_optimizer(self.model, self.config) + + # If multiple-optimizer setup with grad accumulation and without custom optimize method raise an error + if ( + self.grad_accum_steps != 1 + and isinstance(self.optimizer, list) + and not isimplemented(self.model, "optimize") + ): + raise ValueError( + " [!] Coqui Trainer does not support grad_accum_steps for multiple-optimizer setup, please set grad_accum_steps to 1 or implement in your model a custom method called Β΄optimize` that need to deal with dangling gradients in multiple-optimizer setup!" + ) + + # CALLBACK + self.callbacks = TrainerCallback() + self.callbacks.parse_callbacks_dict(callbacks) + self.callbacks.on_init_start(self) + + # init AMP + if self.use_amp_scaler: + if self.use_apex: + self.scaler = None + self.model, self.optimizer = amp.initialize(self.model, self.optimizer, opt_level="O1") + self.scaler = torch.cuda.amp.GradScaler() + else: + self.scaler = None + + # restore model + if self.args.restore_path: + (self.model, self.optimizer, self.scaler, self.restore_step, self.restore_epoch) = self.restore_model( + self.config, args.restore_path, self.model, self.optimizer, self.scaler + ) + self.scaler = torch.cuda.amp.GradScaler() + + # setup scheduler + self.scheduler = self.get_scheduler(self.model, self.config, self.optimizer) + self.scheduler = self.restore_scheduler( + self.scheduler, self.args, self.config, self.restore_epoch, self.restore_step + ) + + # DISTRIBUTED + if self.use_pt_ddp: + self.model = DDP_th(self.model, device_ids=[args.rank], output_device=args.rank) + + # setup accelerator + self.setup_accelerate() + + # count model size + num_params = count_parameters(self.model) + rank_zero_logger_info(f"\n > Model has {num_params} parameters", logger) + + self.callbacks.on_init_end(self) + self.dashboard_logger.add_config(config) + self.save_training_script() + ping_training_run() + + @property + def use_apex(self): + """Return True if using APEX.""" + return not self.args.use_accelerate and self._is_apex_available() + + @property + def use_pt_ddp(self): + """Return True if using PyTorch DDP.""" + return self.num_gpus > 1 and not self.use_accelerate + + @property + def use_accelerate(self): + """Return True if using HF Accelerate.""" + return self.args.use_accelerate + + def setup_accelerate(self): + if self.use_accelerate: + self.model, self.optimizer, self.train_loader, self.scheduler, self.accelerator = self.init_accelerate( + model=self.model, + optimizer=self.optimizer, + training_dataloader=self.train_loader, + scheduler=self.scheduler, + grad_accum_steps=self.grad_accum_steps, + mixed_precision=self.config.mixed_precision, + precision=self.config.precision, + ) + + def prepare_accelerate_loader(self, data_loader): + """Prepare the accelerator for the training.""" + if self.use_accelerate: + return self.accelerator.prepare_data_loader(data_loader) + return data_loader + + @staticmethod + def init_accelerate(model, optimizer, training_dataloader, scheduler, grad_accum_steps, mixed_precision, precision): + """Setup HF Accelerate for the training.""" + + # check if accelerate is installed + try: + from accelerate import Accelerator # pylint:disable=import-outside-toplevel + except ImportError as e: + raise ImportError("Please install accelerate to use this feature.") from e + + _precision = precision if precision is not None else "f16" if mixed_precision else None + if _precision == "float16": + _precision = "f16" + elif _precision == "float8": + _precision = "f8" + elif _precision == "bfloat16": + _precision = "bf16" + accelerator = Accelerator(gradient_accumulation_steps=grad_accum_steps, mixed_precision=_precision) + if isinstance(model, torch.nn.Module): + model = accelerator.prepare_model(model) + + if isinstance(optimizer, dict): + for key, optim in optimizer.items(): + optimizer[key] = accelerator.prepare_optimizer(optim) + elif isinstance(optimizer, list): + for i, optim in enumerate(optimizer): + optimizer[i] = accelerator.prepare_optimizer(optim) + elif optimizer is not None: + optimizer = accelerator.prepare_optimizer(optimizer) + + if isinstance(training_dataloader, torch.utils.data.DataLoader): + training_dataloader = accelerator.prepare_data_loader(training_dataloader) + + if isinstance(scheduler, dict): + for key, sched in scheduler.items(): + scheduler[key] = accelerator.prepare_scheduler(sched) + elif isinstance(scheduler, list): + for i, sched in enumerate(scheduler): + scheduler[i] = accelerator.prepare_scheduler(sched) + elif scheduler is not None: + scheduler = accelerator.prepare_scheduler(scheduler) + + return model, optimizer, training_dataloader, scheduler, accelerator + + def save_training_script(self): + """Save the training script to tracking dashboard and output path.""" + file_path = sys.argv[0] + if os.path.isfile(file_path): + file_name = os.path.basename(file_path) + self.dashboard_logger.add_artifact(file_or_dir=file_path, name=file_name, artifact_type="file") + with open(file_path, "r", encoding="utf8") as f: + self.dashboard_logger.add_text("training-script", f"{f.read()}", 0) + shutil.copyfile(file_path, os.path.join(self.output_path, file_name)) + + @staticmethod + def parse_argv(args: Union[Coqpit, List]): + """Parse command line arguments to init or override `TrainerArgs()`.""" + if isinstance(args, Coqpit): + parser = args.init_argparse(arg_prefix="") + else: + train_config = TrainerArgs() + parser = train_config.init_argparse(arg_prefix="") + training_args, coqpit_overrides = parser.parse_known_args() + args.parse_args(training_args) + return args, coqpit_overrides + + @staticmethod + def init_loggers(config: "Coqpit", output_path: str, dashboard_logger=None, c_logger=None): + """Init console and dashboard loggers. + Use the given logger if passed externally else use config values to pick the right logger. + Return a dashboard logger only for the rank 0 process in DDP + Define a console logger for each process in DDP + + Args: + config (Coqpit): Model config. + output_path (str): Output path to save the training artifacts. + dashboard_logger (DashboardLogger): Object passed to the trainer from outside. + c_logger (ConsoleLogger): Object passed to the trained from outside. + + Returns: + Initialized dashboard_logger and console_logger objects. + """ + c_logger = ConsoleLogger() if c_logger is None else c_logger + + # only allow dashboard logging for the main process in DDP mode + if get_rank() > 0: + return DummyLogger(), c_logger + if dashboard_logger is None: + dashboard_logger = logger_factory(config, output_path) + return dashboard_logger, c_logger + + def setup_small_run(self, small_run: int = None): + """Use a subset of samples for training, evaluation and testing.""" + if small_run is not None: + logger.info("[!] Small Run, only using %i samples.", small_run) + self.train_samples = None if self.train_samples is None else self.train_samples[:small_run] + self.eval_samples = None if self.eval_samples is None else self.eval_samples[:small_run] + self.test_samples = None if self.test_samples is None else self.test_samples[:small_run] + + @staticmethod + def init_training(args: TrainerArgs, coqpit_overrides: Dict, config: Coqpit = None): + """Initialize training and update model configs from command line arguments. + + Args: + args (argparse.Namespace or dict like): Parsed trainer arguments. + config_overrides (argparse.Namespace or dict like): Parsed config overriding arguments. + config (Coqpit): Model config. If none, it is generated from `args`. Defaults to None. + + Returns: + config (Coqpit): Config paramaters. + """ + # set arguments for continuing training + if args.continue_path: + print("continue_path training from previous run.") + args.config_path = os.path.join(args.continue_path, "config.json") + args.restore_path, best_model = get_last_checkpoint(args.continue_path) + if not args.best_path: + args.best_path = best_model + # use the same config + if config: + # Fixes https://github.com/coqui-ai/TTS/issues/3131 + # Dirty, but they are closed + with open(args.config_path, 'r') as file: + data = json.load(file) + audio_data = data["audio"] + model_data = data["model_args"] + + model_args = {} + for key, value in config["model_args"].items(): + model_args[key] = model_data.get(key, value) + + audio_args = {} + for key, value in config["audio"].items(): + audio_args[key] = audio_data.get(key, value) + + config_args = {} + for key, value in config.items(): + config_args[key] = data.get(key, value) + + config = GPTTrainerConfig(**config_args) + config["model_args"] = GPTArgs(**model_args) + config["audio"] = XttsAudioConfig(**audio_args) + else: + coqpit = Coqpit() + coqpit.load_json(args.config_path) + + # override config values from command-line args + # TODO: Maybe it is better to do it outside + if len(coqpit_overrides) > 0: + config.parse_known_args(coqpit_overrides, relaxed_parser=True) + + # update the config.json fields and copy it to the output folder + new_fields = {} + if args.rank == 0: + if args.restore_path: + new_fields["restore_path"] = args.restore_path + new_fields["github_branch"] = get_git_branch() + return config, new_fields + + @staticmethod + def setup_training_environment(args, config, gpu): + if platform.system() != "Windows": + # https://github.com/pytorch/pytorch/issues/973 + import resource # pylint: disable=import-outside-toplevel + + rlimit = resource.getrlimit(resource.RLIMIT_NOFILE) + resource.setrlimit(resource.RLIMIT_NOFILE, (4096, rlimit[1])) + + # set and initialize Pytorch runtime + use_cuda, num_gpus = setup_torch_training_env( + args=args, + cudnn_enable=config.cudnn_enable, + cudnn_deterministic=config.cudnn_deterministic, + cudnn_benchmark=config.cudnn_benchmark, + use_ddp=args.use_ddp, + training_seed=config.training_seed, + allow_tf32=config.allow_tf32, + gpu=gpu if args.gpu is None else args.gpu, + ) + + print_training_env(args, config) + return use_cuda, num_gpus + + @staticmethod + def run_get_model(config: Coqpit, get_model: Callable) -> nn.Module: + """Run the `get_model` function and return the model. + + Args: + config (Coqpit): Model config. + + Returns: + nn.Module: initialized model. + """ + if len(signature(get_model).sig.parameters) == 1: + model = get_model(config) + else: + model = get_model() + return model + + @staticmethod + def run_get_data_samples(config: Coqpit, get_data_samples: Callable) -> nn.Module: + if callable(get_data_samples): + if len(signature(get_data_samples).sig.parameters) == 1: + train_samples, eval_samples = get_data_samples(config) + else: + train_samples, eval_samples = get_data_samples() + return train_samples, eval_samples + return None, None + + def restore_model( + self, + config: Coqpit, + restore_path: str, + model: nn.Module, + optimizer: torch.optim.Optimizer, + scaler: torch.cuda.amp.GradScaler = None, + ) -> Tuple[nn.Module, torch.optim.Optimizer, torch.cuda.amp.GradScaler, int]: + """Restore training from an old run. It restores model, optimizer, AMP scaler and training stats. + + Args: + config (Coqpit): Model config. + restore_path (str): Path to the restored training run. + model (nn.Module): Model to restored. + optimizer (torch.optim.Optimizer): Optimizer to restore. + scaler (torch.cuda.amp.GradScaler, optional): AMP scaler to restore. Defaults to None. + + Returns: + Tuple[nn.Module, torch.optim.Optimizer, torch.cuda.amp.GradScaler, int]: [description] + """ + + def _restore_list_objs(states, obj): + if isinstance(obj, list): + for idx, state in enumerate(states): + obj[idx].load_state_dict(state) + elif isinstance(obj, dict): + for key, state in states.items(): + obj[key].load_state_dict(state) + else: + obj.load_state_dict(states) + return obj + + logger.info(" > Restoring from %s ...", os.path.basename(restore_path)) + checkpoint = load_fsspec(restore_path, map_location="cpu") + + try: + logger.info(" > Restoring Model...") + model.load_state_dict(checkpoint["model"]) + logger.info(" > Restoring Optimizer...") + try: + optimizer = _restore_list_objs(checkpoint["optimizer"], optimizer) + except (KeyError, TypeError, RuntimeError): + logger.info(" > Optimizer is not compatible with the restored model.") + if "scaler" in checkpoint and self.use_amp_scaler and checkpoint["scaler"]: + logger.info(" > Restoring Scaler...") + scaler = _restore_list_objs(checkpoint["scaler"], scaler) + except (KeyError, RuntimeError, ValueError): + logger.info(" > Partial model initialization...") + model_dict = model.state_dict() + model_dict = set_partial_state_dict(model_dict, checkpoint["model"], config) + model.load_state_dict(model_dict) + del model_dict + + optimizer = self.restore_lr(config, self.args, model, optimizer) + + logger.info(" > Model restored from step %i", checkpoint["step"]) + restore_step = checkpoint["step"] + 1 # +1 not to immediately checkpoint if the model is restored + restore_epoch = checkpoint["epoch"] + torch.cuda.empty_cache() + return model, optimizer, scaler, restore_step, restore_epoch + + def restore_lr(self, config, args, model, optimizer): + # use the same lr if continue training + if not args.continue_path: + if isinstance(optimizer, list): + for idx, optim in enumerate(optimizer): + for group in optim.param_groups: + group["lr"] = self.get_lr(model, config)[idx] + elif isinstance(optimizer, dict): + for optim_name, optim in optimizer.items(): + for group in optim.param_groups: + group["lr"] = self.get_lr(model, config)[optim_name] + else: + for group in optimizer.param_groups: + group["lr"] = self.get_lr(model, config) + return optimizer + + ######################### + # DATA LOADING FUNCTIONS + ######################### + + def _get_loader( + self, + model: nn.Module, + config: Coqpit, + assets: Dict, + is_eval: str, + samples: List, + verbose: bool, + num_gpus: int, + ) -> DataLoader: + if num_gpus > 1: + if isimplemented(model.module, "get_data_loader"): + loader = model.module.get_data_loader( + config, + assets, + is_eval, + samples, + verbose, + num_gpus, + self.args.rank, + ) + else: + if isimplemented(model, "get_data_loader"): + loader = model.get_data_loader( + config=config, assets=assets, is_eval=is_eval, samples=samples, verbose=verbose, num_gpus=num_gpus + ) + + assert ( + len(loader) > 0 + ), " ❗ len(DataLoader) returns 0. Make sure your dataset is not empty or len(dataset) > 0. " + return loader + + def get_train_dataloader(self, training_assets: Dict, samples: List, verbose: bool) -> DataLoader: + """Initialize and return a training data loader. + Call ```model.get_train_data_loader``` if it is implemented, else call ```model.get_data_loader``` + and set ```is_eval=False```. + + Args: + ap (AudioProcessor): Audio processor. + samples (List): Data samples used for training. + verbose (bool): enable/disable printing loader stats at initialization. + + Returns: + DataLoader: Initialized training data loader. + """ + if self.num_gpus > 1: + if isimplemented(self.model.module, "get_train_data_loader"): + loader = self.model.module.get_train_data_loader( + self.config, + self.training_assets, + samples, + verbose, + self.num_gpus, + self.args.rank, + ) + return loader + else: + if isimplemented(self.model, "get_train_data_loader"): + loader = self.model.get_train_data_loader( + self.config, self.training_assets, samples, verbose, self.num_gpus + ) + return loader + + return self._get_loader( + self.model, + self.config, + training_assets, + False, + samples, + verbose, + self.num_gpus, + ) + + def get_eval_dataloader(self, training_assets: Dict, samples: List, verbose: bool) -> DataLoader: + """Initialize and return a evaluation data loader. + Call ```model.get_eval_data_loader``` if it is implemented, else call ```model.get_data_loader``` + and set ```is_eval=True```. + + Args: + ap (AudioProcessor): Audio processor. + samples (List): Data samples used for training. + verbose (bool): enable/disable printing loader stats at initialization. + + Returns: + DataLoader: Initialized training data loader. + """ + if self.num_gpus > 1: + if isimplemented(self.model.module, "get_eval_data_loader"): + loader = self.model.module.get_eval_data_loader( + self.config, + self.training_assets, + samples, + verbose, + self.num_gpus, + self.args.rank, + ) + return loader + else: + if isimplemented(self.model, "get_eval_data_loader"): + loader = self.model.get_eval_data_loader( + self.config, self.training_assets, samples, verbose, self.num_gpus + ) + return loader + + return self._get_loader( + self.model, + self.config, + training_assets, + True, + samples, + verbose, + self.num_gpus, + ) + + def get_test_dataloader(self, training_assets: Dict, samples: List, verbose: bool) -> DataLoader: + """Initialize and return a evaluation data loader. + Call ```model.get_test_data_loader``` if it is implemented, else call ```model.get_data_loader``` + and set ```is_eval=True```. + + Args: + ap (AudioProcessor): Audio processor. + samples (List): Data samples used for training. + verbose (bool): enable/disable printing loader stats at initialization. + + Returns: + DataLoader: Initialized training data loader. + """ + if self.num_gpus > 1: + if isimplemented(self.model.module, "get_test_data_loader"): + loader = self.model.module.get_test_data_loader( + self.config, + self.training_assets, + samples, + verbose, + self.num_gpus, + self.args.rank, + ) + return loader + else: + if isimplemented(self.model, "get_test_data_loader"): + loader = self.model.get_test_data_loader( + self.config, self.training_assets, samples, verbose, self.num_gpus + ) + return loader + + return self._get_loader( + self.model, + self.config, + training_assets, + True, + samples, + verbose, + self.num_gpus, + ) + + def format_batch(self, batch: List) -> Dict: + """Format the dataloader output and return a batch. + + 1. Call ```model.format_batch```. + 2. Pass the batch to the Device. + 3. Call ```model.format_batch_on_device```. + + Args: + batch (List): Batch returned by the dataloader. + + Returns: + Dict: Formatted batch. + """ + try: + if self.num_gpus > 1: + batch = self.model.module.format_batch(batch) + else: + batch = self.model.format_batch(batch) + except NotImplementedError: + pass + + if isinstance(batch, dict): + for k, v in batch.items(): + batch[k] = to_cuda(v) + elif isinstance(batch, list): + batch = [to_cuda(v) for v in batch] + + try: + if self.num_gpus > 1: + batch = self.model.module.format_batch_on_device(batch) + else: + batch = self.model.format_batch_on_device(batch) + except NotImplementedError: + pass + return batch + + ###################### + # TRAIN FUNCTIONS + ###################### + + @staticmethod + def master_params(optimizer: torch.optim.Optimizer): + """Generator over parameters owned by the optimizer. + + Used to select parameters used by the optimizer for gradient clipping. + + Args: + optimizer: Target optimizer. + """ + for group in optimizer.param_groups: + for p in group["params"]: + yield p + + @staticmethod + def _model_train_step( + batch: Dict, model: nn.Module, criterion: nn.Module, optimizer_idx: int = None + ) -> Tuple[Dict, Dict]: + """ + Perform a trainig forward step. Compute model outputs and losses. + + Args: + batch (Dict): [description] + model (nn.Module): [description] + criterion (nn.Module): [description] + optimizer_idx (int, optional): [description]. Defaults to None. + + Returns: + Tuple[Dict, Dict]: [description] + """ + input_args = [batch, criterion] + if optimizer_idx is not None: + input_args.append(optimizer_idx) + # unwrap model in DDP training + if hasattr(model, "module"): + return model.module.train_step(*input_args) + return model.train_step(*input_args) + + def _get_autocast_args(self, mixed_precision: bool, precision: str): + device = "cpu" + dtype = torch.get_autocast_cpu_dtype() + if self.use_cuda: + device = "cuda" + dtype = torch.float32 + if mixed_precision: + if precision == "fp16": + dtype = torch.float16 + elif precision == "bf16": + dtype = torch.bfloat16 + else: + raise ValueError(f" ❗ Unknown precision {precision}") + elif mixed_precision: + dtype = torch.bfloat16 + return device, dtype + + def detach_loss_dict( + self, loss_dict: Dict, step_optimizer: bool, optimizer_idx: int = None, grad_norm: float = None + ): + # detach losses for logging + loss_dict_detached = self._detach_loss_dict(loss_dict) + # loss_dict_detached["loss"] = loss_di`ct_detached["loss"] * float(self.grad_accum_steps) + + if optimizer_idx is not None: + loss_dict_detached[f"loss_{optimizer_idx}"] = loss_dict_detached.pop("loss") + if step_optimizer and grad_norm is not None: + loss_dict_detached[f"grad_norm_{optimizer_idx}"] = grad_norm + else: + if step_optimizer and grad_norm is not None: + loss_dict_detached["grad_norm"] = grad_norm + return loss_dict_detached + + def _compute_loss(self, batch: Dict, model: nn.Module, criterion: nn.Module, config: Coqpit, optimizer_idx: int): + device, dtype = self._get_autocast_args(config.mixed_precision, config.precision) + with torch.autocast(device_type=device, dtype=dtype, enabled=config.mixed_precision): + if optimizer_idx is not None: + outputs, loss_dict = self._model_train_step(batch, model, criterion, optimizer_idx=optimizer_idx) + else: + outputs, loss_dict = self._model_train_step(batch, model, criterion) + return outputs, loss_dict + + @staticmethod + def _set_grad_clip_per_optimizer(config: Coqpit, optimizer_idx: int): + # set gradient clipping threshold + grad_clip = 0.0 # meaning no gradient clipping + if "grad_clip" in config and config.grad_clip is not None: + if optimizer_idx is not None: + try: + grad_clip = config.grad_clip[optimizer_idx] + except TypeError: + logger.info(" [!] You are using multiple optimizers but `grad_clip` is not a list.") + else: + grad_clip = config.grad_clip + return grad_clip + + def _compute_grad_norm(self, optimizer: torch.optim.Optimizer): + return torch.norm(torch.cat([param.grad.view(-1) for param in self.master_params(optimizer)], dim=0), p=2) + + def _grad_clipping(self, grad_clip: float, optimizer: torch.optim.Optimizer, scaler: "AMPScaler"): + """Perform gradient clipping""" + if grad_clip is not None and grad_clip > 0: + if scaler: + scaler.unscale_(optimizer) + self.callbacks.before_gradient_clipping(self) + grad_norm = torch.nn.utils.clip_grad_norm_(self.master_params(optimizer), grad_clip) + else: + grad_norm = self._compute_grad_norm(optimizer) + return grad_norm + + def optimize( + self, + batch: Dict, + model: nn.Module, + optimizer: torch.optim.Optimizer, + scaler: "AMPScaler", + criterion: nn.Module, + scheduler: Union[torch.optim.lr_scheduler._LRScheduler, List, Dict], # pylint: disable=protected-access + config: Coqpit, + optimizer_idx: int = None, + step_optimizer: bool = True, + num_optimizers: int = 1, + ) -> Tuple[Dict, Dict, int]: + """Perform a forward - backward pass and run the optimizer. + + Args: + batch (Dict): Input batch. If + model (nn.Module): Model for training. Defaults to None. + optimizer (Union[nn.optim.Optimizer, List]): Model's optimizer. If it is a list then, `optimizer_idx` must be defined to indicate the optimizer in use. + scaler (AMPScaler): AMP scaler. + criterion (nn.Module): Model's criterion. + scheduler (torch.optim.lr_scheduler._LRScheduler): LR scheduler used by the optimizer. + config (Coqpit): Model config. + optimizer_idx (int, optional): Target optimizer being used. Defaults to None. + step_optimizer (bool, optional): Whether step the optimizer. If False, gradients are accumulated and + model parameters are not updated. Defaults to True. + num_optimizers (int, optional): Number of optimizers. Defaults to 1. + + Raises: + RuntimeError: When the loss is NaN. + + Returns: + Tuple[Dict, Dict, int, torch.Tensor]: model outputs, losses, step time and gradient norm. + """ + + step_start_time = time.time() + + # forward pass and loss computation + outputs, loss_dict = self._compute_loss( + batch=batch, model=model, criterion=criterion, config=config, optimizer_idx=optimizer_idx + ) + + # skip the rest if not outputs from the model + if not loss_dict: + step_time = time.time() - step_start_time + return outputs, {}, step_time + + grad_clip = self._set_grad_clip_per_optimizer(config=config, optimizer_idx=optimizer_idx) + # optimizer step + grad_norm = 0 + update_lr_scheduler = True + + # callback + self.callbacks.before_backward_pass(self, loss_dict) + + # accumulated gradients adjustment + loss_dict["loss"] = loss_dict["loss"] / float(self.grad_accum_steps) + + if self.use_accelerate: + with self.accelerator.accumulate(model): + ctx_mgr = self.accelerator.autocast if config.mixed_precision else nullcontext + with ctx_mgr(): + self.accelerator.backward(loss_dict["loss"]) + grad_norm = self._compute_grad_norm(optimizer) + if self.accelerator.sync_gradients and grad_clip is not None and grad_clip > 0: + self.accelerator.clip_grad_norm_(model.parameters(), grad_clip) + optimizer.step() + if not self.config.scheduler_after_epoch and not self.accelerator.optimizer_step_was_skipped: + scheduler.step() + optimizer.zero_grad(set_to_none=True) + else: + if self.use_amp_scaler: + if self.use_apex: + # TODO: verify AMP use for GAN training in TTS + # https://nvidia.github.io/apex/advanced.html?highlight=accumulate#backward-passes-with-multiple-optimizers + with amp.scale_loss(loss_dict["loss"], optimizer) as scaled_loss: + scaled_loss.backward() + if step_optimizer: + grad_norm = self._grad_clipping(grad_clip=grad_clip, optimizer=optimizer, scaler=None) + else: + # model optimizer step in mixed precision mode + scaler.scale(loss_dict["loss"]).backward() + # gradient accumulation + if step_optimizer: + grad_norm = self._grad_clipping(grad_clip=grad_clip, optimizer=optimizer, scaler=scaler) + scale_prev = scaler.get_scale() + scaler.step(optimizer) + # update the scaler at the end of all the optimizer steps + if optimizer_idx is None or (optimizer_idx + 1 == num_optimizers): + scaler.update() + loss_dict["amp_scaler"] = scaler.get_scale() # for logging + update_lr_scheduler = scale_prev <= scaler.get_scale() + else: + # main model optimizer step + loss_dict["loss"].backward() + # gradient accumulation + if step_optimizer: + self.callbacks.before_gradient_clipping(self) + if grad_clip > 0: + grad_norm = torch.nn.utils.clip_grad_norm_(self.master_params(optimizer), grad_clip) + optimizer.step() + + # setup lr + if ( + scheduler is not None + and update_lr_scheduler + and not self.config.scheduler_after_epoch + and step_optimizer + ): + scheduler.step() + + # zero-out optimizer + if step_optimizer: + optimizer.zero_grad(set_to_none=True) + + # pytorch skips the step when the norm is 0. So ignore the norm value when it is NaN + if isinstance(grad_norm, torch.Tensor) and (torch.isnan(grad_norm) or torch.isinf(grad_norm)): + grad_norm = 0 + + step_time = time.time() - step_start_time + + # detach loss dict + loss_dict_detached = self.detach_loss_dict(loss_dict, step_optimizer, optimizer_idx, grad_norm) + return outputs, loss_dict_detached, step_time + + def train_step(self, batch: Dict, batch_n_steps: int, step: int, loader_start_time: float) -> Tuple[Dict, Dict]: + """Perform a training step on a batch of inputs and log the process. + + Args: + batch (Dict): Input batch. + batch_n_steps (int): Number of steps needed to complete an epoch. Needed for logging. + step (int): Current step number in this epoch. + loader_start_time (float): The time when the data loading is started. Needed for logging. + + Returns: + Tuple[Dict, Dict]: Model outputs and losses. + """ + self.callbacks.on_train_step_start(self) + # format data + batch = self.format_batch(batch) + loader_time = time.time() - loader_start_time + + # conteainers to hold model outputs and losses for each optimizer. + outputs_per_optimizer = None + loss_dict = {} + + # OPTIMIZATION + if isimplemented(self.model, "optimize"): # pylint: disable=too-many-nested-blocks + # custom optimize for the model + step_time = time.time() + device, dtype = self._get_autocast_args(self.config.mixed_precision, self.config.precision) + with torch.autocast(device_type=device, dtype=dtype, enabled=self.config.mixed_precision): + outputs, loss_dict_new = self.model.optimize( + batch, + self, + ) + step_time = time.time() - step_time + # If None, skip the step + if outputs is None: + return None, None + # TODO: find a way to log grad_norm for custom optimize + loss_dict_new = self.detach_loss_dict(loss_dict_new, True, None, None) + loss_dict.update(loss_dict_new) + else: + # gradient accumulation + # TODO: grad accumulation for each optimizer + step_optimizer = True + if ((step + 1) % self.grad_accum_steps != 0) and (step + 1 != batch_n_steps): + step_optimizer = False + + if not isinstance(self.optimizer, list): + # auto training with a single optimizer + outputs, loss_dict_new, step_time = self.optimize( + batch, + self.model, + self.optimizer, + self.scaler, + self.criterion, + self.scheduler, + self.config, + step_optimizer=step_optimizer, + num_optimizers=1, + ) + loss_dict.update(loss_dict_new) + else: + # auto training with multiple optimizers (e.g. GAN) + outputs_per_optimizer = [None] * len(self.optimizer) + total_step_time = 0 + for idx, optimizer in enumerate(self.optimizer): + criterion = self.criterion + # scaler = self.scaler[idx] if self.use_amp_scaler else None + scaler = self.scaler + scheduler = None + if self.scheduler is not None: + scheduler = self.scheduler[idx] + outputs, loss_dict_new, step_time = self.optimize( + batch, + self.model, + optimizer, + scaler, + criterion, + scheduler, + self.config, + idx, + step_optimizer=step_optimizer, + num_optimizers=len(self.optimizer), + ) + # skip the rest if the model returns None + total_step_time += step_time + outputs_per_optimizer[idx] = outputs + # merge loss_dicts from each optimizer + # rename duplicates with the optimizer idx + # if None, model skipped this optimizer + if loss_dict_new is not None: + for k, v in loss_dict_new.items(): + if k in loss_dict: + loss_dict[f"{k}-{idx}"] = v + else: + loss_dict[k] = v + step_time = total_step_time + + outputs = outputs_per_optimizer + + # clear any pesky gradients after gradient accumulation + if step_optimizer: + self.model.zero_grad(set_to_none=True) + + # update avg runtime stats + keep_avg_update = {} + keep_avg_update["avg_loader_time"] = loader_time + keep_avg_update["avg_step_time"] = step_time + self.keep_avg_train.update_values(keep_avg_update) + + # update avg loss stats + update_eval_values = {} + for key, value in loss_dict.items(): + update_eval_values["avg_" + key] = value + self.keep_avg_train.update_values(update_eval_values) + + # print training progress + if self.total_steps_done % self.config.print_step == 0: + # log learning rates + lrs = {} + if isinstance(self.optimizer, list): + for idx, optimizer in enumerate(self.optimizer): + current_lr = self.optimizer[idx].param_groups[0]["lr"] + lrs.update({f"current_lr_{idx}": current_lr}) + elif isinstance(self.optimizer, dict): + for key, optimizer in self.optimizer.items(): + current_lr = self.optimizer[key].param_groups[0]["lr"] + lrs.update({f"current_lr_{key}": current_lr}) + else: + current_lr = self.optimizer.param_groups[0]["lr"] + lrs = {"current_lr": current_lr} + + # log run-time stats + loss_dict.update(lrs) + loss_dict.update( + { + "step_time": round(step_time, 4), + "loader_time": round(loader_time, 4), + } + ) + self.c_logger.print_train_step( + batch_n_steps, + step, + self.total_steps_done, + loss_dict, + self.keep_avg_train.avg_values, + ) + + if self.args.rank == 0: + # Plot Training Iter Stats + # reduce TB load and don't log every step + if self.total_steps_done % self.config.plot_step == 0: + self.dashboard_logger.train_step_stats(self.total_steps_done, loss_dict) + if self.total_steps_done % self.config.save_step == 0 and self.total_steps_done != 0: + if self.config.save_checkpoints: + # checkpoint the model + self.save_checkpoint() + + if self.total_steps_done % self.config.log_model_step == 0: + # log checkpoint as artifact + self.update_training_dashboard_logger(batch=batch, outputs=outputs) + + self.dashboard_logger.flush() + + self.total_steps_done += 1 + self.callbacks.on_train_step_end(self) + return outputs, loss_dict + + def train_epoch(self) -> None: + """Main entry point for the training loop. Run training on the all training samples.""" + # initialize the data loader + if self.train_loader is None: + self.train_loader = self.get_train_dataloader( + self.training_assets, + self.train_samples, + verbose=True, + ) + self.train_loader = self.prepare_accelerate_loader(self.train_loader) + # set model to training mode + torch.set_grad_enabled(True) + if self.num_gpus > 1: + self.model.module.train() + else: + self.model.train() + epoch_start_time = time.time() + + self.callbacks.on_train_epoch_start(self) + + self.c_logger.print_train_start() + loader_start_time = time.time() + # TRAINING EPOCH -> iterate over the training samples + batch_num_steps = len(self.train_loader) + for cur_step, batch in enumerate(self.train_loader): + outputs, _ = self.train_step(batch, batch_num_steps, cur_step, loader_start_time) + if outputs is None: + logger.info(" [!] `train_step()` retuned `None` outputs. Skipping training step.") + continue + del outputs + loader_start_time = time.time() + + # RUN EVAL -> run evaluation epoch in the middle of training. Useful for big datasets. + if self.config.run_eval_steps is not None and (self.total_steps_done % self.config.run_eval_steps == 0): + self.eval_epoch() + if self.num_gpus > 1: + self.model.module.train() + else: + self.model.train() + torch.set_grad_enabled(True) + + epoch_time = time.time() - epoch_start_time + self.callbacks.on_train_epoch_end(self) + + # scheduler step + if self.scheduler is not None and self.config.scheduler_after_epoch: + if isinstance(self.scheduler, list): + for scheduler in self.scheduler: + if scheduler is not None: + if isinstance(scheduler, ReduceLROnPlateau): + scheduler.step(self._pick_target_avg_loss(self.keep_avg_eval) if not None else 999) + else: + scheduler.step() + elif isinstance(self.scheduler, dict): # only with `model.optimize()`` + for scheduler in self.scheduler.values(): + if scheduler is not None: + scheduler.step() + elif isinstance(self.scheduler, ReduceLROnPlateau): + self.scheduler.step(self._pick_target_avg_loss(self.keep_avg_eval) if not None else 999) + else: + self.scheduler.step() + # plot self.epochs_done Stats + if self.args.rank == 0: + epoch_stats = {"epoch_time": epoch_time} + epoch_stats.update(self.keep_avg_train.avg_values) + self.dashboard_logger.train_epoch_stats(self.total_steps_done, epoch_stats) + if self.config.model_param_stats: + self.dashboard_logger.model_weights(self.model, self.total_steps_done) + torch.cuda.empty_cache() + + ####################### + # EVAL FUNCTIONS + ####################### + + def _model_eval_step( + self, batch: Dict, model: nn.Module, criterion: nn.Module, optimizer_idx: int = None + ) -> Tuple[Dict, Dict]: + """ + Perform a evaluation forward pass. Compute model outputs and losses with no gradients. + + Args: + batch (Dict): IBatch of inputs. + model (nn.Module): Model to call evaluation. + criterion (nn.Module): Model criterion. + optimizer_idx (int, optional): Optimizer ID to define the closure in multi-optimizer training. Defaults to None. + + Returns: + Tuple[Dict, Dict]: model outputs and losses. + """ + input_args = [batch, criterion] + + if isimplemented(model, "optimize"): + if hasattr(model, "module"): + return model.module.eval_step(batch, self) + return model.eval_step(batch, self) + + if optimizer_idx is not None: + input_args.append(optimizer_idx) + if hasattr(model, "module"): + return model.module.eval_step(*input_args) + + return model.eval_step(*input_args) + + def eval_step(self, batch: Dict, step: int) -> Tuple[Dict, Dict]: + """Perform a evaluation step on a batch of inputs and log the process. + + Args: + batch (Dict): Input batch. + step (int): Current step number in this epoch. + + Returns: + Tuple[Dict, Dict]: Model outputs and losses. + """ + with torch.no_grad(): + outputs = [] + loss_dict = {} + if not isinstance(self.optimizer, list) or isimplemented(self.model, "optimize"): + outputs, loss_dict = self._model_eval_step(batch, self.model, self.criterion) + if outputs is None: + return None, None + else: + outputs = [None] * len(self.optimizer) + for idx, _ in enumerate(self.optimizer): + criterion = self.criterion + outputs_, loss_dict_new = self._model_eval_step(batch, self.model, criterion, idx) + if outputs_ is None: + return None, None + outputs[idx] = outputs_ + + if loss_dict_new: + loss_dict_new[f"loss_{idx}"] = loss_dict_new.pop("loss") + loss_dict.update(loss_dict_new) + + loss_dict = self._detach_loss_dict(loss_dict) + + # update avg stats + update_eval_values = {} + for key, value in loss_dict.items(): + update_eval_values["avg_" + key] = value + self.keep_avg_eval.update_values(update_eval_values) + if self.config.print_eval: + self.c_logger.print_eval_step(step, loss_dict, self.keep_avg_eval.avg_values) + + return outputs, loss_dict + + def eval_epoch(self) -> None: + """Main entry point for the evaluation loop. Run evaluation on the all validation samples.""" + + # initialize it when eval_epoch is called alone. + self.keep_avg_eval = KeepAverage() if self.keep_avg_eval is None else self.keep_avg_eval + + if self.eval_loader is None: + self.eval_loader = ( + self.get_eval_dataloader( + self.training_assets, + self.eval_samples, + verbose=True, + ) + if self.config.run_eval + else None + ) + + torch.set_grad_enabled(False) + self.model.eval() + self.c_logger.print_eval_start() + loader_start_time = time.time() + batch = None + outputs = None + for cur_step, batch in enumerate(self.eval_loader): + # format data + batch = self.format_batch(batch) + loader_time = time.time() - loader_start_time + self.keep_avg_eval.update_values({"avg_loader_time": loader_time}) + outputs_, _ = self.eval_step(batch, cur_step) + if outputs_ is None: + logger.info(" [!] `eval_step()` retuned `None` outputs. Skipping evaluation step.") + continue + outputs = outputs_ + loader_start_time = time.time() + # plot epoch stats, artifacts and figures + if self.args.rank == 0 and outputs is not None: + if hasattr(self.model, "module") and isimplemented(self.model.module, "eval_log"): + self.model.module.eval_log( + batch, + outputs, + self.dashboard_logger, + self.training_assets, + self.total_steps_done, + ) + elif isimplemented(self.model, "eval_log"): + self.model.eval_log( + batch, + outputs, + self.dashboard_logger, + self.training_assets, + self.total_steps_done, + ) + self.dashboard_logger.eval_stats(self.total_steps_done, self.keep_avg_eval.avg_values) + torch.cuda.empty_cache() + + ################################## + # TESTING + ################################## + def test_run(self) -> None: + """Run model test. + + Test run is expected to pass over test samples and produce logging artifacts. + + If ```model.test_run()``` is defined, it will be called and it is expected to set and execute everything + in the model. + + Else if ```mode.test()``` is defined, it will be called and it takes an test data loader as an argument + and iterate over it. + """ + self.model.eval() + test_outputs = None + if isimplemented(self.model, "test_run") or ( + self.num_gpus > 1 and isimplemented(self.model.module, "test_run") + ): + # handle everything in ```model.test_run()` + if self.num_gpus > 1: + test_outputs = self.model.module.test_run(self.training_assets) + else: + test_outputs = self.model.test_run(self.training_assets) + elif isimplemented(self.model, "test") or (self.num_gpus > 1 and isimplemented(self.model.module, "test")): + self.test_loader = self.get_test_dataloader( + self.training_assets, + self.test_samples if self.test_samples else self.eval_samples, + verbose=True, + ) + # use test_loader to load test samples + if self.num_gpus > 1: + test_outputs = self.model.module.test(self.training_assets, self.test_loader, None) + else: + test_outputs = self.model.test(self.training_assets, self.test_loader, None) + if isimplemented(self.model, "test_log") or ( + self.num_gpus > 1 and isimplemented(self.model.module, "test_log") + ): + if self.num_gpus > 1: + self.model.module.test_log( + test_outputs, self.dashboard_logger, self.training_assets, self.total_steps_done + ) + else: + self.model.test_log(test_outputs, self.dashboard_logger, self.training_assets, self.total_steps_done) + + def _restore_best_loss(self): + """Restore the best loss from the args.best_path if provided else + from the model (`args.continue_path`) used for resuming the training""" + if self.args.continue_path and (self.restore_step != 0 or self.args.best_path): + logger.info(" > Restoring best loss from %s ...", os.path.basename(self.args.best_path)) + ch = load_fsspec(self.args.restore_path, map_location="cpu") + if "model_loss" in ch: + if isinstance(ch["model_loss"], dict): + self.best_loss = ch["model_loss"] + # For backwards-compatibility: + elif isinstance(ch["model_loss"], float): + if self.config.run_eval: + self.best_loss = {"train_loss": None, "eval_loss": ch["model_loss"]} + else: + self.best_loss = {"train_loss": ch["model_loss"], "eval_loss": None} + logger.info(" > Starting with loaded last best loss %s", self.best_loss) + + def test(self, model=None, test_samples=None) -> None: + """Run evaluation steps on the test data split. You can either provide the model and the test samples + explicitly or the trainer use values from the initialization. + + Args: + model (nn.Module, optional): Model to use for testing. If None, use the model given in the initialization. + Defaults to None. + + test_samples (List[str], optional): List of test samples to use for testing. If None, use the test samples + given in the initialization. Defaults to None. + """ + + logger.info(" > USING TEST SET...") + self.keep_avg_eval = KeepAverage() + + if model is not None: + self.model = model + + eval_samples_cache = self.eval_samples + if test_samples is not None: + self.eval_samples = test_samples + else: + self.eval_samples = self.test_samples + + self.eval_epoch() + self.c_logger.print_epoch_end(self.epochs_done, self.keep_avg_eval.avg_values) + self.eval_samples = eval_samples_cache + + ################################### + # FIT FUNCTIONS + ################################### + + def _fit(self) -> None: + """πŸƒ train -> evaluate -> test for the number of epochs.""" + self._restore_best_loss() + + self.total_steps_done = self.restore_step + + for epoch in range(0, self.config.epochs): + if self.num_gpus > 1: + # let all processes sync up before starting with a new epoch of training + dist.barrier() + self.callbacks.on_epoch_start(self) + self.keep_avg_train = KeepAverage() + self.epochs_done = epoch + self.c_logger.print_epoch_start(epoch, self.config.epochs, self.output_path) + if not self.skip_train_epoch and not self.start_with_eval: + self.train_epoch() + if self.config.run_eval: + self.keep_avg_eval = KeepAverage() + self.eval_epoch() + if epoch >= self.config.test_delay_epochs and self.args.rank <= 0: + self.test_run() + + self.c_logger.print_epoch_end( + epoch, + self.keep_avg_eval.avg_values if self.config.run_eval else self.keep_avg_train.avg_values, + ) + if self.args.rank in [None, 0]: + self.save_best_model() + self.callbacks.on_epoch_end(self) + self.start_with_eval = False + + def fit_with_largest_batch_size(self, starting_batch_size=2048) -> None: + cuda_meminfo() + bs = starting_batch_size + while True: + gc.collect() + torch.cuda.empty_cache() + try: + gc.collect() + torch.cuda.empty_cache() + self.config.batch_size = bs + logger.info(" > current batch size: %i", self.config.batch_size) + self._fit() + except RuntimeError as exception: + if bs > 1 and should_reduce_batch_size(exception): + bs //= 2 + gc.collect() + torch.cuda.empty_cache() + else: + raise + except Exception as exception: # pylint: disable=broad-except + # catches the torch.cuda.OutOfMemoryError + if bs > 1 and should_reduce_batch_size(exception): + bs //= 2 + gc.collect() + torch.cuda.empty_cache() + else: + raise + else: + break + + def fit(self) -> None: + """Where the ✨️magic✨️ happens...""" + try: + self._fit() + if self.args.rank == 0: + self.dashboard_logger.finish() + except KeyboardInterrupt: + logger.info(" > Keyboard interrupt detected.") + if self.config.save_on_interrupt: + logger.info(" > Saving model before exiting...") + # save the model on keyboard interrupt + self.save_checkpoint() + # update the training dashboard logger + self.update_training_dashboard_logger() + # call the keyboard interrupt callback + self.callbacks.on_keyboard_interrupt(self) + # if the output folder is empty remove the run. + remove_experiment_folder(self.output_path) + # clear the DDP processes + if self.num_gpus > 1: + dist.destroy_process_group() + # finish the wandb run and sync data + if self.args.rank == 0: + self.dashboard_logger.finish() + # stop without error signal + try: + sys.exit(1) + except SystemExit: + os._exit(1) # pylint: disable=protected-access + except BaseException: # pylint: disable=broad-except + remove_experiment_folder(self.output_path) + traceback.print_exc() + sys.exit(1) + + def profile_fit(self, torch_profiler, epochs=None, small_run=None): + """Run training under the torch profiler. + + Example:: + Run torch profiler to profile CPU, GPU and memory usage with Tensorboard logging. + + >>> import torch + >>> profiler = torch.profiler.profile( + >>> activities=[ + >>> torch.profiler.ProfilerActivity.CPU, + >>> torch.profiler.ProfilerActivity.CUDA, + >>> ], + >>> schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=2), + >>> on_trace_ready=torch.profiler.tensorboard_trace_handler("./profiler/"), + >>> record_shapes=True, + >>> profile_memory=True, + >>> with_stack=True, + >>> ) + >>> prof = trainer.profile_fit(profiler, epochs=1, small_run=64) + """ + self.dashboard_logger = DummyLogger() + # train the model for a custom number of epochs + if epochs: + self.config.epocshs = epochs + # use a smaller set of training samples for profiling + if small_run: + self.setup_small_run(small_run) + # run profiler + self.config.run_eval = False + self.config.test_delay_epochs = 9999999 + self.config.epochs = epochs + # set a callback to progress the profiler + self.callbacks_on_train_step_end = [ # pylint: disable=attribute-defined-outside-init + lambda trainer: trainer.torch_profiler.step() + ] + # set the profiler to access in the Trainer + self.torch_profiler = torch_profiler # pylint: disable=attribute-defined-outside-init + # set logger output for Tensorboard + # self.torch_profiler.on_trace_ready = torch.profiler.tensorboard_trace_handler(self.output_path) + self.torch_profiler.start() + self.fit() + self.torch_profiler.stop() + return self.torch_profiler + + @rank_zero_only + def save_best_model(self) -> None: + """Save the best model. It only saves if the current target loss is smaller then the previous.""" + + eval_loss = self._pick_target_avg_loss(self.keep_avg_eval) + train_loss = self._pick_target_avg_loss(self.keep_avg_train) + + # save the model and update the best_loss + self.best_loss = save_best_model( + {"train_loss": train_loss, "eval_loss": eval_loss}, + self.best_loss, + self.config, + self.model, + self.optimizer, + self.scaler if self.use_amp_scaler else None, + self.total_steps_done, + self.epochs_done, + self.output_path, + keep_all_best=self.config.save_all_best, + keep_after=self.config.save_best_after, + save_func=self.dashboard_logger.save_model, + ) + + @rank_zero_only + def save_checkpoint(self) -> None: + """Save the current model checkpoint.""" + eval_loss = self._pick_target_avg_loss(self.keep_avg_eval) + train_loss = self._pick_target_avg_loss(self.keep_avg_train) + + save_checkpoint( + self.config, + self.model, + self.optimizer, + self.scaler if self.use_amp_scaler else None, + self.total_steps_done, + self.epochs_done, + self.output_path, + model_loss={"train_loss": train_loss, "eval_loss": eval_loss}, + save_n_checkpoints=self.config.save_n_checkpoints, + save_func=self.dashboard_logger.save_model, + ) + + @rank_zero_only + def update_training_dashboard_logger(self, batch=None, outputs=None): + aliases = [ + f"epoch-{self.epochs_done}", + f"step-{self.total_steps_done}", + ] + self.dashboard_logger.add_artifact( + file_or_dir=self.output_path, name="checkpoint", artifact_type="model", aliases=aliases + ) + + # training visualizations + if batch is not None and outputs is not None: + if hasattr(self.model, "module") and isimplemented(self.model.module, "train_log"): + self.model.module.train_log( + batch, + outputs, + self.dashboard_logger, + self.training_assets, + self.total_steps_done, + ) + elif isimplemented(self.model, "train_log"): + self.model.train_log( + batch, + outputs, + self.dashboard_logger, + self.training_assets, + self.total_steps_done, + ) + + ##################### + # GET FUNCTIONS + ##################### + + @staticmethod + def get_optimizer(model: nn.Module, config: Coqpit) -> Union[torch.optim.Optimizer, List]: + """Receive the optimizer from the model if model implements `get_optimizer()` else + check the optimizer parameters in the config and try initiating the optimizer. + + Args: + model (nn.Module): Training model. + config (Coqpit): Training configuration. + + Returns: + Union[torch.optim.Optimizer, List]: A optimizer or a list of optimizers. GAN models define a list. + """ + optimizer = None + if isimplemented(model, "get_optimizer"): + try: + optimizer = model.get_optimizer() + except NotImplementedError: + optimizer = None + if optimizer is None: + optimizer_name = config.optimizer + optimizer_params = {} if config.optimizer_params is None else config.optimizer_params + return get_optimizer(optimizer_name, optimizer_params, config.lr, model) + return optimizer + + @staticmethod + def get_lr(model: nn.Module, config: Coqpit) -> Union[float, List[float]]: + """Set the initial learning rate by the model if model implements `get_lr()` else try setting the learning rate + fromthe config. + + Args: + model (nn.Module): Training model. + config (Coqpit): Training configuration. + + Returns: + Union[float, List[float]]: A single learning rate or a list of learning rates, one for each optimzier. + """ + lr = None + if isimplemented(model, "get_lr"): + try: + lr = model.get_lr() + except NotImplementedError: + lr = None + if lr is None: + lr = config.lr + return lr + + @staticmethod + def get_scheduler( + model: nn.Module, config: Coqpit, optimizer: Union[torch.optim.Optimizer, List, Dict] + ) -> Union[torch.optim.lr_scheduler._LRScheduler, List]: # pylint: disable=protected-access + """Receive the scheduler from the model if model implements `get_scheduler()` else + check the config and try initiating the scheduler. + + Args: + model (nn.Module): Training model. + config (Coqpit): Training configuration. + + Returns: + Union[torch.optim.Optimizer, List, Dict]: A scheduler or a list of schedulers, one for each optimizer. + """ + scheduler = None + if isimplemented(model, "get_scheduler"): + try: + scheduler = model.get_scheduler(optimizer) + except NotImplementedError: + scheduler = None + if isinstance(scheduler, dict) and not isimplemented(model, "optimize"): + raise ValueError( + " [!] Dictionary of schedulers are only supported with the manual optimization `model.optimize()`." + ) + if scheduler is None: + lr_scheduler = config.lr_scheduler + lr_scheduler_params = config.lr_scheduler_params + return get_scheduler(lr_scheduler, lr_scheduler_params, optimizer) + return scheduler + + @staticmethod + def restore_scheduler( + scheduler: Union["Scheduler", List, Dict], args: Coqpit, config: Coqpit, restore_epoch: int, restore_step: int + ) -> Union["Scheduler", List]: + """Restore scheduler wrt restored model.""" + if scheduler is not None: # pylint: disable=too-many-nested-blocks + if args.continue_path: + if isinstance(scheduler, list): + for s in scheduler: + if s is not None: + if config.scheduler_after_epoch: + s.last_epoch = restore_epoch + else: + s.last_epoch = restore_step + elif isinstance(scheduler, dict): + for s in scheduler.values(): + if s is not None: + if config.scheduler_after_epoch: + s.last_epoch = restore_epoch + else: + s.last_epoch = restore_step + else: + if config.scheduler_after_epoch: + scheduler.last_epoch = restore_epoch + else: + scheduler.last_epoch = restore_step + return scheduler + + @staticmethod + def get_criterion(model: nn.Module) -> nn.Module: + """Receive the criterion from the model. Model must implement `get_criterion()`. + + Args: + model (nn.Module): Training model. + + Returns: + nn.Module: Criterion layer. + """ + criterion = None + criterion = model.get_criterion() + return criterion + + #################### + # HELPER FUNCTIONS + #################### + + @staticmethod + def _detach_loss_dict(loss_dict: Dict) -> Dict: + """Detach loss values from autograp. + + Args: + loss_dict (Dict): losses. + + Returns: + Dict: losses detached from autograph. + """ + loss_dict_detached = {} + for key, value in loss_dict.items(): + if isinstance(value, (int, float)): + loss_dict_detached[key] = value + else: + loss_dict_detached[key] = value.detach().cpu().item() + return loss_dict_detached + + def _pick_target_avg_loss(self, keep_avg_target: KeepAverage) -> Dict: + """Pick the target loss to compare models""" + + # if the keep_avg_target is None or empty return None + if keep_avg_target is None or len(list(keep_avg_target.avg_values.keys())) == 0: + return None + + target_avg_loss = None + # return if target loss defined in the model config + # if not available in Dict use loss_1 as by default loss + if "target_loss" in self.config and self.config.target_loss: + if f"avg_{self.config.target_loss}" in keep_avg_target.avg_values.keys(): + return keep_avg_target[f"avg_{self.config.target_loss}"] + + raise ValueError( + " [!] Target loss not found in the keep_avg_target. You might be exiting the training loop before it is computed or set the target_loss in the model config incorrectly." + ) + + # take the average of loss_{optimizer_idx} as the target loss when there are multiple optimizers + if isinstance(self.optimizer, list): + target_avg_loss = 0 + for idx in range(len(self.optimizer)): + if f"avg_loss_{idx}" in keep_avg_target.avg_values: + target_avg_loss += keep_avg_target[f"avg_loss_{idx}"] + target_avg_loss /= len(self.optimizer) + else: + target_avg_loss = keep_avg_target.avg_values.get("avg_loss", 0) + return target_avg_loss + + def _setup_logger_config(self, log_file: str) -> None: + """Set up the logger based on the process rank in DDP.""" + + logger_new = logging.getLogger("trainer") + handler = logging.FileHandler(log_file, mode="a") + fmt = logging.Formatter("") + handler.setFormatter(fmt) + logger_new.addHandler(handler) + + # only log to a file if rank > 0 in DDP + if self.args.rank > 0: + logger_new.handlers = [h for h in logger_new.handlers if not isinstance(h, logging.StreamHandler)] + + @staticmethod + def _is_apex_available() -> bool: + """Check if Nvidia's APEX is available.""" + return importlib.util.find_spec("apex") is not None From a97ad658c1796f9dea29636cdebb58125fd7412d Mon Sep 17 00:00:00 2001 From: bmalaski <5288648+bmalaski@users.noreply.github.com> Date: Thu, 20 Jun 2024 22:54:57 -0400 Subject: [PATCH 2/7] ported changes form original branch --- finetune.py | 452 ++++++++++++++++++++---- start_diagnostics.bat | 6 + system/tts_engines/xtts/model_engine.py | 10 +- tts_server.py | 2 +- 4 files changed, 407 insertions(+), 63 deletions(-) create mode 100644 start_diagnostics.bat diff --git a/finetune.py b/finetune.py index 7bbd94e..095afe4 100644 --- a/finetune.py +++ b/finetune.py @@ -1,3 +1,4 @@ +import math import os import gc import re @@ -6,6 +7,8 @@ import time import glob import json import site +import datetime + import torch import signal import random @@ -25,16 +28,22 @@ from pathlib import Path from faster_whisper import WhisperModel from TTS.tts.configs.xtts_config import XttsConfig from TTS.tts.models.xtts import Xtts + +from metrics_logger import MetricsLogger # Use a local Tokenizer to resolve Japanese support # from TTS.tts.layers.xtts.tokenizer import multilingual_cleaners from system.ft_tokenizer.tokenizer import multilingual_cleaners import importlib.metadata as metadata from packaging import version +from tokenizers import ByteLevelBPETokenizer +from tokenizers.pre_tokenizers import Whitespace # STARTUP VARIABLES this_dir = Path(__file__).parent.resolve() audio_folder = this_dir / "finetune" / "put-voice-samples-in-here" -out_path = this_dir / "finetune" / "tmp-trn" +default_path = this_dir / "finetune" / "tmp-trn" +out_path = default_path + progress = 0 theme = gr.themes.Default() refresh_symbol = 'πŸ”„' @@ -310,13 +319,17 @@ def create_temporary_file(folder, suffix=".wav"): -def format_audio_list(target_language, whisper_model, out_path, max_sample_length, eval_split_number, speaker_name_input, gradio_progress=gr.Progress()): +def format_audio_list(target_language, whisper_model, max_sample_length, eval_split_number, speaker_name_input, create_bpe_tokenizer, gradio_progress=gr.Progress()): global validate_train_metadata_path, validate_eval_metadata_path, validate_audio_folder, validate_whisper_model, validate_target_language buffer = 0.3 max_duration = float(max_sample_length) # Ensure max_duration is a float eval_percentage = eval_split_number / 100.0 speaker_name = speaker_name_input audio_total_size = 0 + global out_path + if speaker_name and speaker_name != 'personsname': + out_path = this_dir / "finetune" / speaker_name + os.makedirs(out_path, exist_ok=True) temp_folder = os.path.join(out_path, "temp") os.makedirs(temp_folder, exist_ok=True) @@ -367,8 +380,13 @@ def format_audio_list(target_language, whisper_model, out_path, max_sample_lengt print(f"[FINETUNE] No audio files found in {original_samples_folder}. Skipping processing.") return None, None, 0 + whisper_words = [] + audio_steps = (0, len(original_audio_files)) + gradio_progress_duration = 0 + gradio_progress(audio_steps, desc="Processing Audio Files", unit="files") + for audio_path in tqdm(original_audio_files): # No start argument for tqdm - gradio_progress((2, 4), desc="Processing Audio Files") + start = datetime.datetime.now() audio_file_name_without_ext, _ = os.path.splitext(os.path.basename(audio_path)) temp_audio_path = os.path.join(temp_folder, f"{audio_file_name_without_ext}.wav") @@ -415,6 +433,9 @@ def format_audio_list(target_language, whisper_model, out_path, max_sample_lengt words_list.extend(words) for word_idx, word in enumerate(words_list): + if create_bpe_tokenizer: + whisper_words.append(word.word) + if first_word: sentence_start = word.start if word_idx == 0: @@ -481,6 +502,14 @@ def format_audio_list(target_language, whisper_model, out_path, max_sample_lengt os.remove(temp_audio_path) + end = datetime.datetime.now() + gradio_progress_duration += (end - start).total_seconds() + audio_steps = (audio_steps[0] + 1, audio_steps[1]) + additional_data_points_needed = audio_steps[1] - audio_steps[0] + avg_duration = gradio_progress_duration / audio_steps[0] + gradio_estimated_duration = (avg_duration * additional_data_points_needed) + gradio_progress(audio_steps, desc=f"Processing. Estimated Completion: {c_logger.format_duration(gradio_estimated_duration)}", unit="files") + # Check if the WAV files folder contains files after processing audio_files = [os.path.join(audio_folder, file) for file in os.listdir(audio_folder) if file.endswith('.wav')] @@ -532,6 +561,17 @@ def format_audio_list(target_language, whisper_model, out_path, max_sample_lengt final_training_set.sort_values('audio_file').to_csv(train_metadata_path, sep='|', index=False) final_eval_set.sort_values('audio_file').to_csv(eval_metadata_path, sep='|', index=False) + if create_bpe_tokenizer: + print(f"Training BPE Tokenizer") + tokenizer = ByteLevelBPETokenizer(str(this_dir / base_path / "xttsv2_2.0.3" / "vocab.json")) + tokenizer.pre_tokenizer = Whitespace() + tokenizer.add_tokens(["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]", "[STOP]", "[SAPCE]"]) + tokenizer.train_from_iterator(whisper_words, vocab_size=30000, show_progress=True, min_frequency=2, special_tokens=[ + "[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]", "[STOP]", "[SPACE]", + ]) + # Save the tokenizer model + tokenizer.save(path=str(out_path / "bpe_tokenizer-vocab.json"), pretty=True) + gradio_progress((4,4), "Finalizing") del asr_model, final_eval_set, final_training_set, new_data_df, existing_metadata @@ -648,7 +688,8 @@ def load_and_display_mismatches(): #### STEP 2 BITS ##### ###################### -from trainer import Trainer, TrainerArgs +from trainer import TrainerArgs, Trainer +from trainer_alltalk.trainer import Trainer from TTS.config.shared_configs import BaseDatasetConfig from TTS.tts.datasets import load_tts_samples @@ -663,10 +704,10 @@ def basemodel_or_finetunedmodel_choice(value): elif value == "Existing finetuned model": basemodel_or_finetunedmodel = False -def train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, learning_rate, model_to_train, output_path, max_audio_length=255995, progress=gr.Progress()): +def train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, learning_rate, model_to_train, continue_run, disable_shared_memory, learning_rate_scheduler, optimizer, max_audio_length=255995, progress=gr.Progress()): pfc_check_fail() - if "No Models Available" in eval_csv: + if "No Models Available" in model_to_train: print(f"[FINETUNE] \033[91mError\033[0m: Cannot train. You have not selected a model. Please download a model.") print(f"[FINETUNE] \033[91mError\033[0m: into a sub folder within the correct models folder.") return @@ -678,7 +719,7 @@ def train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, model_path = this_dir / "models" / "xtts" / model_to_train # Set here the path that the checkpoints will be saved. Default: ./training/ - OUT_PATH = os.path.join(output_path, "training") + OUT_PATH = os.path.join(out_path, "training") print("[FINETUNE] \033[94mStarting Step 2\033[0m - Fine-tuning the XTTS Encoder") print(f"[FINETUNE] \033[94mLanguage: \033[92m{language} \033[94mEpochs: \033[92m{num_epochs} \033[94mBatch size: \033[92m{batch_size}\033[0m \033[94mGrad accumulation steps: \033[92m{grad_acumm}\033[0m") print(f"[FINETUNE] \033[94mTraining : \033[92m{train_csv}\033[0m") @@ -731,8 +772,30 @@ def train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, MEL_NORM_FILE = model_path / "mel_stats.pth" SPEAKERS_FILE = model_path / "speakers_xtts.pth" + training_assets = None + if (out_path / "bpe_tokenizer-vocab.json").exists(): + print("[FINETUNE] Using custom BPE tokenizer") + training_assets = { + 'Tokenizer': str(out_path / "bpe_tokenizer-vocab.json") + } + + continue_path = None + if continue_run: + folders = glob.glob(os.path.join(OUT_PATH, '*/')) + if folders: + last_run = max(folders, key=os.path.getmtime) + if last_run: + checkpoints = glob.glob(os.path.join(last_run, "best_model_*.pth")) + if checkpoints: + latest_checkpoint = max(checkpoints, key=os.path.getmtime) + if latest_checkpoint: + XTTS_CHECKPOINT = None + continue_path = last_run + print(f"[FINETUNE] Continuing previous fine tuning {latest_checkpoint}") + + # Copy the supporting files - destination_dir = this_dir / "finetune" / "tmp-trn" / "chkptandnorm" + destination_dir = out_path / "chkptandnorm" destination_dir.mkdir(parents=True, exist_ok=True) shutil.copy2(DVAE_CHECKPOINT, destination_dir / DVAE_CHECKPOINT.name) shutil.copy2(MEL_NORM_FILE, destination_dir / MEL_NORM_FILE.name) @@ -763,6 +826,95 @@ def train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, if language == "ja": number_of_workers = 0 + lr_scheduler = None + lr_scheduler_params = {} + + if learning_rate_scheduler and learning_rate_scheduler != "None": + lr_gamma_mapping = { + 1e-6: 0.9, + 5e-6: 0.8, + 1e-5: 0.3, + 5e-5: 0.3, + 1e-4: 0.3, + 5e-4: 0.1, + 1e-3: 0.1 + } + lr_scheduler = learning_rate_scheduler + if lr_scheduler == "StepLR": + lr_scheduler_params = {'step_size': 30, 'gamma': 0.1, 'last_epoch': -1} + elif lr_scheduler == "MultiStepLR": + exponent = 3 - int(math.log2(num_epochs) / 2) + base = 2 + num_milestones = min(num_epochs, int(math.pow(base, exponent))) + milestone_interval = num_epochs // (num_milestones + 1) + milestones = [milestone_interval * (i + 1) for i in range(num_milestones)] + lr_scheduler_params = {'milestones': milestones, 'gamma': lr_gamma_mapping[learning_rate], 'last_epoch': -1} + elif lr_scheduler == "ExponentialLR": + lr_scheduler_params = {'gamma': 0.5, 'last_epoch': -1} + elif lr_scheduler == "CosineAnnealingLR": + lr_scheduler_params = {'T_max': num_epochs, 'eta_min': 1e-6, 'last_epoch': -1} + elif lr_scheduler == "ReduceLROnPlateau": + lr_scheduler_params = {'mode': 'min', 'factor': 0.8, 'patience': 1, 'threshold': 0.0001, + 'threshold_mode': 'rel', 'cooldown': 0, 'min_lr': 1e-8, 'eps': 1e-08,} + elif lr_scheduler == "CyclicLR": + lr_scheduler_params = {'base_lr': learning_rate, 'max_lr': 0.1, 'step_size_up': 2000, 'step_size_down': None, + 'mode': 'triangular', 'gamma': 1.0, 'scale_fn': None, 'scale_mode': 'cycle', + 'cycle_momentum': True, 'base_momentum': 0.8, 'max_momentum': 0.9, 'last_epoch': -1} + elif lr_scheduler == "OneCycleLR": + #TODO: Determine total steps / epoch for this + lr_scheduler_params = {'max_lr': learning_rate, 'total_steps': None, 'epochs_up': None, 'steps_per_epoch': None, + 'anneal_strategy': 'cos', 'cycle_momentum': True, 'base_momentum': 0.85, + 'max_momentum': 0.95, 'div_factor': 25.0, 'final_div_factor': 10000.0, + 'last_epoch': -1} + elif lr_scheduler == "CosineAnnealingWarmRestarts": + #Set 4 learning rate restarts + lr_scheduler_params = {'T_0': int(num_epochs / 4), 'T_mult': 1, 'eta_min': 1e-6, 'last_epoch': -1} + + optimizer_params = None + optimizer_wd_only_on_weight = OPTIMIZER_WD_ONLY_ON_WEIGHTS + + if optimizer == "AdamW": + optimizer_params = { + "betas": [0.9, 0.96], + "eps": 1e-8, + "weight_decay": 1e-2 + } + elif optimizer == "RMSprop": + optimizer_params = { + "alpha": 0.99, + "eps": 1e-8, + "weight_decay": 1e-4 + } + elif optimizer == "SGD": + optimizer_params = { + "momentum": 0.9, + "weight_decay": 1e-4 + } + elif optimizer == "Adam": + optimizer_params = { + "betas": [0.9, 0.999], + "eps": 1e-8, + "weight_decay": 1e-4 + } + elif optimizer == "Adagrad": + optimizer_params = { + "lr_decay": 0, + "weight_decay": 1e-4, + "eps": 1e-10 + } + elif optimizer == "RAdam": + optimizer_params = { + "betas": [0.9, 0.999], + "eps": 1e-8, + "weight_decay": 1e-2 + } + elif optimizer == "stepwisegraduallr" or optimizer == "noamlr": + optimizer_params = { } + # These are Hardcoded learning rate schedulers + + + print(f"[FINETUNE] Learning Scheduler {lr_scheduler}, params {lr_scheduler_params}") + # training parameters config config = GPTTrainerConfig( epochs=num_epochs, @@ -788,13 +940,13 @@ def train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, # target_loss="loss", print_eval=False, # Optimizer values like tortoise, pytorch implementation with modifications to not apply WD to non-weight parameters. - optimizer="AdamW", + optimizer=optimizer, optimizer_wd_only_on_weights=OPTIMIZER_WD_ONLY_ON_WEIGHTS, - optimizer_params={"betas": [0.9, 0.96], "eps": 1e-8, "weight_decay": 1e-2}, + optimizer_params=optimizer_params, lr=learning_rate, # learning rate - lr_scheduler="MultiStepLR", + lr_scheduler=lr_scheduler, # it was adjusted accordly for the new step scheme - lr_scheduler_params={"milestones": [50000 * 18, 150000 * 18, 300000 * 18], "gamma": 0.5, "last_epoch": -1}, + lr_scheduler_params=lr_scheduler_params, test_sentences=[], ) progress(0, desc="Model is currently training. See console for more information") @@ -808,11 +960,15 @@ def train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, eval_split_size=config.eval_split_size, ) + global c_logger + c_logger = MetricsLogger() + # init the trainer trainer = Trainer( TrainerArgs( restore_path=None, # xtts checkpoint is restored via xtts_checkpoint key so no need of restore it using Trainer restore_path parameter skip_train_epoch=False, + continue_path=continue_path, start_with_eval=START_WITH_EVAL, grad_accum_steps=GRAD_ACUMM_STEPS, ), @@ -821,7 +977,15 @@ def train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, model=model, train_samples=train_samples, eval_samples=eval_samples, + training_assets=training_assets, + c_logger=c_logger, ) + + if(disable_shared_memory): + # Limit training to GPU memory instead of shared memory + print("[FINETUNE] Limiting GPU memory to 95% to prevent spillover") + torch.cuda.set_per_process_memory_fraction(0.95) + trainer.fit() # get the longest text audio file to use as speaker reference @@ -884,7 +1048,13 @@ def run_tts(lang, tts_text, speaker_audio_file): if XTTS_MODEL is None or not speaker_audio_file: return "You need to run the previous step to load the model !!", None, None speaker_audio_file = str(speaker_audio_file) - gpt_cond_latent, speaker_embedding = XTTS_MODEL.get_conditioning_latents(audio_path=speaker_audio_file, gpt_cond_len=XTTS_MODEL.config.gpt_cond_len, max_ref_length=XTTS_MODEL.config.max_ref_len, sound_norm_refs=XTTS_MODEL.config.sound_norm_refs) + if os.path.isdir(speaker_audio_file): + wavs_files = glob.glob(os.path.join(speaker_audio_file, "*.wav")) + speaker_audio_file = wavs_files[0] + else: + wavs_files = [speaker_audio_file] + + gpt_cond_latent, speaker_embedding = XTTS_MODEL.get_conditioning_latents(audio_path=wavs_files, gpt_cond_len=XTTS_MODEL.config.gpt_cond_len, max_ref_length=XTTS_MODEL.config.max_ref_len, sound_norm_refs=XTTS_MODEL.config.sound_norm_refs) out = XTTS_MODEL.inference( text=tts_text, language=lang, @@ -904,33 +1074,45 @@ def run_tts(lang, tts_text, speaker_audio_file): return "Speech generated !", out_path, speaker_audio_file -def get_available_voices(minimum_size_kb=1200): +def get_available_voices(minimum_size_kb=1200, speaker_name=None): + directory = out_path + if speaker_name and speaker_name != 'personsname': + directory = this_dir / "finetune" / speaker_name + voice_files = [ - voice for voice in Path(f"{this_dir}/finetune/tmp-trn/wavs").glob("*.wav") + voice for voice in Path(f"{directory}/wavs").glob("*.wav") if voice.stat().st_size > minimum_size_kb * 1200 # Convert KB to bytes ] return sorted([str(file) for file in voice_files]) # Return full path as string -def find_best_models(directory): +def find_best_models(directory, speaker_name=None): + if speaker_name and speaker_name != 'personsname': + directory = this_dir / "finetune" / speaker_name + """Find files named 'best_model.pth' in the given directory.""" return [str(file) for file in Path(directory).rglob("best_model.pth")] -def find_models(directory, extension): +def find_models(directory, extension, speaker_name=None): + if speaker_name and speaker_name != 'personsname': + directory = this_dir / "finetune" / speaker_name + """Find files with a specific extension in the given directory.""" return [str(file) for file in Path(directory).rglob(f"*.{extension}")] -def find_jsons(directory, filename): +def find_jsons(directory, filename, speaker_name=None): + if speaker_name and speaker_name != 'personsname': + directory = this_dir / "finetune" / speaker_name + """Find files with a specific filename in the given directory.""" return [str(file) for file in Path(directory).rglob(filename)] -# Your main directory -main_directory = Path(this_dir) / "finetune" / "tmp-trn" + # XTTS checkpoint files (best_model.pth) -xtts_checkpoint_files = find_best_models(main_directory) +xtts_checkpoint_files = find_best_models(out_path) # XTTS config files (config.json) -xtts_config_files = find_jsons(main_directory, "config.json") +xtts_config_files = find_jsons(out_path, "config.json") # XTTS vocab files (vocab.json) -xtts_vocab_files = find_jsons(main_directory, "vocab.json") +xtts_vocab_files = find_jsons(out_path, "vocab.json") ########################## #### STEP 4 AND OTHER #### @@ -976,11 +1158,11 @@ def compact_custom_model(xtts_checkpoint_copy, folder_path, overwrite_existing): files_to_copy = ["vocab.json", "config.json"] copy_files(folder_path_new, target_dir, files_to_copy) # Copy second set of files - chkptandnorm_path = this_dir / "finetune" / "tmp-trn" / "chkptandnorm" + chkptandnorm_path = out_path / "chkptandnorm" files_to_copy2 = ["speakers_xtts.pth", "mel_stats.pth", "dvae.pth"] copy_files(chkptandnorm_path, target_dir, files_to_copy2) # Copy large wav files - source_wavs_dir = this_dir / "finetune" / "tmp-trn" / "wavs" + source_wavs_dir = out_path / "wavs" target_wavs_dir = target_dir / "wavs" target_wavs_dir.mkdir(parents=True, exist_ok=True) for file_path in source_wavs_dir.iterdir(): @@ -991,7 +1173,7 @@ def compact_custom_model(xtts_checkpoint_copy, folder_path, overwrite_existing): def delete_training_data(): # Define the folder to be deleted - folder_to_delete = Path(this_dir / "finetune" / "tmp-trn") + folder_to_delete = Path(out_path) # Check if the folder exists before deleting if folder_to_delete.exists(): @@ -1077,6 +1259,9 @@ logging.basicConfig( ] ) +c_logger = MetricsLogger() +def load_metrics(): + return c_logger.plot_metrics(), f"Running Time: {c_logger.format_duration(c_logger.total_duration)} - Estimated Completion: {c_logger.format_duration(c_logger.estimated_duration)}" def read_logs(): sys.stdout.flush() with open(sys.stdout.log_file, "r") as f: @@ -1088,10 +1273,10 @@ def cleanup_before_exit(signum, frame): sys.exit(0) def create_refresh_button(refresh_components, refresh_methods, elem_class, interactive=True): - def refresh(): + def refresh(speaker_name): updates = {} for component, method in zip(refresh_components, refresh_methods): - args = method() if callable(method) else method + args = method(speaker_name=speaker_name) if callable(method) else method if args and 'choices' in args: # Select the most recent file (last in the sorted list) args['value'] = args['choices'][-1] if args['choices'] else "" @@ -1103,12 +1288,33 @@ def create_refresh_button(refresh_components, refresh_methods, elem_class, inter refresh_button = gr.Button("Refresh Dropdowns", elem_classes=elem_class, interactive=interactive) refresh_button.click( fn=refresh, - inputs=[], + inputs=[speaker_name_input_testing], outputs=refresh_components ) return refresh_button +def create_refresh_button_next(refresh_components, refresh_methods, elem_class, interactive=True): + def refresh_export(speaker_name): + updates = {} + for component, method in zip(refresh_components, refresh_methods): + args = method(speaker_name=speaker_name) if callable(method) else method + if args and 'choices' in args: + # Select the most recent file (last in the sorted list) + args['value'] = args['choices'][-1] if args['choices'] else "" + for k, v in args.items(): + setattr(component, k, v) + updates[component] = gr.update(**(args or {})) + return updates + + refresh_button = gr.Button("Refresh Dropdowns", elem_classes=elem_class, interactive=interactive) + refresh_button.click( + fn=refresh_export, + inputs=[speaker_name_input_export], + outputs=refresh_components + ) + + return refresh_button pfc_markdown = f""" @@ -1315,11 +1521,11 @@ if __name__ == "__main__": ####################### with gr.Tab("πŸ“ Step 1 - Generating the dataset"): with gr.Tab("Generate Dataset"): - out_path = gr.Textbox( - label="Output path (where data and checkpoints will be saved):", - value=out_path, - visible=False, - ) + # out_path = gr.Textbox( + # label="Output path (where data and checkpoints will be saved):", + # value=out_path, + # visible=False, + # ) # Define directories this_dir = Path(__file__).parent.resolve() voice_samples_dir = this_dir / "finetune" / "put-voice-samples-in-here" @@ -1381,6 +1587,13 @@ if __name__ == "__main__": delete_dataset_button.click(delete_existing_training_data, outputs=output_text) with gr.Row(): + speaker_name_input = gr.Textbox( + label="Training Project Name", + value="personsname", + visible=True, + scale=2, + info='Create Unique Training Folder', + ) whisper_model = gr.Dropdown( label="Whisper Model", value="large-v2", @@ -1431,12 +1644,11 @@ if __name__ == "__main__": step=1, # Increment step scale=2, ) - speaker_name_input = gr.Textbox( - label="The name of the speaker/person you are training", - value="personsname", - visible=True, - scale=2, - ) + create_bpe_tokenizer = gr.Checkbox( + label='BPE Tokenizer', + value=False, + info='Custom Tokenizer for training' + ) progress_data = gr.Label( label="Progress:" ) @@ -1449,7 +1661,7 @@ if __name__ == "__main__": prompt_compute_btn = gr.Button(value="Step 1 - Create dataset") - def preprocess_dataset(language, whisper_model, out_path, max_sample_length, eval_split_number, speaker_name_input, progress=gr.Progress()): + def preprocess_dataset(language, whisper_model, max_sample_length, eval_split_number, speaker_name_input, create_bpe_tokenizer, progress=gr.Progress()): clear_gpu_cache() test_for_audio_files = [file for file in os.listdir(audio_folder) if any(file.lower().endswith(ext) for ext in ['.wav', '.mp3', '.flac'])] if not test_for_audio_files: @@ -1457,7 +1669,7 @@ if __name__ == "__main__": else: try: - train_meta, eval_meta, audio_total_size = format_audio_list(target_language=language, whisper_model=whisper_model, out_path=out_path, max_sample_length=max_sample_length, eval_split_number=eval_split_number, speaker_name_input=speaker_name_input, gradio_progress=progress) + train_meta, eval_meta, audio_total_size = format_audio_list(target_language=language, whisper_model=whisper_model, max_sample_length=max_sample_length, eval_split_number=eval_split_number, speaker_name_input=speaker_name_input, create_bpe_tokenizer=create_bpe_tokenizer, gradio_progress=progress) except: traceback.print_exc() error = traceback.format_exc() @@ -1472,7 +1684,7 @@ if __name__ == "__main__": return message, "", "" print("[FINETUNE] Dataset Generated. Either run Dataset Validation or move to Step 2") - return "Dataset Generated. Either run Dataset Validation or move to Step 2", train_meta, eval_meta + return "Dataset Generated. Either run Dataset Validation or move to Step 2", train_meta, eval_meta, speaker_name_input def play_selected_audio(selected_file): if selected_file and selected_file != "Please Generate Your Dataset": @@ -1534,8 +1746,26 @@ if __name__ == "__main__": #### GRADIO STEP 2 #### ####################### with gr.Tab("πŸ’» Step 2 - Training"): - with gr.Tab("Training the model"): + with gr.Tab("Training the model"): with gr.Row(): + speaker_name_input_training = gr.Textbox( + label="Project Name", + value="personsname", + visible=True, + info='Create Unique Training Folder', + scale=1, + ) + with gr.Group(): + continue_run = gr.Checkbox( + value=False, + label="Continue Previous", + scale=1, + ) + disable_shared_memory = gr.Checkbox( + value=False, + label="Disable Shared Memory", + scale=1, + ) model_to_train_choices = list(available_models.keys()) model_to_train = gr.Dropdown( choices=model_to_train_choices, @@ -1550,6 +1780,7 @@ if __name__ == "__main__": label="Eval CSV:", scale=2, ) + with gr.Row(): learning_rates = gr.Dropdown( value=5e-6, label="Learning Rate", @@ -1564,7 +1795,42 @@ if __name__ == "__main__": ], type="value", allow_custom_value=True, - scale=0, + scale=1, + ) + learning_rate_scheduler = gr.Dropdown( + value="None", + label="Learning Rate Scheduler(s)", + choices=[ + ("None", "None"), + ("Cosine Annealing", "CosineAnnealingLR"), + ("Cosine Annealing Warm Restarts", "CosineAnnealingWarmRestarts"), + ("Cyclic", "CyclicLR"), + ("Exponential", "ExponentialLR"), + ("Multi Step", "MultiStepLR"), + ("Reduce on Plateau", "ReduceLROnPlateau"), + ("Step", "StepLR") + #("OneCycleLR", "OneCycleLR"), + ], + type="value", + allow_custom_value=False, + scale=2, + ) + optimizer = gr.Dropdown( + value="AdamW", + label="Optimizer", + choices=[ + ("AdamW", "AdamW"), + ("SGD with Momentum", "SGD"), + ("RMSprop", "RMSprop"), + ("Adagrad", "Adagrad"), + ("Adam", "Adam"), + ("Rectified Adam", "RAdam"), + ("Step Wise", "stepwisegraduallr"), + ("Noam", "noamlr") + ], + type="value", + allow_custom_value=False, + scale=2, ) with gr.Row(): num_epochs = gr.Slider( @@ -1599,24 +1865,41 @@ if __name__ == "__main__": progress_train = gr.Label( label="Progress:" ) + + with gr.Row(): + train_time = gr.Label("Estimated Total Training Time", show_label=False) + + with gr.Row(): + model_data = gr.Image(c_logger.plot_metrics(), show_label=False) + logs_tts_train = gr.Textbox( label="Logs:", interactive=False, lines=10, ) + demo.load(load_metrics, None, [model_data, train_time], every=1) demo.load(read_logs, None, logs_tts_train, every=1) train_btn = gr.Button(value="Step 2 - Run the training") - def train_model(language, train_csv, eval_csv, learning_rates, model_to_train, num_epochs, batch_size, grad_acumm, output_path, max_audio_length, progress=gr.Progress()): + def train_model(language, train_csv, eval_csv, learning_rates, model_to_train, num_epochs, batch_size, grad_acumm, max_audio_length, speaker_name_input_training, continue_run, disable_shared_memory, learning_rate_scheduler, optimizer, progress=gr.Progress()): clear_gpu_cache() + global out_path + if speaker_name_input_training and speaker_name_input_training != 'personsname': + out_path = this_dir / "finetune" / speaker_name_input_training + if not train_csv or not eval_csv: - return "You need to run the data processing step or manually set `Train CSV` and `Eval CSV` fields !", "", "", "", "" + if (out_path / "metadata_eval.csv").exists() and (out_path / "metadata_train.csv").exists(): + train_csv = out_path / "metadata_eval.csv" + eval_csv = out_path / "metadata_train.csv" + print("Using existing metadata and training csv.") + else: + return "You need to run the data processing step or manually set `Train CSV` and `Eval CSV` fields !", "", "", "", "" try: # convert seconds to waveform frames max_audio_length = int(max_audio_length * 22050) learning_rate = float(learning_rates) # Convert the learning rate value to a float progress(0, "Initializing training...") - config_path, original_xtts_checkpoint, vocab_file, exp_path, speaker_wav = train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, learning_rate, model_to_train, output_path=str(output_path), max_audio_length=max_audio_length, progress=gr.Progress()) + config_path, original_xtts_checkpoint, vocab_file, exp_path, speaker_wav = train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, learning_rate, model_to_train, continue_run, disable_shared_memory, learning_rate_scheduler, optimizer, max_audio_length=max_audio_length, progress=gr.Progress()) except: traceback.print_exc() error = traceback.format_exc() @@ -1629,7 +1912,7 @@ if __name__ == "__main__": ft_xtts_checkpoint = os.path.join(exp_path, "best_model.pth") print("[FINETUNE] Model training done. Move to Step 3") clear_gpu_cache() - return "Model training done. Move to Step 3", config_path, vocab_file, ft_xtts_checkpoint, speaker_wav + return "Model training done. Move to Step 3", config_path, vocab_file, ft_xtts_checkpoint, speaker_wav, speaker_name_input_training with gr.Tab("Training the model Instructions"): gr.Markdown( @@ -1699,6 +1982,35 @@ if __name__ == "__main__": The optimal learning rate depends on the model architecture, dataset, and other hyperparameters. """ ) + with gr.Tab("Info - Learning Rate Schedulers"): + gr.Markdown( + f""" + ### 🟩 Learning Rate Schedulers
+ β—½ None: Learning rate remains unchanged throughout the training.
+ β—½ Step: Decreases the learning rate by a factor at specified intervals, offering a straightforward method for controlling learning rate decay.
+ β—½ MultiStep: Decreases the learning rate by a factor at specified milestones, useful for gradually reducing the learning rate during training.
+ β—½ Exponential: Decreases the learning rate exponentially over epochs, providing a simple decay mechanism for training.
+ β—½ CosineAnnealing: Adjusts the learning rate following a cosine annealing schedule, facilitating smoother optimization and potentially better convergence.
+ β—½ Reduce On Plateau: Reduces the learning rate when a monitored metric has stopped improving, helping to fine-tune learning rates during training.
+ β—½ Cyclic: Cycles the learning rate between two boundary values, allowing for dynamic learning rate schedules during training.
+ β—½ One Cycle (disabled): Sets the learning rate using a cyclical learning rate policy and varies momentum inversely, potentially accelerating training and improving performance.
+ β—½ Cosine Annealing Warm Restarts: Adjusts the learning rate following a cosine annealing schedule with warm restarts, offering a balance between exploration and exploitation during optimization.

+ The choice of learning rate scheduler depends on the specific characteristics of the model, dataset, and training objectives. + """ + ) + with gr.Tab("Info - Optimizers"): + gr.Markdown( + f""" + ### 🟩 Optimizers
+ β—½ SGD with Momentum: Good for convergence, requires careful tuning of the learning rate and momentum. + β—½ RMSprop: Adaptively adjusts learning rates, good for varying gradient magnitudes. + β—½ Adam: Widely used, performs well with default parameters. + β—½ Adagrad: Good for sparse data, adapts learning rates based on historical gradients. + β—½ AdamW: Similar to Adam but with decoupled weight decay for better regularization. + When picking an optimizer, you need to experiment with different options to find the one that achieves the best performance and convergence for your task. + AdamW has been set to default as its a sensible pick that will handle most needs. + """ + ) with gr.Tab("Info - Epochs"): gr.Markdown( f""" @@ -1778,6 +2090,13 @@ if __name__ == "__main__": allow_custom_value=True, # Allow custom values scale=1, ) + speaker_name_input_testing = gr.Textbox( + label="Project Name", + value="personsname", + visible=True, + info='Create Unique Training Folder', + scale=1, + ) with gr.Row(): tts_language = gr.Dropdown( @@ -1804,12 +2123,13 @@ if __name__ == "__main__": ) # Create refresh button refresh_button = create_refresh_button( - [xtts_checkpoint, xtts_config, xtts_vocab, speaker_reference_audio], + [xtts_checkpoint, xtts_config, xtts_vocab, speaker_reference_audio, + speaker_name_input_testing], [ - lambda: {"choices": find_best_models(main_directory), "value": ""}, - lambda: {"choices": find_jsons(main_directory, "config.json"), "value": ""}, - lambda: {"choices": find_jsons(main_directory, "vocab.json"), "value": ""}, - lambda: {"choices": get_available_voices(), "value": ""}, + lambda speaker_name: {"choices": find_best_models(out_path, speaker_name=speaker_name), "value": ""}, + lambda speaker_name: {"choices": find_jsons(out_path, "config.json", speaker_name=speaker_name), "value": ""}, + lambda speaker_name: {"choices": find_jsons(out_path, "vocab.json", speaker_name=speaker_name),"value": ""}, + lambda speaker_name: {"choices": get_available_voices(speaker_name=speaker_name),"value": ""}, ], elem_class="refresh-button-class" ) @@ -1878,11 +2198,18 @@ if __name__ == "__main__": allow_custom_value=True, scale=2, ) + speaker_name_input_export = gr.Textbox( + label="Project Name", + value="personsname", + visible=True, + info='Unique Training Folder', + scale=1, + ) # Create refresh button - refresh_button = create_refresh_button( - [xtts_checkpoint_copy,], + refresh_button = create_refresh_button_next( + [xtts_checkpoint_copy,speaker_name_input_export], [ - lambda: {"choices": find_best_models(main_directory), "value": ""}, + lambda speaker_name: {"choices": find_best_models(out_path, speaker_name), "value": ""}, ], elem_class="refresh-button-class") with gr.Row(): @@ -1901,15 +2228,16 @@ if __name__ == "__main__": inputs=[ lang, whisper_model, - out_path, max_sample_length, eval_split_number, speaker_name_input, + create_bpe_tokenizer ], outputs=[ progress_data, train_csv, eval_csv, + speaker_name_input_training ], ) @@ -1925,10 +2253,14 @@ if __name__ == "__main__": num_epochs, batch_size, grad_acumm, - out_path, max_audio_length, + speaker_name_input_training, + continue_run, + disable_shared_memory, + learning_rate_scheduler, + optimizer ], - outputs=[progress_train, xtts_config, xtts_vocab, xtts_checkpoint, speaker_reference_audio], + outputs=[progress_train, xtts_config, xtts_vocab, xtts_checkpoint, speaker_reference_audio, speaker_name_input_testing], ) load_btn.click( diff --git a/start_diagnostics.bat b/start_diagnostics.bat new file mode 100644 index 0000000..f36efbc --- /dev/null +++ b/start_diagnostics.bat @@ -0,0 +1,6 @@ +@echo off +cd /D "C:\AI\text-generation-webui\extensions\alltalk_tts\" +set CONDA_ROOT_PREFIX=C:\AI\text-generation-webui\extensions\alltalk_tts\alltalk_environment\conda +set INSTALL_ENV_DIR=C:\AI\text-generation-webui\extensions\alltalk_tts\alltalk_environment\env +call "C:\AI\text-generation-webui\extensions\alltalk_tts\alltalk_environment\conda\condabin\conda.bat" activate "C:\AI\text-generation-webui\extensions\alltalk_tts\alltalk_environment\env" +call python diagnostics.py diff --git a/system/tts_engines/xtts/model_engine.py b/system/tts_engines/xtts/model_engine.py index 7c14395..b7abc5f 100644 --- a/system/tts_engines/xtts/model_engine.py +++ b/system/tts_engines/xtts/model_engine.py @@ -1,6 +1,7 @@ ############################################### # DONT CHANGE # These are base imports needed # ############################################### +import glob import os import sys import json @@ -471,10 +472,15 @@ class tts_class: print(f"[{self.branding}Debug] Deciding if streaming or not") if self.debug_tts else None print(f"[{self.branding}Debug] self.current_model_loaded", self.current_model_loaded) if self.debug_tts else None self.current_model_loaded + if os.path.isdir(voice): + wavs_files = glob.glob(os.path.join(voice, "*.wav")) + else: + wavs_files = [f"{self.main_dir}/voices/{voice}"] + if self.current_model_loaded.startswith ("xtts"): print(f"[{self.branding}Debug] Text arriving at engine {text}") if self.debug_tts else None gpt_cond_latent, speaker_embedding = self.model.get_conditioning_latents( - audio_path=[f"{self.main_dir}/voices/{voice}"], + audio_path=wavs_files, gpt_cond_len=self.model.config.gpt_cond_len, max_ref_length=self.model.config.max_ref_len, sound_norm_refs=self.model.config.sound_norm_refs, @@ -551,7 +557,7 @@ class tts_class: self.model.tts_to_file( text=text, file_path=output_file, - speaker_wav=[f"{self.main_dir}/voices/{voice}"], + speaker_wav=wavs_files, language=language, temperature=temperature, length_penalty=self.model.config.length_penalty, diff --git a/tts_server.py b/tts_server.py index 71424fe..dcc4140 100644 --- a/tts_server.py +++ b/tts_server.py @@ -1089,7 +1089,7 @@ def process_rvc_narrator(part_type, voice_gen, pitch, default_model_file, output class JSONInput(BaseModel): text_input: str = Field(..., max_length=int(api_defaults["api_max_characters"]), description=f"text_input needs to be {api_defaults['api_max_characters']} characters or less.") text_filtering: str = Field(..., pattern="^(none|standard|html)$", description="text_filtering needs to be 'none', 'standard' or 'html'.") - character_voice_gen: str = Field(..., pattern=r'^[\(\)a-zA-Z0-9\_\-./\s]+$', description="character_voice_gen needs to be the name of a valid voice for the loaded TTS engine.") + #character_voice_gen: str = Field(..., pattern=r'^[\(\)a-zA-Z0-9\_\-./\s]+$', description="character_voice_gen needs to be the name of a valid voice for the loaded TTS engine.") rvccharacter_voice_gen: str = Field(..., description="rvccharacter_voice_gen needs to be the name of a valid pth file in the 'folder\\file.pth' format or the word 'Disabled'.") rvccharacter_pitch: float = Field(..., description="RVC Character pitch needs to be a number between -24 and 24") narrator_enabled: str = Field(..., pattern="^(true|false|silent)$", description="narrator_enabled needs to be 'true', 'false' or 'silent'.") From e081442bf39394fee2988012e8b875651544a902 Mon Sep 17 00:00:00 2001 From: bmalaski <5288648+bmalaski@users.noreply.github.com> Date: Tue, 25 Jun 2024 15:49:06 -0400 Subject: [PATCH 3/7] Update finetune.py --- finetune.py | 14 +++++++------- start_diagnostics.bat | 6 ------ 2 files changed, 7 insertions(+), 13 deletions(-) delete mode 100644 start_diagnostics.bat diff --git a/finetune.py b/finetune.py index 095afe4..1809a98 100644 --- a/finetune.py +++ b/finetune.py @@ -843,12 +843,12 @@ def train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, if lr_scheduler == "StepLR": lr_scheduler_params = {'step_size': 30, 'gamma': 0.1, 'last_epoch': -1} elif lr_scheduler == "MultiStepLR": - exponent = 3 - int(math.log2(num_epochs) / 2) - base = 2 - num_milestones = min(num_epochs, int(math.pow(base, exponent))) - milestone_interval = num_epochs // (num_milestones + 1) - milestones = [milestone_interval * (i + 1) for i in range(num_milestones)] - lr_scheduler_params = {'milestones': milestones, 'gamma': lr_gamma_mapping[learning_rate], 'last_epoch': -1} + # exponent = 3 - int(math.log2(num_epochs) / 2) + # base = 2 + # num_milestones = min(num_epochs, int(math.pow(base, exponent))) + # milestone_interval = num_epochs // (num_milestones + 1) + # milestones = [milestone_interval * (i + 1) for i in range(num_milestones)] + lr_scheduler_params = {"milestones": [50000 * 18, 150000 * 18, 300000 * 18], "gamma": 0.5, "last_epoch": -1} elif lr_scheduler == "ExponentialLR": lr_scheduler_params = {'gamma': 0.5, 'last_epoch': -1} elif lr_scheduler == "CosineAnnealingLR": @@ -1850,7 +1850,7 @@ if __name__ == "__main__": grad_acumm = gr.Slider( label="Grad accumulation steps:", minimum=1, - maximum=32, + maximum=1024, step=1, value=args.grad_acumm, ) diff --git a/start_diagnostics.bat b/start_diagnostics.bat deleted file mode 100644 index f36efbc..0000000 --- a/start_diagnostics.bat +++ /dev/null @@ -1,6 +0,0 @@ -@echo off -cd /D "C:\AI\text-generation-webui\extensions\alltalk_tts\" -set CONDA_ROOT_PREFIX=C:\AI\text-generation-webui\extensions\alltalk_tts\alltalk_environment\conda -set INSTALL_ENV_DIR=C:\AI\text-generation-webui\extensions\alltalk_tts\alltalk_environment\env -call "C:\AI\text-generation-webui\extensions\alltalk_tts\alltalk_environment\conda\condabin\conda.bat" activate "C:\AI\text-generation-webui\extensions\alltalk_tts\alltalk_environment\env" -call python diagnostics.py From 797a4b996a1f5dc10019ea07f523d41f1c6d2dfb Mon Sep 17 00:00:00 2001 From: erew123 <35898566+erew123@users.noreply.github.com> Date: Mon, 1 Jul 2024 08:32:56 +0100 Subject: [PATCH 4/7] Update finetune.py --- finetune.py | 222 +++++++++++++++++++++++++++++++--------------------- 1 file changed, 134 insertions(+), 88 deletions(-) diff --git a/finetune.py b/finetune.py index 1809a98..48906f9 100644 --- a/finetune.py +++ b/finetune.py @@ -843,12 +843,12 @@ def train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, if lr_scheduler == "StepLR": lr_scheduler_params = {'step_size': 30, 'gamma': 0.1, 'last_epoch': -1} elif lr_scheduler == "MultiStepLR": - # exponent = 3 - int(math.log2(num_epochs) / 2) - # base = 2 - # num_milestones = min(num_epochs, int(math.pow(base, exponent))) - # milestone_interval = num_epochs // (num_milestones + 1) - # milestones = [milestone_interval * (i + 1) for i in range(num_milestones)] - lr_scheduler_params = {"milestones": [50000 * 18, 150000 * 18, 300000 * 18], "gamma": 0.5, "last_epoch": -1} + exponent = 3 - int(math.log2(num_epochs) / 2) + base = 2 + num_milestones = min(num_epochs, int(math.pow(base, exponent))) + milestone_interval = num_epochs // (num_milestones + 1) + milestones = [milestone_interval * (i + 1) for i in range(num_milestones)] + lr_scheduler_params = {'milestones': milestones, 'gamma': lr_gamma_mapping[learning_rate], 'last_epoch': -1} elif lr_scheduler == "ExponentialLR": lr_scheduler_params = {'gamma': 0.5, 'last_epoch': -1} elif lr_scheduler == "CosineAnnealingLR": @@ -867,6 +867,10 @@ def train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, 'max_momentum': 0.95, 'div_factor': 25.0, 'final_div_factor': 10000.0, 'last_epoch': -1} elif lr_scheduler == "CosineAnnealingWarmRestarts": + if num_epochs < 4: + error_message = "For Cosine Annealing Warm Restarts, epochs must be at least 4. Please set a minimum of 4 epochs." + progress(1.0, desc=f"Error: {error_message}") + raise ValueError(error_message) #Set 4 learning rate restarts lr_scheduler_params = {'T_0': int(num_epochs / 4), 'T_mult': 1, 'eta_min': 1e-6, 'last_epoch': -1} @@ -1592,7 +1596,6 @@ if __name__ == "__main__": value="personsname", visible=True, scale=2, - info='Create Unique Training Folder', ) whisper_model = gr.Dropdown( label="Whisper Model", @@ -1631,13 +1634,13 @@ if __name__ == "__main__": scale=1, ) max_sample_length = gr.Dropdown( - label="Maximum Audio Length Created by Whisper (seconds)", + label="Max Audio Length (in seconds)", value="30", choices=["10", "15", "20", "25", "30", "35", "40", "45", "50", "55", "60", "65", "70", "75", "80", "85", "90"], scale=2 ) eval_split_number = gr.Number( - label="Evaluation data Split (the % to use for Evaluation data)", + label="Evaluation Data Split", value=15, # Default value minimum=5, # Minimum value maximum=95, # Maximum value @@ -1711,36 +1714,57 @@ if __name__ == "__main__": with gr.Tab("Generate Dataset Instructions"): gr.Markdown( f""" - ### πŸ“ Generating the dataset
- ### πŸŸ₯ Important Note - Windows - "UserWarning: huggingface_hub cache-system uses symlinks." - β—½ This error is caused by Huggingfaces download software downloading the Whisper model. If you get this error, please restart your Windows command prompt with "Run as Administrator" and restart finetuning.
- β—½ This should only occur the 1st time it downloads the Whisper model. [Huggingface Reference here](https://huggingface.co/docs/huggingface_hub/en/guides/manage-cache#limitations) - ### πŸŸ₯ Important Note - Language support. - β—½ Although I have done my best to help automate this step, the Whisper model is not great with all languages and may fail to build your training data correctly. The Large-v3 may be better at certain languages.
- β—½ You can find information about the Whisper model [here](https://github.com/openai/whisper?tab=readme-ov-file#available-models-and-languages) and you can find data about manually building training data [here](https://docs.coqui.ai/en/latest/formatting_your_dataset.html), as well as details below about the file structure this step performs.
- ### 🟦 What you need to do - β—½ Please read Coqui's guide on what makes a good dataset [here](https://docs.coqui.ai/en/latest/what_makes_a_good_dataset.html#what-makes-a-good-dataset)
- β—½ Place your audio files in {str(audio_folder)} - β—½ Your audio samples can be in the format mp3, wav, or flac.
- β—½ You will need a minimum of 2 minutes of audio in either one or multiple audio files. 5 to 10 minutes of audio would probably be better, allowing for more varied sample data to be generated.
- β—½ Very small sample files cause errors, so I would recommend that the samples are at least 30 seconds and longer.
- β—½ FYI Anecdotal evidence suggests that the Whisper 2 model may yield superior results in audio splitting and dataset creation.
- ### 🟨 What this step is doing
- β—½ With step one, we are going to be stripping your audio file(s) into smaller files, using Whisper to find spoken words/sentences, compile that into excel sheets of training data, ready for Step 2.
- β—½ Whisper is making a best effort to find spoken audio and break it down into smaller audio files. The content of these audio files is then transcribed into CSV fles (which you can edit in Excel or similar).
- β—½ These files (audio and CSV) are used at the next step to train the model "this is what the audio sounds like and these are the words being spoken".
- β—½ If you are using Audio that has multiple speakers/people speaking, you will HAVE to edit the generated CSV files and remove the **OTHER** person out of the training list.
- β—½ If you wish to manually look at the CSV files before running the next step, you are welcome to do so and edit them as necessary. The greater the accuracy of the text the better the training will be.
-           β—½ /alltalk_tts/finetune/tmp-trn/lang.txt Contains a two digit langauage code e.g. `en`, `fr` etc.
-           β—½ /alltalk_tts/finetune/tmp-trn/metadata_train.csv Contains the list of wavs to text that will be used. This is what the model is trained with.
-           β—½ /alltalk_tts/finetune/tmp-trn/metadata_eval.csv Lists the evaluation wav to text, used to generate TTS while training and evaluate the quality of generated TTS to the original sample.
-           β—½ /alltalk_tts/finetune/tmp-trn/wavs/ These are the wavs split out from your original training samples and match the contents of the CSV files.
- β—½ Whilst you can choose multiple Whisper models, its best only to use the 1 model as each one is about 3GB in size and will download to your local huggingface cache on first-time use.
- β—½ If and when you have completed training, should you wish to delete this 3GB Whisper model from your system, you are welcome to do so.
- ### 🟩 How long will this take?
- β—½ First time, it needs to download the Whisper model which is 3GB. After that a few minutes on an average 3-4 year old system.
+ #### 🟦 What You Need to Do + 1. Read [Coqui's guide on creating a good dataset](https://docs.coqui.ai/en/latest/what_makes_a_good_dataset.html).
+ 2. Place your audio files (MP3, WAV, or FLAC) in `{str(audio_folder)}` Or Upload them through the interface.
+ 3. Provide at least 2 minutes of audio, preferably 5-10 minutes for better results.
+ 4. Maximum Audio Length will ensure the audio samples do not exceed the set value in seconds. The AI model can only process so much audio per audio sample, per epoch.
+ 5. Enter a unique Project Name if you wish to create a dedicated training folder. You will need to specify this name as you move through the interfaces and click the Refresh buttons where needed.
+ 6. Select the Whisper model (large-v2 recommended for English use cases).
+ 7. Choose your Dataset Language.
+ 8. Adjust the Evaluation Split percentage (default 15% is suitable for most cases).
+ 9. Optionally enable the BPE Tokenizer for custom tokenization, which can improve the model's ability to handle unique words, accents, or languages by creating a vocabulary tailored to your specific dataset. This is especially useful for non-standard speech patterns or languages with complex morphology.
+ + #### 🟨 What This Step Does + 1. Uses Whisper to transcribe and segment your audio files.
+ 2. Creates smaller audio clips and corresponding transcriptions.
+ 3. Generates two CSV files: `metadata_train.csv` and `metadata_eval.csv`.
+ 4. Saves processed audio in `/finetune/[project_name]/wavs/`.
+ 5. Stores CSV files and metadata in `/finetune/[project_name]/`.
+ + Key files generated:
+ β—½ `/finetune/[project_name]/lang.txt`: Contains the two-letter language code.
+ β—½ `/finetune/[project_name]/metadata_train.csv`: List of audio files and transcriptions for training.
+ β—½ `/finetune/[project_name]/metadata_eval.csv`: Evaluation data for assessing training quality.
+ + #### Important Considerations + β—½ For multi-speaker audio, manually edit CSV files to remove other speakers.
+ β—½ You can manually edit CSV files to improve transcription accuracy & use the `Dataset Validation` to help identify possible bad transcriptions.
+ β—½ The Whisper 2 model may provide better results in audio splitting and dataset creation for English languages, however please see the Whisper site for more details on other languages.
+ β—½ Each Whisper model is about 3GB. It's recommended to stick with one model to save disk space.
""" ) + + with gr.Tab("Generate Dataset Notes"): + gr.Markdown( + f""" + #### πŸŸ₯ Important Notes + β—½ **Windows Users:** If you encounter a `UserWarning: huggingface_hub cache-system uses symlinks` error, restart your Windows command prompt with `Run as Administrator` and relaunch the finetuning process. This typically only occurs during the first Whisper model download.
+ β—½ **Language Support:** Whisper's performance varies across languages. The Large-v3 model may offer better results for certain languages. For problematic languages, consider manual dataset creation.
+ + #### 🟩 Processing Time + β—½ First run: Requires downloading the 3GB Whisper model.
+ β—½ Subsequent runs: A few minutes on an average 3-4 year old system.
+ + After processing, use the 'Dataset Validation' tab to check transcription accuracy and make necessary corrections.
+ + For more details on Whisper models or manual dataset creation, refer to:
+ β—½ [Whisper Models Information](https://github.com/openai/whisper?tab=readme-ov-file#available-models-and-languages)
+ β—½ [Manual Dataset Creation Guide](https://docs.coqui.ai/en/latest/formatting_your_dataset.html) + """ + ) + + ####################### #### GRADIO STEP 2 #### @@ -1752,26 +1776,19 @@ if __name__ == "__main__": label="Project Name", value="personsname", visible=True, - info='Create Unique Training Folder', scale=1, ) with gr.Group(): continue_run = gr.Checkbox( value=False, - label="Continue Previous", + label="Continue Previous Project", scale=1, ) disable_shared_memory = gr.Checkbox( value=False, - label="Disable Shared Memory", + label="Disable Shared Memory Use", scale=1, ) - model_to_train_choices = list(available_models.keys()) - model_to_train = gr.Dropdown( - choices=model_to_train_choices, - label="Select the Model to train", - value=model_to_train_choices[0] if model_to_train_choices else None - ) train_csv = gr.Textbox( label="Train CSV:", scale=2, @@ -1781,6 +1798,14 @@ if __name__ == "__main__": scale=2, ) with gr.Row(): + model_to_train_choices = list(available_models.keys()) + model_to_train = gr.Dropdown( + choices=model_to_train_choices, + label="Select the Model to train", + value=model_to_train_choices[0] if model_to_train_choices else None, + scale=2, + ) + learning_rates = gr.Dropdown( value=5e-6, label="Learning Rate", @@ -1798,7 +1823,7 @@ if __name__ == "__main__": scale=1, ) learning_rate_scheduler = gr.Dropdown( - value="None", + value="CosineAnnealingWarmRestarts", label="Learning Rate Scheduler(s)", choices=[ ("None", "None"), @@ -1850,7 +1875,7 @@ if __name__ == "__main__": grad_acumm = gr.Slider( label="Grad accumulation steps:", minimum=1, - maximum=1024, + maximum=32, step=1, value=args.grad_acumm, ) @@ -1867,7 +1892,8 @@ if __name__ == "__main__": ) with gr.Row(): - train_time = gr.Label("Estimated Total Training Time", show_label=False) + train_time = gr.Label("Estimated Total Training Time", show_label=False, scale=2) + train_btn = gr.Button(value="Step 2 - Run the training", scale=1) with gr.Row(): model_data = gr.Image(c_logger.plot_metrics(), show_label=False) @@ -1879,7 +1905,6 @@ if __name__ == "__main__": ) demo.load(load_metrics, None, [model_data, train_time], every=1) demo.load(read_logs, None, logs_tts_train, every=1) - train_btn = gr.Button(value="Step 2 - Run the training") def train_model(language, train_csv, eval_csv, learning_rates, model_to_train, num_epochs, batch_size, grad_acumm, max_audio_length, speaker_name_input_training, continue_run, disable_shared_memory, learning_rate_scheduler, optimizer, progress=gr.Progress()): clear_gpu_cache() @@ -1900,24 +1925,30 @@ if __name__ == "__main__": learning_rate = float(learning_rates) # Convert the learning rate value to a float progress(0, "Initializing training...") config_path, original_xtts_checkpoint, vocab_file, exp_path, speaker_wav = train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, learning_rate, model_to_train, continue_run, disable_shared_memory, learning_rate_scheduler, optimizer, max_audio_length=max_audio_length, progress=gr.Progress()) - except: - traceback.print_exc() - error = traceback.format_exc() - return f"The training was interrupted due an error !! Please check the console to check the full error message! \n Error summary: {error}", "", "", "", "" + + # copy original files to avoid parameters changes issues + shutil.copy(config_path, exp_path) + shutil.copy(vocab_file, exp_path) - # copy original files to avoid parameters changes issues - shutil.copy(config_path, exp_path) - shutil.copy(vocab_file, exp_path) + ft_xtts_checkpoint = os.path.join(exp_path, "best_model.pth") + print("[FINETUNE] Model training done. Move to Step 3") + clear_gpu_cache() + return "Model training done. Move to Step 3", config_path, vocab_file, ft_xtts_checkpoint, speaker_wav, speaker_name_input_training - ft_xtts_checkpoint = os.path.join(exp_path, "best_model.pth") - print("[FINETUNE] Model training done. Move to Step 3") - clear_gpu_cache() - return "Model training done. Move to Step 3", config_path, vocab_file, ft_xtts_checkpoint, speaker_wav, speaker_name_input_training + except ValueError as ve: + # This will catch the specific ValueError we raised in train_gpt + error_message = str(ve) + print(f"[FINETUNE] Error: {error_message}") + return f"Training error: {error_message}", "", "", "", "", "" + except Exception as e: + # This will catch any other unexpected errors + error_message = f"An unexpected error occurred: {str(e)}" + print(f"[FINETUNE] Error: {error_message}") + return f"Training error: {error_message}", "", "", "", "", "" with gr.Tab("Training the model Instructions"): gr.Markdown( f""" - ### πŸ’» Training
### πŸŸ₯ Important Note - Language support. β—½ If this step is failing/erroring you may wish to check your training data was created correctly (Detailed in Step 1), confirming that wav files have been generated and your `metadata_train.csv` and `metadata_eval.csv` files have been populated.
β—½ Ignore 'fatal: not a git repository (or any of the parent directories): .git' as this is a legacy training script issue.
@@ -1928,6 +1959,23 @@ if __name__ == "__main__":          - What you are training (A human voice, A cartoon voice, A new language entirely etc)
         - How much audio you have (you may want less or more eval or epochs)
β—½ The key indicator of sufficient model training is whether it sounds right to you. Coqui suggests beginning with the base settings for training. If the resulting audio doesn't meet your expectations, additional training sessions may be necessary.
+ #### πŸŸͺ Continue Previous Project + This option allows you to resume training from where it left off if your previous session was interrupted (e.g., due to a power outage or system crash). When enabled:
+ β—½ The system will attempt to load the most recent checkpoint from your previous training session.
+ β—½ It will continue training from that point, preserving your progress and saving time.
+ β—½ This is particularly useful for long training sessions or when you want to extend training on a previously trained model.
+ #### πŸŸͺ Disable Shared Memory + This option affects how the system manages memory during training:
+ β—½ When enabled, it limits GPU memory usage to 95% of its capacity.
+ β—½ This prevents the training process from spilling over into your system's shared memory.
+ β—½ You're more likely to get a clear "out of memory" error if your GPU runs out of memory, rather than experiencing slower performance due to memory spillover.
+ β—½ This option is useful if you prefer to keep training confined to GPU memory for better performance and clearer error reporting.
+ + Important Note for Windows Users with Low VRAM GPUs:
+ β—½ If you're using Windows and have a GPU with lower VRAM (e.g., 8GB or less), it's recommended to leave this option disabled.
+ β—½ Windows can utilize system RAM as extended VRAM, which can be crucial for completing the training process on lower VRAM GPUs.
+ β—½ Disabling this option allows Windows to manage memory allocation more flexibly, potentially preventing out-of-memory errors on GPUs with limited VRAM. + ### 🟨 What this step is doing
β—½ Very simply put, it's taking all our wav files generated in Step 1, along with our recorded speech in out excel documents and its training the model on that voice e.g. listen to this audio file and this is what is being said in it, so learn to reproduce it.
### 🟩 How long will this take?
@@ -1972,13 +2020,13 @@ if __name__ == "__main__": gr.Markdown( f""" ### 🟩 Learning Rate
- β—½ 1e-6: Very small learning rate, slow but stable learning.
- β—½ 5e-6: Small learning rate, slow but stable learning.
- β—½ 1e-5: Moderate learning rate, balanced between stability and convergence speed.
- β—½ 5e-5: Higher learning rate, faster convergence but potential instability.
- β—½ 1e-4: High learning rate, faster convergence but increased risk of instability.
- β—½ 5e-4: Very high learning rate, fast convergence but higher risk of instability.
- β—½ 1e-3: Extremely high learning rate, very fast convergence but high risk of instability.

+ β—½ **1e-6:** Very small learning rate, slow but stable learning.
+ β—½ **5e-6:** Small learning rate, slow but stable learning.
+ β—½ **1e-5:** Moderate learning rate, balanced between stability and convergence speed.
+ β—½ **5e-5:** Higher learning rate, faster convergence but potential instability.
+ β—½ **1e-4:** High learning rate, faster convergence but increased risk of instability.
+ β—½ **5e-4:** Very high learning rate, fast convergence but higher risk of instability.
+ β—½ **1e-3:** Extremely high learning rate, very fast convergence but high risk of instability.

The optimal learning rate depends on the model architecture, dataset, and other hyperparameters. """ ) @@ -1986,15 +2034,15 @@ if __name__ == "__main__": gr.Markdown( f""" ### 🟩 Learning Rate Schedulers
- β—½ None: Learning rate remains unchanged throughout the training.
- β—½ Step: Decreases the learning rate by a factor at specified intervals, offering a straightforward method for controlling learning rate decay.
- β—½ MultiStep: Decreases the learning rate by a factor at specified milestones, useful for gradually reducing the learning rate during training.
- β—½ Exponential: Decreases the learning rate exponentially over epochs, providing a simple decay mechanism for training.
- β—½ CosineAnnealing: Adjusts the learning rate following a cosine annealing schedule, facilitating smoother optimization and potentially better convergence.
- β—½ Reduce On Plateau: Reduces the learning rate when a monitored metric has stopped improving, helping to fine-tune learning rates during training.
- β—½ Cyclic: Cycles the learning rate between two boundary values, allowing for dynamic learning rate schedules during training.
- β—½ One Cycle (disabled): Sets the learning rate using a cyclical learning rate policy and varies momentum inversely, potentially accelerating training and improving performance.
- β—½ Cosine Annealing Warm Restarts: Adjusts the learning rate following a cosine annealing schedule with warm restarts, offering a balance between exploration and exploitation during optimization.

+ β—½ **None:** Learning rate remains unchanged throughout the training.
+ β—½ **Step:** Decreases the learning rate by a factor at specified intervals, offering a straightforward method for controlling learning rate decay.
+ β—½ **MultiStep:** Decreases the learning rate by a factor at specified milestones, useful for gradually reducing the learning rate during training.
+ β—½ **Exponential:** Decreases the learning rate exponentially over epochs, providing a simple decay mechanism for training.
+ β—½ **CosineAnnealing:** Adjusts the learning rate following a cosine annealing schedule, facilitating smoother optimization and potentially better convergence.
+ β—½ **Reduce On Plateau:** Reduces the learning rate when a monitored metric has stopped improving, helping to fine-tune learning rates during training.
+ β—½ **Cyclic:** Cycles the learning rate between two boundary values, allowing for dynamic learning rate schedules during training.
+ β—½ **One Cycle (disabled):** Sets the learning rate using a cyclical learning rate policy and varies momentum inversely, potentially accelerating training and improving performance.
+ β—½ **Cosine Annealing Warm Restarts:** Adjusts the learning rate following a cosine annealing schedule with warm restarts, offering a balance between exploration and exploitation during optimization.

The choice of learning rate scheduler depends on the specific characteristics of the model, dataset, and training objectives. """ ) @@ -2002,12 +2050,12 @@ if __name__ == "__main__": gr.Markdown( f""" ### 🟩 Optimizers
- β—½ SGD with Momentum: Good for convergence, requires careful tuning of the learning rate and momentum. - β—½ RMSprop: Adaptively adjusts learning rates, good for varying gradient magnitudes. - β—½ Adam: Widely used, performs well with default parameters. - β—½ Adagrad: Good for sparse data, adapts learning rates based on historical gradients. - β—½ AdamW: Similar to Adam but with decoupled weight decay for better regularization. - When picking an optimizer, you need to experiment with different options to find the one that achieves the best performance and convergence for your task. + β—½ **SGD with Momentum:** Good for convergence, requires careful tuning of the learning rate and momentum.
+ β—½ **RMSprop:** Adaptively adjusts learning rates, good for varying gradient magnitudes.
+ β—½ **Adam:** Widely used, performs well with default parameters.
+ β—½ **Adagrad:** Good for sparse data, adapts learning rates based on historical gradients.
+ β—½ **AdamW:** Similar to Adam but with decoupled weight decay for better regularization.

+ When picking an optimizer, you need to experiment with different options to find the one that achieves the best performance and convergence for your task.

AdamW has been set to default as its a sensible pick that will handle most needs. """ ) @@ -2091,10 +2139,9 @@ if __name__ == "__main__": scale=1, ) speaker_name_input_testing = gr.Textbox( - label="Project Name", + label="Project Name (Refresh Dropdowns on change)", value="personsname", visible=True, - info='Create Unique Training Folder', scale=1, ) @@ -2199,10 +2246,9 @@ if __name__ == "__main__": scale=2, ) speaker_name_input_export = gr.Textbox( - label="Project Name", + label="Project Name (Refresh Dropdowns on change)", value="personsname", visible=True, - info='Unique Training Folder', scale=1, ) # Create refresh button From e347f1504624c15d8a9e78f07a099460bdd8fcbf Mon Sep 17 00:00:00 2001 From: bmalaski <5288648+bmalaski@users.noreply.github.com> Date: Tue, 2 Jul 2024 00:14:54 -0400 Subject: [PATCH 5/7] add warmup --- finetune.py | 22 +++++++++++++------- trainer_alltalk/WarmUpScheduler.py | 33 ++++++++++++++++++++++++++++++ trainer_alltalk/trainer.py | 15 ++++++++++++++ 3 files changed, 63 insertions(+), 7 deletions(-) create mode 100644 trainer_alltalk/WarmUpScheduler.py diff --git a/finetune.py b/finetune.py index 48906f9..f0e8dd5 100644 --- a/finetune.py +++ b/finetune.py @@ -704,7 +704,7 @@ def basemodel_or_finetunedmodel_choice(value): elif value == "Existing finetuned model": basemodel_or_finetunedmodel = False -def train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, learning_rate, model_to_train, continue_run, disable_shared_memory, learning_rate_scheduler, optimizer, max_audio_length=255995, progress=gr.Progress()): +def train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, learning_rate, model_to_train, continue_run, disable_shared_memory, learning_rate_scheduler, optimizer, warm_up, max_audio_length=255995, progress=gr.Progress()): pfc_check_fail() if "No Models Available" in model_to_train: @@ -983,6 +983,8 @@ def train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, eval_samples=eval_samples, training_assets=training_assets, c_logger=c_logger, + warmup=warm_up, + ) if(disable_shared_memory): @@ -1789,6 +1791,11 @@ if __name__ == "__main__": label="Disable Shared Memory Use", scale=1, ) + warm_up = gr.Checkbox( + value=False, + label="Perform Small a Warmup Learning", + scale=1, + ) train_csv = gr.Textbox( label="Train CSV:", scale=2, @@ -1875,7 +1882,7 @@ if __name__ == "__main__": grad_acumm = gr.Slider( label="Grad accumulation steps:", minimum=1, - maximum=32, + maximum=128, step=1, value=args.grad_acumm, ) @@ -1906,7 +1913,7 @@ if __name__ == "__main__": demo.load(load_metrics, None, [model_data, train_time], every=1) demo.load(read_logs, None, logs_tts_train, every=1) - def train_model(language, train_csv, eval_csv, learning_rates, model_to_train, num_epochs, batch_size, grad_acumm, max_audio_length, speaker_name_input_training, continue_run, disable_shared_memory, learning_rate_scheduler, optimizer, progress=gr.Progress()): + def train_model(language, train_csv, eval_csv, learning_rates, model_to_train, num_epochs, batch_size, grad_acumm, max_audio_length, speaker_name_input_training, continue_run, disable_shared_memory, learning_rate_scheduler, optimizer, warm_up, progress=gr.Progress()): clear_gpu_cache() global out_path if speaker_name_input_training and speaker_name_input_training != 'personsname': @@ -1914,8 +1921,8 @@ if __name__ == "__main__": if not train_csv or not eval_csv: if (out_path / "metadata_eval.csv").exists() and (out_path / "metadata_train.csv").exists(): - train_csv = out_path / "metadata_eval.csv" - eval_csv = out_path / "metadata_train.csv" + train_csv = out_path / "metadata_train.csv" + eval_csv = out_path / "metadata_eval.csv" print("Using existing metadata and training csv.") else: return "You need to run the data processing step or manually set `Train CSV` and `Eval CSV` fields !", "", "", "", "" @@ -1924,7 +1931,7 @@ if __name__ == "__main__": max_audio_length = int(max_audio_length * 22050) learning_rate = float(learning_rates) # Convert the learning rate value to a float progress(0, "Initializing training...") - config_path, original_xtts_checkpoint, vocab_file, exp_path, speaker_wav = train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, learning_rate, model_to_train, continue_run, disable_shared_memory, learning_rate_scheduler, optimizer, max_audio_length=max_audio_length, progress=gr.Progress()) + config_path, original_xtts_checkpoint, vocab_file, exp_path, speaker_wav = train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, learning_rate, model_to_train, continue_run, disable_shared_memory, learning_rate_scheduler, optimizer, warm_up, max_audio_length=max_audio_length, progress=gr.Progress()) # copy original files to avoid parameters changes issues shutil.copy(config_path, exp_path) @@ -2304,7 +2311,8 @@ if __name__ == "__main__": continue_run, disable_shared_memory, learning_rate_scheduler, - optimizer + optimizer, + warm_up ], outputs=[progress_train, xtts_config, xtts_vocab, xtts_checkpoint, speaker_reference_audio, speaker_name_input_testing], ) diff --git a/trainer_alltalk/WarmUpScheduler.py b/trainer_alltalk/WarmUpScheduler.py new file mode 100644 index 0000000..9406c82 --- /dev/null +++ b/trainer_alltalk/WarmUpScheduler.py @@ -0,0 +1,33 @@ +from torch.optim.lr_scheduler import LRScheduler + +class WarmUpScheduler(LRScheduler): + def __init__(self, optimizer, total_epochs, warmup_lr, target_lr, after_scheduler=None): + self.total_epochs = total_epochs + self.warmup_epochs = int(total_epochs * 0.1) + self.warmup_lr = warmup_lr + self.target_lr = target_lr + self.after_scheduler = after_scheduler + print(f"[WARMUP] {self.warmup_epochs} epochs") + super(WarmUpScheduler, self).__init__(optimizer) + + def get_lr(self): + if self.last_epoch < self.warmup_epochs: + progress = self.last_epoch / self.warmup_epochs + return [self.warmup_lr + progress * (self.target_lr - self.warmup_lr) for _ in self.base_lrs] + if self.after_scheduler: + return self.after_scheduler.get_last_lr() + return [base_lr for base_lr in self.base_lrs] + + def step(self, epoch=None): + if epoch is not None and epoch < self.warmup_epochs: + self.last_epoch = epoch + elif self.after_scheduler and epoch is not None: + if epoch < self.warmup_epochs: + self.last_epoch = epoch + else: + self.after_scheduler.step(epoch - self.warmup_epochs) + self.last_epoch = epoch + else: + super(WarmUpScheduler, self).step(epoch) + if self.after_scheduler and self.last_epoch >= self.warmup_epochs: + self.after_scheduler.step(epoch - self.warmup_epochs if epoch is not None else None) \ No newline at end of file diff --git a/trainer_alltalk/trainer.py b/trainer_alltalk/trainer.py index e786a2e..9b345ec 100644 --- a/trainer_alltalk/trainer.py +++ b/trainer_alltalk/trainer.py @@ -58,6 +58,8 @@ from trainer.utils.distributed import ( rank_zero_only, ) +from trainer_alltalk.WarmUpScheduler import WarmUpScheduler + logger = logging.getLogger("trainer") if is_apex_available(): @@ -309,6 +311,7 @@ class Trainer: parse_command_line_args: bool = True, callbacks: Dict[str, Callable] = {}, gpu: int = None, + warmup: bool = False, ) -> None: """Simple yet powerful πŸΈπŸ’¬ TTS trainer for PyTorch. It can train all the available `tts` and `vocoder` models or easily be customized. @@ -429,6 +432,7 @@ class Trainer: self.overfit_batch = args.overfit_batch self.skip_train_epoch = args.skip_train_epoch self.start_with_eval = args.start_with_eval + self.warmup = warmup assert self.grad_accum_steps > 0, " [!] grad_accum_steps must be greater than 0." @@ -569,6 +573,16 @@ class Trainer: self.scheduler, self.args, self.config, self.restore_epoch, self.restore_step ) + if self.warmup: + warmup_scheduler = WarmUpScheduler( + optimizer=self.optimizer, + total_epochs=self.config.epochs, + warmup_lr=1e-7, + target_lr=self.config.lr, + after_scheduler=self.scheduler + ) + self.scheduler = warmup_scheduler + # DISTRIBUTED if self.use_pt_ddp: self.model = DDP_th(self.model, device_ids=[args.rank], output_device=args.rank) @@ -1542,6 +1556,7 @@ class Trainer: self.model.train() torch.set_grad_enabled(True) + epoch_time = time.time() - epoch_start_time self.callbacks.on_train_epoch_end(self) From 6a992cb8f4a87a59786479efc83c5aa373cf4e6a Mon Sep 17 00:00:00 2001 From: bmalaski <5288648+bmalaski@users.noreply.github.com> Date: Wed, 3 Jul 2024 09:59:25 -0400 Subject: [PATCH 6/7] Update WarmUpScheduler.py --- trainer_alltalk/WarmUpScheduler.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/trainer_alltalk/WarmUpScheduler.py b/trainer_alltalk/WarmUpScheduler.py index 9406c82..d6d1d20 100644 --- a/trainer_alltalk/WarmUpScheduler.py +++ b/trainer_alltalk/WarmUpScheduler.py @@ -3,7 +3,9 @@ from torch.optim.lr_scheduler import LRScheduler class WarmUpScheduler(LRScheduler): def __init__(self, optimizer, total_epochs, warmup_lr, target_lr, after_scheduler=None): self.total_epochs = total_epochs - self.warmup_epochs = int(total_epochs * 0.1) + self.warmup_epochs = int(total_epochs * 0.05) + if self.warmup_epochs <= 0: + self.warmup_epochs = 1 self.warmup_lr = warmup_lr self.target_lr = target_lr self.after_scheduler = after_scheduler From d5c20262a8c9ecb106bb6ccb246d1e855210cd4a Mon Sep 17 00:00:00 2001 From: bmalaski <5288648+bmalaski@users.noreply.github.com> Date: Wed, 3 Jul 2024 09:59:25 -0400 Subject: [PATCH 7/7] Update WarmUpScheduler.py