auraloss

To evaluate the pre-trained models on the test set, use the examples/test_comp.py script. This requires the SignalTrain LA2A dataset and the provided model checkpoints (downloaded as a .tgz and extracted).

Hardware Requirements:

• Approximately 12 GB of VRAM (based on a batch size of 128 and patch length of 262,144 samples at 44.1 kHz).
• It is recommended to use half precision (--precision 16) as the models were trained in half precision.

Key Arguments:

• --root_dir: Path to the extracted SignalTrain LA2A Dataset.
• --logdir: Path to the extracted model checkpoints.
• --eval_subset: Set to "test" for evaluation.
• --preload: Set to True to load audio files into RAM for faster processing.

python examples/test_comp.py \
--root_dir /path/to/SignalTrain_LA2A_Dataset_1.1 \
--logdir path/to/checkpoints/version_9 \
--batch_size 128 \
--sample_rate 44100 \
--eval_subset "test" \
--eval_length 262144 \
--num_workers 8 \
--gpus 1 \
--shuffle False \
--precision 16 \
--preload True

You can retrain the compressor models using the examples/train_comp.py script. Training time varies by hardware; for example, training one of the six models for 20 epochs takes approximately 6.5 hours on an NVIDIA Quadro RTX 6000.

Key Arguments:

• --root_dir: Path to the extracted SignalTrain LA2A Dataset.
• --max_epochs: Number of training epochs.
• --train_length: Length of the training audio patches.
• --eval_length: Length of the evaluation audio patches.
• --kernel_size, --channel_width, --dilation_growth: Model architecture hyperparameters.
• --lr: Learning rate.
• --precision: Use 16 for half precision training.

python examples/train_comp.py \
--root_dir /path/to/SignalTrain_LA2A_Dataset_1.1 \
--max_epochs 20 \
--batch_size 128 \
--sample_rate 44100 \
--train_length 32768 \
--eval_length 262144 \
--num_workers 8 \
--kernel_size 15 \
--channel_width 32 \
--dilation_growth 2 \
--lr 0.001 \
--gpus 1 \
--shuffle True \
--precision 16 \
--preload True

You can enhance the MultiResolutionSTFTLoss by applying Mel-scaled spectrograms and perceptual weighting. This requires specifying fft_sizes, hop_sizes, win_lengths, scale="mel", n_bins, and the sample_rate.

import torch
import auraloss

bs = 8
chs = 1
seq_len = 131072
sample_rate = 44100

target = torch.rand(bs, chs, seq_len)
pred = torch.rand(bs, chs, seq_len)

# Define the loss function with perceptual weighting
loss_fn = auraloss.freq.MultiResolutionSTFTLoss(
    fft_sizes=[1024, 2048, 8192],
    hop_sizes=[256, 512, 2048],
    win_lengths=[1024, 2048, 8192],
    scale="mel",
    n_bins=128,
    sample_rate=sample_rate,
    perceptual_weighting=True,
)

# Compute loss
loss = loss_fn(pred, target)

bs = 8
chs = 1
seq_len = 131072
sample_rate = 44100

# some audio you want to compare
target = torch.rand(bs, chs, seq_len)
pred = torch.rand(bs, chs, seq_len)

# define the loss function
loss_fn = auraloss.freq.MultiResolutionSTFTLoss(
    fft_sizes=[1024, 2048, 8192],
    hop_sizes=[256, 512, 2048],
    win_lengths=[1024, 2048, 8192],
    scale="mel",
    n_bins=128,
    sample_rate=sample_rate,
    perceptual_weighting=True,
)

# compute
loss = loss_fn(pred, target)

auraloss provides loss functions categorized into Time domain, Frequency domain, and Perceptual transforms.

Time domain

• Error-to-signal ratio (ESR): auraloss.time.ESRLoss()
• DC error (DC): auraloss.time.DCLoss()
• Log hyperbolic cosine (Log-cosh): auraloss.time.LogCoshLoss()
• Signal-to-noise ratio (SNR): auraloss.time.SNRLoss()
• Scale-invariant signal-to-distortion ratio (SI-SDR): auraloss.time.SISDRLoss()
• Scale-dependent signal-to-distortion ratio (SD-SDR): auraloss.time.SDSDRLoss()

Frequency domain

• Aggregate STFT: auraloss.freq.STFTLoss()
• Aggregate Mel-scaled STFT: auraloss.freq.MelSTFTLoss(sample_rate)
• Multi-resolution STFT: auraloss.freq.MultiResolutionSTFTLoss()
• Random-resolution STFT: auraloss.freq.RandomResolutionSTFTLoss()
• Sum and difference STFT loss: auraloss.freq.SumAndDifferenceSTFTLoss()

Perceptual transforms

• Sum and difference signal transform: auraloss.perceptual.SumAndDifference()
• FIR pre-emphasis filters: auraloss.perceptual.FIRFilter()

The FIRFilter module implements perceptual pre-emphasis filters based on Wright & Välimäki, 2019. It applies the same filter to both an input (predicted) signal and a target (ground truth) signal.

Filter Types (filter_type):

• "hp": First-order highpass filter. Uses the coef parameter.
• "fd": Folded differentiator. Uses the coef parameter.
• "aw": A-weighting filter (based on IEC/CD 1672). Uses least-squares fitting to create a digital FIR filter from an analog definition.

Arguments:

• filter_type (str): Type of filter ("hp", "fd", or "aw"). Default: "hp".
• coef (float): Coefficient for "hp" and "fd". Default: 0.85 (optimized for 44.1 kHz).
• fs (int): Sampling rate. Default: 44100.
• ntaps (int): Number of FIR filter taps. Must be an odd integer. Default: 101.
• plot (bool): If True, plots the magnitude response of the filter. Default: False.

Returns:

• A tuple of (filtered_input, filtered_target) tensors, both with shape (B, #channels, #samples).

import torch
from auraloss.perceptual import FIRFilter

# Setup parameters
B, C, S = 1, 2, 16000
fs = 16000

# Initialize an A-weighting filter
filter_module = FIRFilter(filter_type="aw", fs=fs, ntaps=101)

# Dummy signals
input_signal = torch.randn(B, C, S)
target_signal = torch.randn(B, C, S)

# Apply filter
filtered_input, filtered_target = filter_module(input_signal, target_signal)