Geometric Loss functions between sampled measures, images and volumes

The GeomLoss library provides efficient GPU implementations for:

It is hosted on GitHub and distributed under the permissive MIT license.
PyPi version Downloads

GeomLoss functions are available through the custom PyTorch layers SamplesLoss, ImagesLoss and VolumesLoss which allow you to work with weighted point clouds (of any dimension), density maps and volumetric segmentation masks. Geometric losses come with three backends each:

  • A simple tensorized implementation, for small problems (< 5,000 samples).

  • A reference online implementation, with a linear (instead of quadratic) memory footprint, that can be used for finely sampled measures.

  • A very fast multiscale code, which uses an octree-like structure for large-scale problems in dimension <= 3.

A typical sample of code looks like:

import torch
from geomloss import SamplesLoss  # See also ImagesLoss, VolumesLoss

# Create some large point clouds in 3D
x = torch.randn(100000, 3, requires_grad=True).cuda()
y = torch.randn(200000, 3).cuda()

# Define a Sinkhorn (~Wasserstein) loss between sampled measures
loss = SamplesLoss(loss="sinkhorn", p=2, blur=.05)

L = loss(x, y)  # By default, use constant weights = 1/number of samples
g_x, = torch.autograd.grad(L, [x])  # GeomLoss fully supports autograd!

GeomLoss is a simple interface for cutting-edge Optimal Transport algorithms. It provides:

  • Support for batchwise computations.

  • Linear (instead of quadratic) memory footprint for large problems, relying on the KeOps library for map-reduce operations on the GPU.

  • Fast kernel truncation for small bandwidths, using an octree-based structure.

  • Log-domain stabilization of the Sinkhorn iterations, eliminating numeric overflows for small values of \(\varepsilon\).

  • Efficient computation of the gradients, which bypasses the naive backpropagation algorithm.

  • Support for unbalanced Optimal Transport, with a softening of the marginal constraints through a maximum reach parameter.

  • Support for the ε-scaling heuristic in the Sinkhorn loop, with kernel truncation in dimensions 1, 2 and 3. On typical 3D problems, our implementation is 50-100 times faster than the standard SoftAssign/Sinkhorn algorithm.

Note, however, that SamplesLoss does not implement the Fast Multipole or Fast Gauss transforms. If you are aware of a well-packaged implementation of these algorithms on the GPU, please contact me!

The divergences implemented here are all symmetric, positive definite and therefore suitable for measure-fitting applications. For positive input measures \(\alpha\) and \(\beta\), our \(\text{Loss}\) functions are such that

\[\begin{split}\text{Loss}(\alpha,\beta) ~&=~ \text{Loss}(\beta,\alpha), \\ 0~=~\text{Loss}(\alpha,\alpha) ~&\leqslant~ \text{Loss}(\alpha,\beta), \\ 0~=~\text{Loss}(\alpha,\beta)~&\Longleftrightarrow~ \alpha = \beta.\end{split}\]

GeomLoss can be used in a wide variety of settings, from shape analysis (LDDMM, optimal transport…) to machine learning (kernel methods, GANs…) and image processing. Details and examples are provided below:

Author and Contributors

Feel free to contact us for any bug report or feature request:

Licensing, academic use

This library is licensed under the permissive MIT license, which is fully compatible with both academic and commercial applications. If you use this code in a research paper, please cite:

@inproceedings{feydy2019interpolating,
    title={Interpolating between Optimal Transport and MMD using Sinkhorn Divergences},
    author={Feydy, Jean and S{\'e}journ{\'e}, Thibault and Vialard, Fran{\c{c}}ois-Xavier and Amari, Shun-ichi and Trouve, Alain and Peyr{\'e}, Gabriel},
    booktitle={The 22nd International Conference on Artificial Intelligence and Statistics},
    pages={2681--2690},
    year={2019}
}

Table of contents