# Geometric Loss functions between sampled measures, images and volumes

The GeomLoss library provides efficient GPU implementations for:  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

# Create some large point clouds in 3D
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


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:

@inproceedings{feydy2019interpolating,