Geometric Loss functions between sampled measures, images and volumes

N.B.: This is still an alpha release! Please send me your feedback: I will polish the user interface, implement Hausdorff divergences, add support for images/volumes and clean the documentation before a first stable release by mid-May 2019 - coming back from AiStats 2019.

The GeomLoss library provides efficient GPU implementations for:

These loss functions, defined between positive measures, 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:

GeomLoss is licensed under the MIT license.


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