def fit(x_train):
# Setup the K-NN estimator:
start = time.time()
# Encoding as KeOps LazyTensors:
D = x_train.shape[1]
X_i = Vi(0, D) # Purely symbolic "i" variable, without any data array
X_j = Vj(1, D) # Purely symbolic "j" variable, without any data array
# Symbolic distance matrix:
if metric == "euclidean":
D_ij = ((X_i - X_j)**2).sum(-1)
# K-NN query operator:
KNN_fun = D_ij.argKmin(K, dim=1)
# N.B.: The "training" time here should be negligible.
elapsed = time.time() - start
def f(x_test):
start = time.time()
# Actual K-NN query:
indices = KNN_fun(x_test, x_train)
elapsed = time.time() - start
indices = indices.cpu().numpy()
return indices, elapsed
return f, elapsed
def plot_kernel(gamma):
""" Samples 'x -> ∑_j b_j * k_j(x - y_j)' on the grid, and displays it as a heatmap. """
if gamma.dim() == 2:
heatmap = (-Vi(x).weightedsqdist(Vj(y), Vj(gamma))).exp() @ b
else:
heatmap = (-Vi(x).weightedsqdist(Vj(y), Pm(gamma))).exp() @ b
heatmap = heatmap.view(
res, res).cpu().numpy() # reshape as a 'background' image
plt.imshow(
-heatmap,
interpolation="bilinear",
origin="lower",
vmin=-1,
vmax=1,
cmap=cm.RdBu,
extent=(0, 1, 0, 1),
)
plt.show()
def lse_lazytensor(p, D, batchdims=(1, )):
"""Modern LazyTensor implementation."""
x_i = Vi(0, D)
y_j = Vj(1, D)
f_j = Vj(2, 1)
epsinv = Pm(3, 1)
x_i.batchdims = batchdims
y_j.batchdims = batchdims
f_j.batchdims = batchdims
epsinv.batchdims = batchdims
if p == 2:
D_ij = ((x_i - y_j)**2).sum(-1) / 2
elif p == 1:
D_ij = ((x_i - y_j)**2).sum(-1).sqrt()
smin = (f_j - epsinv * D_ij).logsumexp(2)
return smin
def test_logSumExp_kernels_feature(self):
############################################################
from pykeops.torch import Vi, Vj, Pm
kernels = {
"gaussian": lambda xc, yc, sigmac: (
-Pm(1 / sigmac ** 2) * Vi(xc).sqdist(Vj(yc))
),
"laplacian": lambda xc, yc, sigmac: (
-(Pm(1 / sigmac ** 2) * Vi(xc).sqdist(Vj(yc))).sqrt()
),
"cauchy": lambda xc, yc, sigmac: (
1 + Pm(1 / sigmac ** 2) * Vi(xc).sqdist(Vj(yc))
)
.power(-1)
.log(),
"inverse_multiquadric": lambda xc, yc, sigmac: (
1 + Pm(1 / sigmac ** 2) * Vi(xc).sqdist(Vj(yc))
)
.sqrt()
.power(-1)
.log(),
}
for k in ["gaussian", "laplacian", "cauchy", "inverse_multiquadric"]:
with self.subTest(k=k):
# Call cuda kernel
gamma_lazy = kernels[k](self.xc, self.yc, self.sigmac)
gamma_lazy = gamma_lazy.logsumexp(dim=1, weight=Vj(self.gc.exp())).cpu()
# gamma = kernel_product(params, self.xc, self.yc, self.gc).cpu()
# Numpy version
log_K = log_np_kernel(self.x, self.y, self.sigma, kernel=k)
log_KP = log_K + self.g.T
gamma_py = log_sum_exp(log_KP, axis=1)
# compare output
self.assertTrue(
np.allclose(gamma_lazy.data.numpy().ravel(), gamma_py, atol=1e-6)
)
def fit(x_train):
# Setup the K-NN estimator:
x_train = tensor(x_train)
start = timer()
# Encoding as KeOps LazyTensors:
D = x_train.shape[1]
X_i = Vi(0, D) # Purely symbolic "i" variable, without any data array
X_j = Vj(1, D) # Purely symbolic "j" variable, without any data array
# Symbolic distance matrix:
if metric == "euclidean":
D_ij = ((X_i - X_j)**2).sum(-1)
elif metric == "manhattan":
D_ij = (X_i - X_j).abs().sum(-1)
elif metric == "angular":
D_ij = -(X_i | X_j)
elif metric == "hyperbolic":
D_ij = ((X_i - X_j)**2).sum(-1) / (X_i[0] * X_j[0])
else:
raise NotImplementedError(
f"The '{metric}' distance is not supported.")
# K-NN query operator:
KNN_fun = D_ij.argKmin(K, dim=1)
# N.B.: The "training" time here should be negligible.
elapsed = timer() - start
def f(x_test):
x_test = tensor(x_test)
start = timer()
# Actual K-NN query:
indices = KNN_fun(x_test, x_train)
elapsed = timer() - start
indices = indices.cpu().numpy()
return indices, elapsed
return f, elapsed
def log_likelihoods(self, sample):
"""Log-density, sampled on a given point cloud."""
self.update_covariances()
K_ij = -Vi(sample).weightedsqdist(Vj(self.mu), Vj(
self.params["gamma"]))
return K_ij.logsumexp(dim=1, weight=Vj(self.weights()))
def likelihoods(self, sample):
"""Samples the density on a given point cloud."""
self.update_covariances()
return (-Vi(sample).weightedsqdist(Vj(self.mu), Vj(
self.params["gamma"]))).exp() @ self.weights()
###############################################################################
# .. note::
# This operator uses a conjugate gradient solver and assumes
# that **formula** defines a **symmetric**, positive and definite
# **linear** reduction with respect to the alias ``"b"``
# specified trough the third argument.
#
# Apply our solver on arbitrary point clouds:
#
print(
"Solving a Gaussian linear system, with {} points in dimension {}.".format(
N, D))
start = time.time()
K_xx = keops.exp(-keops.sum((Vi(x) - Vj(x))**2, dim=2) / (2 * sigma**2))
cfun = keops.solve(K_xx, Vi(b), alpha=alpha, call=False)
c = cfun()
end = time.time()
print("Timing (KeOps implementation):", round(end - start, 5), "s")
###############################################################################
# Compare with a straightforward PyTorch implementation:
#
start = time.time()
K_xx = alpha * torch.eye(N) + torch.exp(-torch.sum(
(x[:, None, :] - x[None, :, :])**2, dim=2) / (2 * sigma**2))
c_py = torch.solve(b, K_xx)[0]
end = time.time()
print("Timing (PyTorch implementation):", round(end - start, 5), "s")
sigmac = torch.tensor(sigma, dtype=torchtype, device=device)
except:
pass
####################################################################
# Convolution Benchmarks
# ----------------------
#
# We loop over four different kernels:
#
kernel_to_test = ['gaussian', 'laplacian', 'cauchy', 'inverse_multiquadric']
kernels = {
'gaussian':
lambda xc, yc, sigmac: (-Pm(1 / sigmac**2) * Vi(xc).sqdist(Vj(yc))).exp(),
'laplacian':
lambda xc, yc, sigmac:
(-(Pm(1 / sigmac**2) * Vi(xc).sqdist(Vj(yc))).sqrt()).exp(),
'cauchy':
lambda xc, yc, sigmac:
(1 + Pm(1 / sigmac**2) * Vi(xc).sqdist(Vj(yc))).power(-1),
'inverse_multiquadric':
lambda xc, yc, sigmac:
(1 + Pm(1 / sigmac**2) * Vi(xc).sqdist(Vj(yc))).sqrt().power(-1)
}
#####################################################################
# With four backends: Numpy, vanilla PyTorch, Generic KeOps reductions
# and a specific, handmade legacy CUDA code for kernel convolutions:
#
# Here is how it works, if we # want to perform a simple gaussian convolution. # # We first create the dataset using 2D tensors: M, N = (100000, 200000) if use_cuda else (1000, 2000) D = 3 x = torch.randn(M, D).type(tensor) y = torch.randn(N, D).type(tensor) ############################################################################ # Then we use :func:`Vi <pykeops.torch.Vi>` and :func:`Vj <pykeops.torch.Vj>` to convert to KeOps LazyTensor objects from pykeops.torch import Vi, Vj x_i = Vi(x) # (M, 1, D) LazyTensor, equivalent to LazyTensor( x[:,None,:] ) y_j = Vj(y) # (1, N, D) LazyTensor, equivalent to LazyTensor( y[None,:,:] ) ############################################################################ # and perform our operations: D2xy = ((x_i - y_j)**2).sum() gamma = D2xy.sum_reduction(dim=1) ######################################################################### # Note that in the first line we used ``sum`` without any axis or dim parameter. # This is equivalent to ``sum(-1)`` or ``sum(dim=2)``, because # the axis parameter is set to ``2`` by default. But speaking about ``dim=2`` # here with the :func:`Vi <pykeops.torch.Vi>`, :func:`Vj <pykeops.torch.Vj>` helpers could be misleading. # Similarly we used ``sum_reduction`` instead of ``sum`` to make it clear # that we perform a reduction, but sum and sum_reduction with ``dim=0`` or ``1`` # are equivalent (however ``sum_reduction`` with ``dim=2`` is forbidden)
def GaussLinKernel(sigma):
x, y, u, v, b = Vi(0, 3), Vj(1, 3), Vi(2, 3), Vj(3, 3), Vj(4, 1)
gamma = 1 / (sigma * sigma)
D2 = x.sqdist(y)
K = (-D2 * gamma).exp() * (u * v).sum()**2
return (K * b).sum_reduction(axis=1)
def GaussKernel(sigma):
x, y, b = Vi(0, 3), Vj(1, 3), Vj(2, 3)
gamma = 1 / (sigma * sigma)
D2 = x.sqdist(y)
K = (-D2 * gamma).exp()
return (K * b).sum_reduction(axis=1)
def gaussian_kernel(x, y, sigma=0.2):
x_i = Vi(x)
y_j = Vj(y)
D_ij = ((x_i - y_j)**2).sum(-1)
return (-D_ij / (2 * sigma**2)).exp()
###############################################################
# Generate some data, stored on the CPU (host) memory:
#
M = 1000
N = 2000
x = np.random.randn(M, 3).astype(dtype)
y = np.random.randn(N, 3).astype(dtype)
a = np.random.randn(N, 1).astype(dtype)
p = np.random.randn(1).astype(dtype)
#########################################
# Launch our routine on the CPU:
#
xi, yj, aj = Vi(x), Vj(y), Vj(a)
c = ((p - aj)**2 * (xi + yj).exp()).sum(axis=1, backend='CPU')
#########################################
# And on our GPUs, with copies between
# the Host and Device memories:
#
if IsGpuAvailable():
for gpuid in gpuids:
d = ((p - aj)**2 * (xi + yj).exp()).sum(axis=1,
backend='GPU',
device_id=gpuid)
print('Relative error on gpu {}: {:1.3e}'.format(
gpuid, float(np.sum(np.abs(c - d)) / np.sum(np.abs(c)))))
# Plot the results next to each other: