Package: auraloss
auralossA collection of audio-focused loss functions in PyTorch.[PDF]Setuppip install auralossIf you want to useMelSTFTLoss()orFIRFilter()you will need to specify the extra install (librosa and scipy).pip install auraloss[all]Usageimporttorchimportauralossmrstft=auraloss.freq.MultiResolutionSTFTLoss()input=torch.rand(8,1,44100)target=torch.rand(8,1,44100)loss=mrstft(input,target)NEW: Perceptual weighting with mel scaled spectrograms.bs=8chs=1seq_len=131072sample_rate=44100# some audio you want to comparetarget=torch.rand(bs,chs,seq_len)pred=torch.rand(bs,chs,seq_len)# define the loss functionloss_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,)# computeloss=loss_fn(pred,target)CitationIf you use this code in your work please consider citing us.@inproceedings{steinmetz2020auraloss,title={auraloss: {A}udio focused loss functions in {PyTorch}},author={Steinmetz, Christian J. and Reiss, Joshua D.},booktitle={Digital Music Research Network One-day Workshop (DMRN+15)},year={2020}}Loss functionsWe categorize the loss functions as either time-domain or frequency-domain approaches.
Additionally, we include perceptual transforms.Loss functionInterfaceReferenceTime domainError-to-signal ratio (ESR)auraloss.time.ESRLoss()Wright & Välimäki, 2019DC error (DC)auraloss.time.DCLoss()Wright & Välimäki, 2019Log hyperbolic cosine (Log-cosh)auraloss.time.LogCoshLoss()Chen et al., 2019Signal-to-noise ratio (SNR)auraloss.time.SNRLoss()Scale-invariant signal-to-distortionratio (SI-SDR)auraloss.time.SISDRLoss()Le Roux et al., 2018Scale-dependent signal-to-distortionratio (SD-SDR)auraloss.time.SDSDRLoss()Le Roux et al., 2018Frequency domainAggregate STFTauraloss.freq.STFTLoss()Arik et al., 2018Aggregate Mel-scaled STFTauraloss.freq.MelSTFTLoss(sample_rate)Multi-resolution STFTauraloss.freq.MultiResolutionSTFTLoss()Yamamoto et al., 2019*Random-resolution STFTauraloss.freq.RandomResolutionSTFTLoss()Steinmetz & Reiss, 2020Sum and difference STFT lossauraloss.freq.SumAndDifferenceSTFTLoss()Steinmetz et al., 2020Perceptual transformsSum and difference signal transformauraloss.perceptual.SumAndDifference()FIR pre-emphasis filtersauraloss.perceptual.FIRFilter()Wright & Välimäki, 2019*Wang et al., 2019also propose a multi-resolution spectral loss (thatEngel et al., 2020follow),
but they do not include both the log magnitude (L1 distance) and spectral convergence terms, introduced inArik et al., 2018, and then extended for the multi-resolution case inYamamoto et al., 2019.ExamplesCurrently we include an example using a set of the loss functions to train a TCN for modeling an analog dynamic range compressor.
For details please refer to the details inexamples/compressor.
We provide pre-trained models, evaluation scripts to compute the metrics in thepaper, as well as scripts to retrain models.There are some more advanced things you can do based upon theSTFTLossclass.
For example, you can compute both linear and log scaled STFT errors as inEngel et al., 2020.
In this case we do not include the spectral convergence term.stft_loss=auraloss.freq.STFTLoss(w_log_mag=1.0,w_lin_mag=1.0,w_sc=0.0,)There is also a Mel-scaled STFT loss, which has some special requirements.
This loss requires you set the sample rate as well as specify the correct device.sample_rate=44100melstft_loss=auraloss.freq.MelSTFTLoss(sample_rate,device="cuda")You can also build a multi-resolution Mel-scaled STFT loss with 64 bins easily.
Make sure you pass the correct device where the tensors you are comparing will be.loss_fn=auraloss.freq.MultiResolutionSTFTLoss(scale="mel",n_bins=64,sample_rate=sample_rate,device="cuda")If you are computing a loss on stereo audio you may want to consider the sum and difference (mid/side) loss.
Below we have shown an example of using this loss function with the perceptual weighting and mel scaling for
further perceptual relevance.target=torch.rand(8,2,44100)pred=torch.rand(8,2,44100)loss_fn=auraloss.freq.SumAndDifferenceSTFTLoss(fft_sizes=[1024,2048,8192],hop_sizes=[256,512,2048],win_lengths=[1024,2048,8192],perceptual_weighting=True,sample_rate=44100,scale="mel",n_bins=128,)loss=loss_fn(pred,target)DevelopmentRun tests locally with pytest.python -m pytest