You can visit KeOps official website for documentation and tutorials.
The KeOps library lets you compute generic reductions of very large arrays whose entries are given by a mathematical formula. It combines a tiled reduction scheme with an automatic differentiation engine, and can be used through Matlab, Python (NumPy or PyTorch) or R backends. It is perfectly suited to the computation of Kernel matrix-vector products and the associated gradients, even when the full kernel matrix does not fit into the GPU memory.
Math libraries understand variables as matrices, also known as tensors. (a) These are usually dense and encoded as explicit numerical arrays that can have a large memory footprint. (b) Alternatively, some operators can be encoded as sparse matrices: libraries store in memory the indices (in,jn) and values Mn = Min,jn that correspond to a small number of non-zero coefficients. Reduction operations are then implemented using indexing methods and scattered memory accesses. (c) We provide support for a third class of tensors: symbolic matrices whose coefficients are given by a formula Mi,j = F(xi,yj) that is evaluated on data arrays (xi) and (yj). Reduction operations are implemented using parallel schemes that compute the coefficients Mi,j on-the-fly. We take advantage of the structure of CUDA registers to bypass costly memory transfers and achieve optimal runtimes on a wide range of applications.
Using the PyTorch backend, a typical sample of code looks like:
# Create two arrays with 3 columns and a (huge) number of lines, on the GPU
import torch
x = torch.randn(1000000, 3, requires_grad=True).cuda()
y = torch.randn(2000000, 3).cuda()
# Turn our Tensors into KeOps symbolic variables:
from pykeops.torch import LazyTensor
x_i = LazyTensor( x[:,None,:] ) # x_i.shape = (1e6, 1, 3)
y_j = LazyTensor( y[None,:,:] ) # y_j.shape = ( 1, 2e6,3)
# We can now perform large-scale computations, without memory overflows:
D_ij = ((x_i - y_j)**2).sum(dim=2) # Symbolic (1e6,2e6,1) matrix of squared distances
K_ij = (- D_ij).exp() # Symbolic (1e6,2e6,1) Gaussian kernel matrix
# And come back to vanilla PyTorch Tensors or NumPy arrays using
# reduction operations such as .sum(), .logsumexp() or .argmin().
# Here, the kernel density estimation a_i = sum_j exp(-|x_i-y_j|^2)
# is computed using a CUDA online map-reduce routine that has a linear
# memory footprint and outperforms standard PyTorch implementations
# by two orders of magnitude.
a_i = K_ij.sum(dim=1) # Genuine torch.cuda.FloatTensor, a_i.shape = (1e6, 1),
g_x = torch.autograd.grad((a_i ** 2).sum(), [x]) # KeOps supports autograd!KeOps allows you to leverage your GPU without compromising on usability. It provides:
- Linear (instead of quadratic) memory footprint for Kernel operations.
- Support for a wide range of mathematical formulas.
- Seamless computation of derivatives, up to arbitrary orders.
- Sum, LogSumExp, Min, Max but also ArgMin, ArgMax or K-min reductions.
- A conjugate gradient solver for e.g. large-scale spline interpolation or kriging, Gaussian process regression.
- An interface for block-sparse and coarse-to-fine strategies.
- Support for multi GPU configurations.
KeOps can thus be used in a wide variety of settings, from shape analysis (LDDMM, optimal transport...) to machine learning (kernel methods, k-means...) or kriging (aka. Gaussian process regression). More details are provided below:
KeOps is one of the fastest solution to compute generic Kernel formula. See our benchmark
When using KeOps for your scientific work, please cite the preprint or directly use the following bibtex:
@article{charlier2020kernel,
title={Kernel operations on the {GPU}, with autodiff, without memory overflows},
author={Charlier, Benjamin and Feydy, Jean and Glaun{\`e}s, Joan Alexis and Collin, Fran{\c{c}}ois-David and Durif, Ghislain},
journal={arXiv preprint arXiv:2004.11127},
year={2020}
}As of today, KeOps provides core routines for:
- Deformetrica, a shape analysis software developed by the Aramis Inria team.
- GeomLoss, a multiscale implementation of Kernel and Wasserstein distances that scales up to millions of samples on modern hardware.
- FshapesTk and the Shapes toolbox, two research-oriented LDDMM toolkits.
Feel free to contact us for any bug report or feature request, you can also fill an issue report on GitHub.
- François-David Collin