Trainer API

The TTS.trainer.Trainer provides a lightweight, extensible, and feature-complete training run-time. We optimized it for 🐸 but can also be used for any DL training in different domains. It supports distributed multi-gpu, mixed-precision (apex or torch.amp) training.

Trainer

class TTS.trainer.Trainer(args, config, output_path, c_logger=None, dashboard_logger=None, model=None, get_model=None, get_data_samples=None, train_samples=None, eval_samples=None, cudnn_benchmark=False, training_assets={}, parse_command_line_args=True)

eval_epoch()

Main entry point for the evaluation loop. Run evaluation on the all validation samples.

eval_step(batch, step)

Perform a evaluation step on a batch of inputs and log the process.

Parameters

• batch (Dict) – Input batch.

• step (int) – Current step number in this epoch.

Returns

Model outputs and losses.

Return type

Tuple[Dict, Dict]

fit()

Where the ✨️magic✨️ happens…

format_batch(batch)

Format the dataloader output and return a batch.

Parameters

batch (List) – Batch returned by the dataloader.

Returns

Formatted batch.

Return type

Dict

static get_criterion(model)

Receive the criterion from the model. Model must implement get_criterion().

Parameters

model (nn.Module) – Training model.

Returns

Criterion layer.

Return type

nn.Module

static get_lr(model, config)

Set the initial learning rate by the model if model implements get_lr() else try setting the learning rate fromthe config.

Parameters

• model (nn.Module) – Training model.

• config (Coqpit) – Training configuration.

Returns

A single learning rate or a list of learning rates, one for each optimzier.

Return type

Union[float, List[float]]

static get_optimizer(model, config)

Receive the optimizer from the model if model implements get_optimizer() else check the optimizer parameters in the config and try initiating the optimizer.

Parameters

• model (nn.Module) – Training model.

• config (Coqpit) – Training configuration.

Returns

A optimizer or a list of optimizers. GAN models define a list.

Return type

Union[torch.optim.Optimizer, List]

static get_scheduler(model, config, optimizer)

Receive the scheduler from the model if model implements get_scheduler() else check the config and try initiating the scheduler.

Parameters

• model (nn.Module) – Training model.

• config (Coqpit) – Training configuration.

Returns

A scheduler or a list of schedulers, one for each optimizer.

Return type

Union[torch.optim.Optimizer, List]

get_train_dataloader(training_assets, data_items, verbose)

Initialize and return a training data loader.

Parameters

• ap (AudioProcessor) – Audio processor.

• data_items (List) – Data samples used for training.

• verbose (bool) – enable/disable printing loader stats at initialization.

Returns

Initialized training data loader.

Return type

DataLoader

init_training(args, coqpit_overrides, config=None)

Initialize training and update model configs from command line arguments.

Parameters

• args (argparse.Namespace or dict like) – Parsed input 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 paramaters.

Return type

c (TTS.utils.io.AttrDict)

static master_params(optimizer)

Generator over parameters owned by the optimizer.

Used to select parameters used by the optimizer for gradient clipping.

Parameters

optimizer – Target optimizer.

static parse_argv(args)

Parse command line arguments to init or override TrainingArgs().

restore_model(config, restore_path, model, optimizer, scaler=None)

Restore training from an old run. It restores model, optimizer, AMP scaler and training stats.

Parameters

• 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

[description]

Return type

Tuple[nn.Module, torch.optim.Optimizer, torch.cuda.amp.GradScaler, int]

static run_get_model(config, get_model)

Run the get_model function and return the model.

Parameters

config (Coqpit) – Model config.

Returns

initialized model.

Return type

nn.Module

save_best_model()

Save the best model. It only saves if the current target loss is smaller then the previous.

test_run()

Run test and log the results. Test run must be defined by the model. Model must return figures and audios to be logged by the Tensorboard.

train_epoch()

Main entry point for the training loop. Run training on the all training samples.

train_step(batch, batch_n_steps, step, loader_start_time)

Perform a training step on a batch of inputs and log the process.

Parameters

• 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

Model outputs and losses.

Return type

Tuple[Dict, Dict]

TrainingArgs

class TTS.trainer.TrainingArgs(continue_path='', restore_path='', best_path='', skip_train_epoch=False, config_path='', rank=0, group_id='', use_ddp=False)

Trainer arguments to be defined externally. It helps integrating the Trainer with the higher level APIs and set the values for distributed training.