GeomLoss : Geometric Loss functions between sampled measures, images and volumes

[Geo-Sci-Info]

GeomLoss : Geometric Loss functions between sampled measures, images and volumes

Find all the docs and tutorials of the version 0.2.3 in the read the docs website:

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 meshes, images, volumes and clean the documentation over the summer of 2020.

The GeomLoss library provides efficient GPU implementations for:

• Kernel norms (also known as Maximum Mean Discrepancies).

• Hausdorff divergences, which are positive definite generalizations of the ICP loss, analogous to log-likelihoods of Gaussian Mixture Models.

• Unbiased Sinkhorn divergences, which are cheap yet positive definite approximations of Optimal Transport (Wasserstein) costs.

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.

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 𝜀
• 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 𝛼 and 𝛽, our Loss

functions are such that
Loss(𝛼,𝛽) = Loss(𝛽,𝛼),
0 = Loss(𝛼,𝛼) ⩽ Loss(𝛼,𝛽),
0 = Loss(𝛼,𝛽) ⟺ 𝛼=𝛽.

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:

• Maths and algorithms
• PyTorch API
• Source code
• Examples

GeomLoss is licensed under the MIT license.

Author and Contributors

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

• Jean Feydy
• Pierre Roussillon (extensions to brain tractograms and normal cycles)

Related projects

You may be interested by:

• The KeOps library, which provides efficient CUDA routines for point cloud processing, with full PyTorch support.

• Rémi Flamary and Nicolas Courty’s Python Optimal Transport library, which provides a reference implementation of OT-related methods for small problems.

• Bernhard Schmitzer’s Optimal Transport toolbox, which provides a reference multiscale solver for the OT problem, on the CPU.