Machine learning

Layers
- Base layers
  - BaseModel
- Spectral gate
  - SpectralGate
    - SpectralGate.forward()
    - SpectralGate.training

ML Accelerator

class audiotools.ml.accelerator.Accelerator(amp: bool = False)[source]

Bases: object

This class is used to prepare models and dataloaders for usage with DDP or DP. Use the functions prepare_model, prepare_dataloader to prepare the respective objects. In the case of models, they are moved to the appropriate GPU and SyncBatchNorm is applied to them. In the case of dataloaders, a sampler is created and the dataloader is initialized with that sampler.

If the world size is 1, prepare_model and prepare_dataloader are no-ops. If the environment variable LOCAL_RANK is not set, then the script was launched without torchrun, and DataParallel will be used instead of DistributedDataParallel (not recommended), if the world size (number of GPUs) is greater than 1.

Parameters: amp (bool, optional) – Whether or not to enable automatic mixed precision, by default False

autocast(*args, **kwargs)[source]: Context manager for autocasting. Arguments go to torch.cuda.amp.autocast.

backward(loss: Tensor)[source]

Backwards pass, after scaling the loss if amp is enabled.

Parameters: loss (torch.Tensor) – Loss value.

prepare_dataloader(dataset: Iterable, start_idx: Optional[int] = None, **kwargs)[source]

Wraps a dataset with a DataLoader, using the correct sampler if DDP is enabled.

Parameters

dataset (Iterable) – Dataset to build Dataloader around.
start_idx (int, optional) – Start index of sampler, useful if resuming from some epoch, by default None

Returns

_description_

Return type

_type_

prepare_model(model: Module, **kwargs)[source]

Prepares model for DDP or DP. The model is moved to the device of the correct rank.

Parameters: model (torch.nn.Module) – Model that is converted for DDP or DP.
Returns: Wrapped model, or original model if DDP and DP are turned off.
Return type: torch.nn.Module

step(optimizer: Optimizer)[source]

Steps the optimizer, using a scaler if amp is enabled.

Parameters: optimizer (torch.optim.Optimizer) – Optimizer to step forward.

static unwrap(model)[source]: Unwraps the model if it was wrapped in DDP or DP, otherwise just returns the model. Use this to unwrap the model returned by audiotools.ml.accelerator.Accelerator.prepare_model().

update()[source]: Updates the scale factor.

ML Loggers / decorators

class audiotools.ml.decorators.Mean[source]

Bases: object

Keeps track of the running mean, along with the latest value.

reset()[source]

update(val)[source]

class audiotools.ml.decorators.Tracker(writer: Optional[SummaryWriter] = None, log_file: Optional[str] = None, rank: int = 0, console_width: int = 100, step: int = 0)[source]

Bases: object

A tracker class that helps to monitor the progress of training and logging the metrics.

metrics

A dictionary containing the metrics for each label.

Type: dict

history

A dictionary containing the history of metrics for each label.

Type: dict

writer

A SummaryWriter object for logging the metrics.

Type: SummaryWriter

rank

The rank of the current process.

Type: int

step

The current step of the training.

Type: int

tasks

A dictionary containing the progress bars and tables for each label.

Type: dict

pbar

A progress bar object for displaying the progress.

Type: Progress

consoles

A list of console objects for logging.

Type: list

live

A Live object for updating the display live.

Type: Live

print(msg: str)[source]: Prints the given message to all consoles.

update(label: str, fn_name: str)[source]: Updates the progress bar and table for the given label.

done(label: str, title: str)[source]: Resets the progress bar and table for the given label and prints the final result.

track(label: str, length: int, completed: int = 0, op: dist.ReduceOp = dist.ReduceOp.AVG, ddp_active: bool = "LOCAL_RANK" in os.environ)[source]: A decorator for tracking the progress and metrics of a function.

log(label: str, value_type: str = 'value', history: bool = True)[source]: A decorator for logging the metrics of a function.

is_best(label: str, key: str) → bool[source]: Checks if the latest value of the given key in the label is the best so far.

state_dict() → dict[source]: Returns a dictionary containing the state of the tracker.

load_state_dict(state_dict: dict) → Tracker[source]: Loads the state of the tracker from the given state dictionary.

