Spectral Optimal Transport Losses for PyTorch

This repository contains an implementation of Spectral Optimal Transport (SOT) loss functions for PyTorch, with a pip-installable package sot-loss. SOT loss functions are differentiable spectral losses which compare the spectra of two audio signals using optimal transport principles. These loss functions can be used for training neural networks in audio processing tasks, particularly those involving DDSP. It can also be used more generally as a metric for audio signal comparison.

SOT does this

Multi-Scale Spectral loss and others do this

Installation

You can install the sot-loss package using pip:

pip install sot-loss

Usage

The primary components of this package are the Wasserstein1DLoss and MultiResolutionSOTLoss classes, which can be used as PyTorch loss functions. Here is a basic example of how to use the Wasserstein1DLoss:

import torch
from sot import Wasserstein1DLoss

# Create some dummy audio signals
x = torch.randn(4, 16000)
y = torch.randn(4, 16000)

# Initialize the SOT loss
sot_loss = Wasserstein1DLoss(transform='stft', 
                             fft_size=2048,
                             hop_length=512, 
                             sample_rate=16000, 
                             window='flattop', 
                             square_magnitude=True)

# Compute the loss
loss = sot_loss(x, y)
print(loss)

Using your own mapping audio -> 2D representation:

x_spec = custom_transform(x) # batch, channels, time
y_spec = custom_transform(y) 
x_positions = get_custom_positions(x_spec) # channels
y_positions = get_custom_positions(y_spec)

sot_loss = Wasserstein1DLoss(transform='identity',
                             # other non-transform parameters can go here
                             balanced=True,
                             normalize=True,
                             )
loss = sot_loss(x_spec, y_spec, x_positions=x_positions, y_positions=y_positions)
print(loss)

Advanced Usage

The Wasserstein1DLoss and MultiResolutionSOTLoss classes offer a range of parameters to customize the spectral representation and the loss calculation.

Spectral Transform Parameters

These parameters are available in both Wasserstein1DLoss and MultiResolutionSOTLoss.

Transform parameters (if using built-in transforms):

Loss parameters (applies even if using custom transforms):

The MultiResolutionSOTLoss combines multiple Wasserstein1DLoss instances, each with a different set of STFT parameters.

About the Paper

This is the also the official repository for the paper "Unsupervised Harmonic Parameter Estimation Using Differentiable DSP and Spectral Optimal Transport.", by Bernardo Torres, Geoffroy Peeters, and Gaël Richard. Check out the poster here.

For repoducing the results from the paper, please check out the paper branch.

Citation

If you find our work useful or use it in your research, you can cite it using:

@inproceedings{torres2024unsupervised,
  title={Unsupervised harmonic parameter estimation using differentiable DSP and spectral optimal transport},
  author={Torres, Bernardo and Peeters, Geoffroy and Richard, Ga{\"e}l},
  booktitle={ICASSP 2024-2024 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
  pages={1176--1180},
  year={2024},
  organization={IEEE}
}