Running experiments

The Experiment base class and the scripts/run.py entry point together provide a complete, configuration-driven training loop for experiments with libdamp. The main goal is to reduce training boilerplate code while maximizing flexibility of the individual experiment setup and ensuring reproducibility.

The Experiment class

Experiment is a Lightning LightningModule subclass that defines the training/validation/test step structure common to all libdamp experiments. To define a new experiment, subclass it and implement the usual Lightning hooks (forward(), training_step(), validation_step(), configure_optimizers(), …), composing the model out of libdamp.generators, libdamp.processors, and libdamp.models:

import gin
import libdamp

@gin.configurable
class MyExperiment(libdamp.Experiment):
    def __init__(self, fs, **kwargs):
        super().__init__(**kwargs)
        self.fs = fs
        self.osc = libdamp.generators.HarmonicOsc(N=512, fs=fs)
        self.envelope = libdamp.processors.GainEnvelope()
        self.loss_fn = libdamp.RMSLoss()

    def forward(self, f0, amplitudes, gain):
        self.osc.update(f0=f0, a=amplitudes)
        x = self.osc.generate()
        self.envelope.update(g=gain)
        return self.envelope.process(x)

    def training_step(self, batch, batch_idx):
        y, f0, amplitudes, gain = batch
        y_hat = self(f0, amplitudes, gain)
        loss = self.loss_fn(y, y_hat)
        self.log("train_loss", loss)
        return loss

All constructor parameters of Experiment itself (batch size, number of epochs, checkpointing, logging, …) are marked gin.REQUIRED and are meant to be set in a gin-config configuration file rather than hardcoded, so the same experiment class can be reused across many training runs that only differ in configuration. log_audio() is a ready-made helper for logging example audio (predictions and, once, the reference) to disk and optionally MLflow.

Running an experiment with run.py

scripts/run.py is the command-line entry point that turns an Experiment subclass and a gin config into a full training run. The gin config selects which experiment and datasets to use via the libdamp() binding:

libdamp.experiment = @MyExperiment()
libdamp.train_dataset = @MyTrainDataset()
libdamp.val_dataset = @MyValDataset()    # optional
libdamp.test_dataset = @MyTestDataset()  # optional

MyExperiment.fs = 16000.0
# ... plus all the gin.REQUIRED Experiment parameters (batch_size, max_epochs, save_path, ...)

and is then run with:

python scripts/run.py --config path/to/config.gin --seed 0

Command-line options:

  • -c, --config (required): one or more gin config files, merged in order.

  • --config-path: additional directories to search for gin files included from a config (via gin’s include statement).

  • --seed: fixed random seed for reproducible runs (default: 0).

Given the resolved configuration, run.py then takes care of the rest of the training boilerplate so individual experiments don’t have to:

  • builds the train/validation DataLoader\ s from the configured datasets;

  • sets up logging (CSV always, plus TensorBoard and/or MLflow if enabled) and writes the fully-resolved (“operative”) gin config alongside the run for reproducibility;

  • configures checkpointing (best-n and/or last-epoch) and early stopping based on val_loss, and optionally resumes from a previous checkpoint’s weights;

  • runs training via a Lightning Trainer configured from the experiment’s parameters (accelerator, devices, gradient accumulation/clipping, validation interval, fast_dev_run for quick smoke tests, …);

  • runs a final test pass with both the best and the last checkpoint, if a test dataset was given.

API reference

Base class for an experiment that can be run with run.py

class libdamp.experiment.Experiment(accelerator=<object object>, accumulate_grad_batches=<object object>, batch_size=<object object>, devices=<object object>, early_stopping=<object object>, early_stopping_patience=<object object>, enable_tensorboard=<object object>, enable_mlflow=<object object>, enable_audio_logging=<object object>, audio_log_batch_idx=<object object>, audio_log_item_idx=<object object>, fast_dev_run=<object object>, full_initial_checkpoint=<object object>, gradient_clip_val=<object object>, limit_val_batches=<object object>, load_checkpoint_weights=<object object>, max_epochs=<object object>, mlflow_tracking_uri=<object object>, name=<object object>, profiler=<object object>, save_best_n_ckpts=<object object>, save_last_epoch_ckpt=<object object>, save_path=<object object>, shuffle_data=<object object>, val_check_interval=<object object>)[source]

Bases: LightningModule

Base class for an experiment that can be run with run.py.

Defaults for all properties are set in configs/base.gin.

property name

Replace ‘{date}’ in name with the current date and time

log_audio(name: str, prediction, reference=None, current_batch_idx=None, split='val')[source]

Helper function to log audio prediction and reference (if provided) after each training, validation or testing epoch. The item to log is specified by self.audio_log_item_idx and self.audio_log_batch_idx. Audio files are saved to self.save_path / self.name / “audio” and optionally logged to MLflow if self.enable_mlflow is True.