done(label: str, title: str)[source]

Resets the progress bar and table for the given label and prints the final result.

Parameters

label (str) – The label of the progress bar and table to be reset.
title (str) – The title to be displayed when printing the final result.

is_best(label, key)[source]

Checks if the latest value of the given key in the label is the best so far.

Parameters

label (str) – The label of the metrics to be checked.
key (str) – The key of the metric to be checked.

Returns

True if the latest value is the best so far, otherwise False.

Return type

bool

load_state_dict(state_dict)[source]

Loads the state of the tracker from the given state dictionary.

Parameters: state_dict (dict) – A dictionary containing the history and step of the tracker.
Returns: The tracker object with the loaded state.
Return type: Tracker

log(label: str, value_type: str = 'value', history: bool = True)[source]

A decorator for logging the metrics of a function.

Parameters

label (str) – The label to be associated with the logging.
value_type (str, optional) – The type of value to be logged, by default “value”.
history (bool, optional) – Whether to save the history of the metrics, by default True.

print(msg)[source]

Prints the given message to all consoles.

Parameters: msg (str) – The message to be printed.

state_dict()[source]

Returns a dictionary containing the state of the tracker.

Returns: A dictionary containing the history and step of the tracker.
Return type: dict

track(label: str, length: int, completed: int = 0, op: ~torch.distributed.distributed_c10d.ReduceOp = <RedOpType.AVG: 1>, ddp_active: bool = False)[source]

A decorator for tracking the progress and metrics of a function.

Parameters

label (str) – The label to be associated with the progress and metrics.
length (int) – The total number of iterations to be completed.
completed (int, optional) – The number of iterations already completed, by default 0.
op (dist.ReduceOp, optional) – The reduce operation to be used, by default dist.ReduceOp.AVG.
ddp_active (bool, optional) – Whether the DistributedDataParallel is active, by default “LOCAL_RANK” in os.environ.

update(label, fn_name)[source]

Updates the progress bar and table for the given label.

Parameters

label (str) – The label of the progress bar and table to be updated.
fn_name (str) – The name of the function associated with the label.

audiotools.ml.decorators.default_list()[source]

audiotools.ml.decorators.timer(prefix: str = 'time')[source]

Adds execution time to the output dictionary of the decorated function. The function decorated by this must output a dictionary. The key added will follow the form “[prefix]/[name_of_function]”

Parameters: prefix (str, optional) – The key added will follow the form “[prefix]/[name_of_function]”, by default “time”.

audiotools.ml.decorators.when(condition)[source]

Runs a function only when the condition is met. The condition is a function that is run.

Parameters: condition (Callable) – Function to run to check whether or not to run the decorated function.

Example

Checkpoint only runs every 100 iterations, and only if the local rank is 0.

>>> i = 0
>>> rank = 0
>>>
>>> @when(lambda: i % 100 == 0 and rank == 0)
>>> def checkpoint():
>>>     print("Saving to /runs/exp1")
>>>
>>> for i in range(1000):
>>>     checkpoint()

Experiment tracking/saving

Useful class for Experiment tracking, and ensuring code is saved alongside files.

class audiotools.ml.experiment.Experiment(exp_directory: str = 'runs/', exp_name: Optional[str] = None)[source]

Bases: object

This class contains utilities for managing experiments. It is a context manager, that when you enter it, changes your directory to a specified experiment folder (which optionally can have an automatically generated experiment name, or a specified one), and changes the CUDA device used to the specified device (or devices).

Parameters

exp_directory (str) – Folder where all experiments are saved, by default “runs/”.
exp_name (str, optional) – Name of the experiment, by default uses the current time, date, and hostname to save.

static generate_exp_name()[source]

Generates a random experiment name based on the date and a randomly generated adjective-noun tuple.

Returns: Randomly generated experiment name.
Return type: str

snapshot(filter_fn: ~typing.Callable = <function Experiment.<lambda>>)[source]

Captures a full snapshot of all the files tracked by git at the time the experiment is run. It also captures the diff against the committed code as a separate file.

Parameters: filter_fn (Callable, optional) – Function that can be used to exclude some files from the snapshot, by default accepts all files