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, 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": 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 log_likelihoods(self, sample):
"""Log-density, sampled on a given point cloud."""
self.update_covariances()
return kernel_product(self.params,
sample,
self.mu,
self.weights_log(),
mode='lse')
def likelihoods(self, sample):
"""Samples the density on a given point cloud."""
self.update_covariances()
return kernel_product(self.params,
sample,
self.mu,
self.weights(),
mode='sum')
def log_likelihoods(self, sample):
"""Log-density, sampled on a given point cloud."""
self.update_covariances()
#print([obj.shape for obj in [self.params['gamma'], sample, self.mu, self.weights_log()]])
return kernel_product(self.params,
sample,
self.mu,
self.weights_log(),
mode='lse',
backend="pytorch")
def plot_kernel(params):
""" Samples 'x -> ∑_j b_j * k_j(x - y_j)' on the grid, and displays it as a heatmap. """
heatmap = kernel_product(params, x, y, 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))
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.
"""
# Instead of "gamma" or "sigma", the Sinkhorn algorithm is best understood in terms
# of a regularization strength "epsilon", which divides the cost.
eps = params["epsilon"]
# The kernel_product convention is that gamma is the *squared distance multiplier*
# (just like the \Sigma/2 of multivariate Gaussian laws)
params["gamma"] = 1 / eps**(2/a)
mu_i, x_i = Mu ; nu_j, y_j = Nu
mu_i, nu_j = mu_i.view(-1,1), nu_j.view(-1,1)
mu_log, nu_log = mu_i.log(), nu_j.log()
# Initialize the dual variables to zero:
U, V = torch.zeros_like(mu_log), torch.zeros_like(nu_log)
# The Sinkhorn loop... is best implemented in the log-domain ! ----------------------------
for it in range(params["nits"]) :
# Kernel products + pointwise divisions, in the log-domain.
# Mathematically speaking, we're alternating Kullback-Leibler projections.
# N.B.: By convention, U is the deformable source and V is the fixed target.
# If we break before convergence, it is thus important to finish
# with a "projection on the mu-constraint"!
V = - kernel_product(params, y_j, x_i, mu_log+U, mode="lse" )
U = - kernel_product(params, x_i, y_j, nu_log+V, mode="lse" )
# To compute the full mass of the regularized transport plan (used in the corrective term),
# we use a "bonus" kernel_product mode, "log_scaled". Using generic_sum would have been possible too.
Gamma1 = kernel_product(params, x_i, y_j, torch.ones_like(V), mu_log+U, nu_log+V, mode="log_scaled")
# The Sinkhorn cost is homogeneous to C(x,y)
return eps * ( dot(mu_i, U) + dot(nu_j, V) - dot( Gamma1, torch.ones_like(Gamma1) ) )
def showcase_params(params, title, ind):
"""Samples "x -> ∑_j b_j * k_j(x - y_j)" on the grid, and displays it as a heatmap."""
heatmap = kernel_product(params, x, y, b)
heatmap = heatmap.view(
res, res).cpu().numpy() # reshape as a "background" image
plt.subplot(2, 3, ind)
plt.imshow(-heatmap,
interpolation='bilinear',
origin='lower',
vmin=-1,
vmax=1,
cmap=cm.RdBu,
extent=(0, 1, 0, 1))
plt.title(title, fontsize=20)
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)
print("Kernel dot product : ", disp(aKxy_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
# higher order derivatives.
[grad_x, grad_y, grad_s] = grad(aKxy_b, [x, y, sigma], create_graph=True)
print("Gradient wrt. x : ", disp(grad_x[:2,:]) )
print("Gradient wrt. y : ", disp(grad_y[:2,:]) )
print("Gradient wrt. s : ", disp(grad_s) )
grad_x_norm = scalprod(grad_x, grad_x)
[grad_xx, grad_xy] = grad(grad_x_norm, [x,y], create_graph=True)
print("Arbitrary formula 1 : ", disp(grad_xx[:2,:]) )
print("Arbitrary formula 2 : ", disp(grad_xy[:2,:]) )
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
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])),
end='')
print(' (absolute error: ',
np.max(np.abs(g_pytorch.numpy() - g_numpy)), ')')
except:
print('Time for PyTorch: Not Done')
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)
print("Time for Keops+pytorch: {:.4f}s".format(speed_keops), end="")
print(" (absolute error: ",
np.max(np.abs(g3.data.numpy() - gnumpy)), ")")
except:
pass
# vanilla pytorch (with cuda if available else uses cpu)
try:
# Define a kernel: Wrap it (and its parameters) into a JSON dict structure
# 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
# higher order derivatives.
[grad_x, grad_s] = grad(aKxy_b, [x, sigma], create_graph=True)
grad_x_norm = scalprod(grad_x, grad_x)
[grad_xx] = grad(grad_x_norm, [x], create_graph=True)
print("Backend = {:^7}: cost = {:.4f}, grad wrt. s = {:.4f}".format(
label, aKxy_b.item(), grad_s.item()))
# Fancy display: plot the results next to each other.
axes[0].plot(grad_x.detach().cpu().numpy()[:40, 0], linestyle, label=label)
# 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. / torch.autograd.Variable(sigmac,
requires_grad=False).type(dtype)**2,
"backend":
"auto",
}
g1 = kernel_product(params, xc, yc, bc, mode=mode).cpu()
speed_pykeops_gen = timeit.Timer(
'g1 = kernel_product( params,xc,yc,bc, mode=mode).cpu()',
GC,
globals=globals(),
timer=time.time).timeit(LOOPS)
print(
"Time for keops generic: {:.4f}s".format(speed_pykeops_gen),
end="")
print(" (absolute error: ",
np.max(np.abs(g1.data.numpy() - gnumpy)), ")")
except:
pass
# vanilla pytorch (with cuda if available else uses cpu)
try:
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) \
- kernel_product( kernel, t, y_j, nu_j)
fs.append(f)
header = "t gaussian laplacian energy_distance"
lines = [ f.view(-1).data.cpu().numpy() for f in fs ]
data = np.stack(lines).T
np.savetxt("output/graphs/kernel_1D.csv", data, fmt='%-9.5f', header=header, comments = "")