def K(x, y, u, v, b):
params = {
"id": Kernel("gaussian(x,y) * linear(u,v)**2"),
"gamma": (CpuOrGpu(torch.FloatTensor([1 / sigma**2])), None),
"backend": backend_keops
}
return kernel_product(params, (x, u), (y, v), b)
def K(x, y, b):
params = {
'id': Kernel('gaussian(x,y)'),
'gamma': 1 / (sigma * sigma),
'backend': 'auto'
}
return kernel_product(params, x, y, b)
def K(x, y, b):
params = {
"id": Kernel("gaussian(x,y)"),
"gamma": CpuOrGpu(torch.FloatTensor([1 / sigma**2])),
"backend": backend_keops
}
return kernel_product(params, x, y, b)
def K(x, y, u, v, b):
params = {
'id': Kernel('gaussian(x,y) * linear(u,v)**2'),
'gamma': (1 / (sigma * sigma), None),
'backend': 'auto'
}
return kernel_product(params, (x, u), (y, v), b)
def __init__(self, M, sparsity=0, D=2):
super(GaussianMixture, self).__init__()
self.params = {'id': Kernel('gaussian(x,y)')}
# We initialize our model with random blobs scattered across
# the unit square, with a small-ish radius:
self.mu = Parameter(torch.rand(M, D).type(dtype))
self.A = 15 * torch.ones(M, 1, 1) * torch.eye(D, D).view(1, D, D)
self.A = Parameter((self.A).type(dtype).contiguous())
self.w = Parameter(torch.ones(M, 1).type(dtype))
self.sparsity = sparsity
interpolation='bilinear',
origin='lower',
vmin=-1,
vmax=1,
cmap=cm.RdBu,
extent=(0, 1, 0, 1))
plt.title(title, fontsize=20)
plt.figure()
# TEST ===================================================================================
# Let's use a "gaussian" kernel, i.e.
# k(x_i,y_j) = exp( - WeightedSquareNorm(gamma, x_i-y_j ) )
params = {
"id": Kernel("gaussian(x,y)"),
}
# The type of kernel is inferred from the shape of the parameter "gamma",
# used as a "metric multiplier".
# Denoting D == x.shape[1] == y.shape[1] the size of the feature space, rules are :
# - if "gamma" is a vector (gamma.shape = [K]), it is seen as a fixed parameter
# - if "gamma" is a 2d-tensor (gamma.shape = [M,K]), it is seen as a "j"-variable "gamma_j"
#
# - if K == 1 , gamma is a scalar factor in front of a simple euclidean squared norm :
# WeightedSquareNorm( g, x-y ) = g * |x-y|^2
# - if K == D , gamma is a diagonal matrix:
# WeightedSquareNorm( g, x-y ) = < x-y, diag(g) * (x-y) >
# = \sum_d ( g[d] * ((x-y)[d])**2 )
# - if K == D*D, gamma is a (symmetric) matrix:
def benchmark(bench_name, N, dev, backend, loops = 10, enable_GC=True, fidelity=None) :
importlib.reload(torch)
device = torch.device(dev)
x_i = torch.randn(N, D, dtype=torch.float32, device=device, requires_grad=True)
y_j = torch.randn(N, D, dtype=torch.float32, device=device)
mu_i = torch.randn(N, 1, dtype=torch.float32, device=device)
nu_j = torch.randn(N, 1, dtype=torch.float32, device=device)
mu_i = mu_i.abs() ; nu_j = nu_j.abs()
mu_i = mu_i / mu_i.sum() ; nu_j = nu_j / nu_j.sum()
s2v = lambda x : torch.tensor([x], dtype=torch.float32, device=device)
if bench_name == "gaussian_conv" :
k = { "id" : Kernel("gaussian(x,y)"),
"gamma" : s2v( .25 ),
"backend" : backend, }
from pykeops.torch import kernel_product
_ = kernel_product(k, x_i, y_j, nu_j)
import gc
GC = 'gc.enable();' if enable_GC else 'pass;'
print("{:3} NxN-gaussian-convs, with N ={:7}: {:3}x".format(loops, N, loops), end="")
elapsed = timeit.Timer('_ = kernel_product(k,x_i,y_j,nu_j)', GC,
globals = locals(), timer = time.time).timeit(loops)
elif bench_name == "fidelities" :
from divergences import kernel_divergence, regularized_ot, hausdorff_divergence, sinkhorn_divergence
if fidelity == "energy_distance" :
params = ("energy", None)
c = kernel_divergence(mu_i,x_i, nu_j,y_j, k=params ) ; c.backward()
code = "c = kernel_divergence(mu_i,x_i, nu_j,y_j, k=params ) ; c.backward()"
elif fidelity == "gaussian_kernel" :
params = ("gaussian", .25)
c = kernel_divergence(mu_i,x_i, nu_j,y_j, k=params ) ; c.backward()
code = "c = kernel_divergence(mu_i,x_i, nu_j,y_j, k=params ) ; c.backward()"
elif fidelity == "log_kernel" :
params = {
"p" : 1,
"eps" : .1,
"nits" : 1,
"tol" : 0.,
}
c = hausdorff_divergence(mu_i,x_i, nu_j,y_j, **params ) ; c.backward()
code = "c = hausdorff_divergence(mu_i,x_i, nu_j,y_j, **params ) ; c.backward()"
elif fidelity == "hausdorff" :
params = {
"p" : 1,
"eps" : .1,
"nits" : 3,
"tol" : 0.,
}
c = hausdorff_divergence(mu_i,x_i, nu_j,y_j, **params ) ; c.backward()
code = "c = hausdorff_divergence(mu_i,x_i, nu_j,y_j, **params ) ; c.backward()"
elif fidelity == "sinkhorn" :
params = {
"p" : 1,
"eps" : .1,
"nits" : 20,
"assume_convergence" : True, # This is true in practice, and lets us win a x2 factor
"tol" : 0.,
}
c = sinkhorn_divergence(mu_i,x_i, nu_j,y_j, **params ) ; c.backward()
code = "c = sinkhorn_divergence(mu_i,x_i, nu_j,y_j, **params ) ; c.backward()"
import gc
GC = 'gc.enable();' if enable_GC else 'pass;'
print("{:3} NxN fidelities, with N ={:7}: {:3}x".format(loops, N, loops), end="")
elapsed = timeit.Timer(code, GC, globals = locals(), timer = time.time).timeit(loops)
print("{:3.6f}s".format(elapsed/loops))
return elapsed / loops
###############################################
# Kernel definition
# ---------------------
# Let's use a **Gaussian kernel** given through
#
# .. math::
#
# k(x_i,y_j) = \exp( -\|x - y\|_{\Gamma}^2) = \exp( - (x_i - y_j)^t \Gamma (x_i-y_j) ),
#
# which is equivalent to the KeOps formula ``exp(-WeightedSquareNorm(gamma, x_i-y_j ))``.
# Using the high-level :class:`pykeops.torch.Kernel`
# and :func:`pykeops.torch.kernel_product` wrappers, we can simply define:
params = {'id': Kernel('gaussian(x,y)')}
###############################################
# The precise meaning of the computation is then defined through the extra entry ``gamma`` of the ``params`` dictionary,
# which will is to be used as a 'metric multiplier'. Denoting ``D == x.shape[1] == y.shape[1]`` the size of the feature space, the integer ``K`` can be ``1``, ``D`` or ``D*D``. Rules are:
#
# - if ``gamma`` is a vector (``gamma.shape = [K]``), it is seen as a fixed parameter
# - if ``gamma`` is a 2d-tensor (``gamma.shape = [M,K]``), it is seen as a ``j``-variable
#
# N.B.: Beware of ``Shape([K]) != Shape([1,K])`` confusions!
###############################################
#
# Isotropic Kernels
# -----------------
#
# Pure numpy
g_numpy = np.matmul(np_kernel(x, y, sigma, kernel=k), b)
speed_numpy[k] = timeit.repeat(
'gnumpy = np.matmul( np_kernel(x, y, sigma, kernel=k), b)',
globals=globals(),
repeat=5,
number=1)
print('Time for NumPy: {:.4f}s'.format(
np.median(speed_numpy[k])))
# Vanilla pytorch (with cuda if available, and cpu otherwise)
try:
from pykeops.torch import Kernel, kernel_product
params = {
'id': Kernel(k + '(x,y)'),
'gamma': 1. / (sigmac**2),
'backend': 'pytorch',
}
g_pytorch = kernel_product(params, xc, yc, bc, mode='sum').cpu()
torch.cuda.synchronize()
speed_pytorch[k] = np.array(
timeit.repeat(
"kernel_product(params, xc, yc, bc, mode='sum'); torch.cuda.synchronize()",
globals=globals(),
repeat=REPEAT,
number=4)) / 4
print('Time for PyTorch: {:.4f}s'.format(
np.median(speed_pytorch[k])),
globals=globals(),
timer=time.time).timeit(LOOPS)
print("Time for keops specific: {:.4f}s".format(speed_pykeops),
end="")
print(" (absolute error: ", np.max(np.abs(g1 - gnumpy)),
")")
except:
pass
else:
print("Time for keops specific: skipping (no Gpu detected)")
# keops + pytorch : generic tiled implementation (with cuda if available else uses cpu)
try:
# Define a kernel: Wrap it (and its parameters) into a JSON dict structure
mode = "sum"
kernel = Kernel(k + "(x,y)")
params = {
"id": kernel,
"gamma": 1. / sigmac**2,
"backend": "auto",
}
aKxy_b = torch.dot(
ac.view(-1),
kernel_product(params, xc, yc, bc, mode=mode).view(-1))
g3 = torch.autograd.grad(aKxy_b, xc, create_graph=False)[0].cpu()
speed_keops = timeit.Timer(
'g3 = torch.autograd.grad(torch.dot(ac.view(-1), kernel_product( params, xc,yc,bc, mode=mode).view(-1)), xc, create_graph=False)[0]',
GC,
globals=globals(),
timer=time.time).timeit(LOOPS)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# N.B.: PyTorch's default dtype is float32
a = torch.randn(npoints_x, dimsignal, device=device)
x = torch.randn(npoints_x, dimpoints, requires_grad=True, device=device)
y = torch.randn(npoints_y, dimpoints, device=device)
b = torch.randn(npoints_y, dimsignal, device=device)
####################################################################
# A Gaussian convolution
# ----------------------
# Wrap the kernel's parameters into a JSON dict structure
sigma = scal_to_var(-1.5)
params = {
'id': Kernel('gaussian(x,y)'),
'gamma': .5 / sigma**2,
}
####################################################################
# Test, using both **PyTorch** and **KeOps** (online) backends:
axes = plt.subplot(2, 1, 1), plt.subplot(2, 1, 2)
for backend, linestyle, label in [("auto", "-", "KeOps"),
("pytorch", "--", "PyTorch")]:
Kxy_b = kernel_product(params, x, y, b, backend=backend)
aKxy_b = scalprod(a, Kxy_b)
# Computing a gradient is that easy - we can also use the 'aKxy_b.backward()' syntax.
# Notice the 'create_graph=True', which will allow us to compute
import torch
from pykeops.torch import Kernel, kernel_product
s2v = lambda x : torch.tensor([x])
Nt, Nx = 501, 201
t = torch.linspace( 0, 1, Nt ).view(-1,1)
x_i = torch.linspace( .2, .35, Nx).view(-1,1)
y_j = torch.linspace( .65, .8, Nx).view(-1,1)
mu_i = .15 * torch.ones( Nx ).view(-1,1) / Nx
nu_j = .15 * torch.ones( Nx ).view(-1,1) / Nx
kernels = {
"gaussian" : {
"id" : Kernel("gaussian(x,y)"),
"gamma" : s2v( 1 / .1**2 ),
},
"laplacian" : {
"id" : Kernel("laplacian(x,y)"),
"gamma" : s2v( 1 / .1**2 ),
},
"energy_distance" : {
"id" : Kernel("-distance(x,y)"),
"gamma" : s2v( 1. ),
},
}
fs = [ t ]
for name, kernel in kernels.items() :
f = kernel_product( kernel, t, x_i, mu_i) \
from scipy import misc
from time import time
tensor = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor
# ================================================================================================
# ================================ The Sinkhorn algorithm ======================================
# ================================================================================================
# Parameters of our optimal transport cost -------------------------------------------------------
a = 1 # Use a cost function "C(x,y) = |x-y|^a"
# The Sinkhorn algorithm relies on a kernel "k(x,y) = exp(-C(x,y))" :
kernel_names = { 1 : "laplacian" , # exp(-|x|), Earth-mover distance
2 : "gaussian" } # exp(-|x|^2), Quadratic "Wasserstein" distance
params = {
"id" : Kernel( kernel_names[a] + "(x,y)" ),
# Default set of parameters for the Sinkhorn cost - they make sense on the unit square/cube:
"nits" : 30, # Number of iterations in the Sinkhorn loop
"epsilon" : tensor([ .05**a ]), # Regularization strength, homogeneous to C(x,y)
}
# The Sinkhorn loop ------------------------------------------------------------------------------
dot = lambda a,b : torch.dot(a.view(-1), b.view(-1))
def OT_distance(params, Mu, Nu) :
"""
Computes an optimal transport cost using the Sinkhorn algorithm, stabilized in the log-domain.
See the section 4.4 of
"Gabriel Peyré and Marco Cuturi, Computational Optimal Transport, ArXiv:1803.00567, 2018"
for reference.
"""
a = torch.randn(npoints_x,dimsignal, requires_grad=True, device=device)
x = torch.randn(npoints_x,dimpoints, requires_grad=True, device=device)
y = torch.randn(npoints_y,dimpoints, requires_grad=True, device=device)
b = torch.randn(npoints_y,dimsignal, requires_grad=True, device=device)
#--------------------------------------------------------------#
# A first Kernel #
#--------------------------------------------------------------#
# N.B.: "sum" is default, "lse" is for "log-sum-exp"
modes = ["sum", "lse"] if dimsignal==1 else ["sum"]
# Wrap the kernel's parameters into a JSON dict structure
sigma = scal_to_var(-1.5)
params = {
"id" : Kernel("gaussian(x,y)"),
"gamma" : 1./sigma**2,
"mode" : "sum",
}
# Test, using a pytorch or keops
for mode in modes :
params["mode"] = mode
print("Mode :", mode, "========================================")
for backend in ["pytorch", "auto"] :
params["backend"] = backend
print("Backend :", backend, "--------------------------")
Kxy_b = kernel_product( params, x,y,b )
aKxy_b = scalprod(a, Kxy_b)