Note
Go to the end to download the full example code
Benchmarking Gaussian convolutions in high dimensions
Let’s compare the performances of PyTorch and KeOps on simple Gaussian RBF kernel products, as the dimension grows.
Setup
import torch
from matplotlib import pyplot as plt
from benchmark_utils import random_normal, full_benchmark
use_cuda = torch.cuda.is_available()
Benchmark specifications:
N = 10000 # Number of samples
# Dimensions to test:
Dims = [1, 3, 5, 10, 20, 30, 50, 80, 100, 120, 150, 200, 300, 500, 1000, 2000, 3000]
Synthetic dataset.
def generate_samples(D, device="cuda", lang="torch", batchsize=1, **kwargs):
"""Generates two point clouds x, y and a scalar signal b of size N.
Args:
D (int): dimension of the ambient space.
device (str, optional): "cuda", "cpu", etc. Defaults to "cuda".
lang (str, optional): "torch", "numpy", etc. Defaults to "torch".
batchsize (int, optional): number of experiments to run in parallel. Defaults to None.
Returns:
3-uple of arrays: x, y, b
"""
randn = random_normal(device=device, lang=lang)
x = randn((batchsize, N, D))
y = randn((batchsize, N, D))
b = randn((batchsize, N, 1))
return x, y, b
Define a simple Gaussian RBF product, using a tensorized implementation. Note that expanding the squared norm \(\|x-y\|^2\) as a sum \(\|x\|^2 - 2 \langle x, y \rangle + \|y\|^2\) allows us to leverage the fast matrix-matrix product of the BLAS/cuBLAS libraries.
def gaussianconv_pytorch(x, y, b, tf32=False, **kwargs):
"""(B,N,D), (B,N,D), (B,N,1) -> (B,N,1)"""
# If False, we stick to float32 computations.
# If True, we use TensorFloat32 whenever possible.
torch.backends.cuda.matmul.allow_tf32 = tf32
D_xx = (x * x).sum(-1).unsqueeze(2) # (B,N,1)
D_xy = torch.matmul(x, y.permute(0, 2, 1)) # (B,N,D) @ (B,D,M) = (B,N,M)
D_yy = (y * y).sum(-1).unsqueeze(1) # (B,1,M)
D_xy = D_xx - 2 * D_xy + D_yy # (B,N,M)
K_xy = (-D_xy).exp() # (B,N,M)
return K_xy @ b # (B,N,1)
Define a simple Gaussian RBF product, using an online implementation:
from pykeops.torch import generic_sum
def gaussianconv_keops(x, y, b, backend="GPU", **kwargs):
D = x.shape[-1]
fun = generic_sum(
"Exp(X|Y) * B", # Formula
"A = Vi(1)", # Output
"X = Vi({})".format(D), # 1st argument
"Y = Vj({})".format(D), # 2nd argument
"B = Vj(1)", # 3rd argument
)
ex = (-(x * x).sum(-1)).exp()[:, :, None]
ey = (-(y * y).sum(-1)).exp()[:, :, None]
return ex * fun(2 * x, y, b * ey, backend=backend)
Same, but without the chunked computation mode:
def gaussianconv_keops_nochunks(x, y, b, backend="GPU", **kwargs):
D = x.shape[-1]
fun = generic_sum(
"Exp(X|Y) * B", # Formula
"A = Vi(1)", # Output
"X = Vi({})".format(D), # 1st argument
"Y = Vj({})".format(D), # 2nd argument
"B = Vj(1)", # 3rd argument
enable_chunks=False,
)
ex = (-(x * x).sum(-1)).exp()[:, :, None]
ey = (-(y * y).sum(-1)).exp()[:, :, None]
return ex * fun(2 * x, y, b * ey, backend=backend)
PyTorch vs. KeOps (Gpu)
routines = [
(gaussianconv_pytorch, "PyTorch (GPU, TF32=False)", {"tf32": False}),
(gaussianconv_pytorch, "PyTorch (GPU, TF32=True)", {"tf32": True}),
(gaussianconv_keops_nochunks, "KeOps < 1.4.2 (GPU)", {}),
(gaussianconv_keops, "KeOps >= 1.4.2 (GPU)", {}),
]
full_benchmark(
f"Gaussian Matrix-Vector products in high dimension, with N={N:,} (GPU)",
routines,
generate_samples,
problem_sizes=Dims,
xlabel="Dimension of the points",
)
plt.show()
Benchmarking : Gaussian Matrix-Vector products in high dimension, with N=10,000 (GPU) ===============================
PyTorch (GPU, TF32=False) -------------
1x100 loops of size 1 : 1x100x 2.9 ms
1x100 loops of size 3 : 1x100x 2.9 ms
1x100 loops of size 5 : 1x100x 2.9 ms
1x100 loops of size 10 : 1x100x 2.9 ms
1x100 loops of size 20 : 1x100x 3.0 ms
1x100 loops of size 30 : 1x100x 3.1 ms
1x100 loops of size 50 : 1x100x 3.3 ms
1x100 loops of size 80 : 1x100x 3.6 ms
1x100 loops of size 100 : 1x100x 3.8 ms
1x100 loops of size 120 : 1x100x 4.0 ms
1x100 loops of size 150 : 1x100x 4.3 ms
1x100 loops of size 200 : 1x100x 4.8 ms
1x100 loops of size 300 : 1x100x 6.0 ms
1x100 loops of size 500 : 1x100x 8.1 ms
1x100 loops of size 1 k: 1x100x 13.5 ms
1x100 loops of size 2 k: 1x100x 24.4 ms
1x 10 loops of size 3 k: 1x 10x 35.2 ms
PyTorch (GPU, TF32=True) -------------
1x100 loops of size 1 : 1x100x 2.9 ms
1x100 loops of size 3 : 1x100x 2.9 ms
1x100 loops of size 5 : 1x100x 2.9 ms
1x100 loops of size 10 : 1x100x 2.9 ms
1x100 loops of size 20 : 1x100x 3.0 ms
1x100 loops of size 30 : 1x100x 3.1 ms
1x100 loops of size 50 : 1x100x 3.1 ms
1x100 loops of size 80 : 1x100x 3.0 ms
1x100 loops of size 100 : 1x100x 3.1 ms
1x100 loops of size 120 : 1x100x 3.1 ms
1x100 loops of size 150 : 1x100x 3.4 ms
1x100 loops of size 200 : 1x100x 3.2 ms
1x100 loops of size 300 : 1x100x 3.5 ms
1x100 loops of size 500 : 1x100x 3.7 ms
1x100 loops of size 1 k: 1x100x 4.5 ms
1x100 loops of size 2 k: 1x100x 6.2 ms
1x100 loops of size 3 k: 1x100x 7.9 ms
KeOps < 1.4.2 (GPU) -------------
1x100 loops of size 1 : 1x100x 489.7 µs
1x100 loops of size 3 : 1x100x 547.2 µs
1x100 loops of size 5 : 1x100x 589.2 µs
1x100 loops of size 10 : 1x100x 848.1 µs
1x100 loops of size 20 : 1x100x 1.4 ms
1x100 loops of size 30 : 1x100x 1.6 ms
1x100 loops of size 50 : 1x100x 3.2 ms
1x100 loops of size 80 : 1x100x 4.9 ms
1x100 loops of size 100 : 1x100x 7.6 ms
1x100 loops of size 120 : 1x100x 4.9 ms
1x100 loops of size 150 : 1x100x 8.3 ms
1x100 loops of size 200 : 1x100x 10.0 ms
1x100 loops of size 300 : 1x100x 62.4 ms
1x 10 loops of size 500 : 1x 10x 512.2 ms
1x 1 loops of size 1 k: 1x 1x 2.8 s
1x 1 loops of size 2 k: 1x 1x 12.2 s
** Too slow!
KeOps >= 1.4.2 (GPU) -------------
1x100 loops of size 1 : 1x100x 482.7 µs
1x100 loops of size 3 : 1x100x 542.2 µs
1x100 loops of size 5 : 1x100x 588.0 µs
1x100 loops of size 10 : 1x100x 847.4 µs
1x100 loops of size 20 : 1x100x 1.4 ms
1x100 loops of size 30 : 1x100x 1.6 ms
1x100 loops of size 50 : 1x100x 3.2 ms
1x100 loops of size 80 : 1x100x 4.9 ms
1x100 loops of size 100 : 1x100x 7.4 ms
1x100 loops of size 120 : 1x100x 8.3 ms
1x100 loops of size 150 : 1x100x 12.2 ms
1x100 loops of size 200 : 1x100x 15.5 ms
1x100 loops of size 300 : 1x100x 20.4 ms
1x 10 loops of size 500 : 1x 10x 35.5 ms
1x 10 loops of size 1 k: 1x 10x 81.4 ms
1x 10 loops of size 2 k: 1x 10x 127.6 ms
1x 10 loops of size 3 k: 1x 10x 222.5 ms
Total running time of the script: ( 1 minutes 16.627 seconds)