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 = "")