DOCSEARCHbyAlgolia
Submit Docs

GeomLoss

repository·main·Indexed 20 days ago

https://github.com/jeanfeydy/geomloss

A library providing geometric loss functions for computing distances between point clouds, images, and volumes. It implements kernel norms (MMD), Hausdorff divergences, and debiased Sinkhorn divergences. GeomLoss integrates with PyTorch via SamplesLoss, ImagesLoss, and VolumesLoss, utilizing the KeOps library for a linear memory footprint on GPUs and supporting multiscale implementations for large-scale problems.

Tokens
8K
Snippets
21
Records
34
Agent score
71%

What's inside geomloss

  1. Overview of GeomLoss capabilities

    main

    GeomLoss provides efficient GPU implementations for geometric loss functions between sampled measures, images, and volumes. It supports:

    • Kernel norms (Maximum Mean Discrepancies).
    • Hausdorff divergences (generalizations of Chamfer-ICP loss).
    • Debiased Sinkhorn divergences (approximations of Optimal Transport/Wasserstein distances).

    All implemented divergences are symmetric and positive definite, making them suitable for measure-fitting applications.

  3. Solve optimal transport with geomloss.ot

    main

    The geomloss.ot submodule provides scalable optimal transport solvers compatible with the Python Optimal Transport library. It allows for solving optimal transport problems using either pre-computed cost matrices or cost functions evaluated on points.

    Warning: The geomloss.ot API is currently in development and subject to change.

  4. How to use GeomLoss with PyTorch

    main

    GeomLoss integrates with PyTorch through three main custom layer classes:

    • SamplesLoss: For weighted point clouds of any dimension.
    • ImagesLoss: For density maps.
    • VolumesLoss: For volumetric segmentation masks.

    Each loss class supports three backends depending on the problem scale:

    1. tensorized: For small problems (< 5,000 samples).
    2. online: A reference implementation with a linear memory footprint, suitable for finely sampled measures.
    3. multiscale: A fast implementation using an octree-like structure for large-scale problems in dimensions $\le 3$.

    GeomLoss supports batchwise computations, autograd, and efficient gradient computation.

    import torch
from geomloss import SamplesLoss

# 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!

  5. Key features of GeomLoss

    main

    GeomLoss is designed for high-performance Optimal Transport tasks with the following features:

    • Linear Memory Footprint: Uses the KeOps library for map-reduce operations on the GPU to avoid quadratic memory scaling.
    • Fast Kernel Truncation: Uses an octree-based structure for small bandwidths in dimensions 1, 2, and 3.
    • Numerical Stability: Log-domain stabilization of Sinkhorn iterations prevents overflows for small $\varepsilon$.
    • Unbalanced Optimal Transport: Supports softening marginal constraints via a reach parameter.
    • $\varepsilon$-scaling heuristic: Implements scaling in the Sinkhorn loop for efficiency.
    • Efficient Gradients: Bypasses naive backpropagation for faster gradient computation.

  6. Use the SamplesLoss interface for geometric loss

    main
    The geomloss module provides geometric loss functions between point clouds, images, and volumes. While a new geomloss.ot API is under development, the SamplesLoss interface is the stable, long-term supported method for computing these losses. Use SamplesLoss for current production or research needs until the ot API is stabilized.

  7. Install GeomLoss via pip

    main

    To install GeomLoss and its optional dependencies, follow these steps:

    1. Install PyTorch.
    2. Install the KeOps library.
    3. Install GeomLoss using pip.

    For Google Colab users, you can install the full suite of dependencies in one command.

    # Standard installation
pip install geomloss

# Google Colab installation (includes optional dependencies)
!pip install geomloss[full]

  8. Compare Kernel, Hausdorff, and Sinkhorn loss functions

    main

    The geomloss library provides three distinct types of geometric loss functions for comparing point clouds, images, or volumes. You can use the examples in the examples/comparisons/ directory to understand the mathematical and practical differences between them:

    • Kernel loss: Uses a kernel-based approach (often related to optimal transport) to measure distances.
    • Hausdorff loss: Measures the maximum distance from a point in one set to the nearest point in the other set.
    • Sinkhorn loss: An entropic regularized version of optimal transport, typically used for its computational efficiency and smoothness.

  9. Install GeomLoss from Git

    main

    You can install a specific version of GeomLoss directly from the GitHub repository using pip syntax, or by cloning the repository manually and updating your PYTHONPATH.

    # Install the main branch with full dependencies using pip
pip install git+https://github.com/jeanfeydy/geomloss.git@main#egg=project[full]

# Manual installation via cloning
git clone https://github.com/jeanfeydy/geomloss.git
# Then add the path to your environment:
export PYTHONPATH=$PYTHONPATH:/path/to/geomloss

  10. Reproduce brain tractogram experiments using unbalanced regularized Optimal Transport

    main
    You can use geomloss to process white matter fiber tracks by applying unbalanced, regularized Optimal Transport. The provided scripts in the examples/brain_tractograms/ directory allow for the reproduction of experiments from the paper 'Fast and scalable Optimal Transport for brain tractograms' (Miccai 2019) by Jean Feydy, Pierre Roussillon, Alain Trouvé, and Pietro Gori.

  11. Build the documentation on Google Colab

    main

    You can generate the documentation website using a Google Colab notebook. This process involves mounting Google Drive to save the output, installing Sphinx and its dependencies, cloning the repositories, and running the make html command twice (once to compile KeOps routines and once to get accurate timings).

    # 1. Mount Google Drive
from google.colab import drive
drive.mount('/content/gdrive')

# 2. Install Sphinx and dependencies
!pip uninstall sphinx
!pip3 install sphinx
!pip install numpy GPUtil cmake ninja sphinx-gallery recommonmark sphinxcontrib-httpdomain sphinx_rtd_theme plyfile

# 3. Install KeOps and GeomLoss
!pip install pykeops
!git clone https://github.com/jeanfeydy/geomloss.git

# 4. Configure environment paths
import os
import sys
os.environ['PYTHONPATH'] += ":/content/geomloss/"
sys.path.append('/content/geomloss/')

# 5. Build documentation
%cd /content/geomloss/doc
!make html
!make clean
!make html

# 6. Save to Google Drive
!zip -r geomloss_documentation.zip _build
!cp geomloss_documentation.zip /content/gdrive/'My Drive'