Skip to content

Latest commit

 

History

History
74 lines (56 loc) · 2.28 KB

File metadata and controls

74 lines (56 loc) · 2.28 KB

Kernel methods in PyTorch

Fast implementations of standard utilities for kernel machines

Installation

pip install -I git+https://github.com/parthe/torchkernels

Requires a PyTorch installation

Stable behavior

Currently this code has been tested with n=10,000 samples.
with Python 3.11 and PyTorch 2.4

Test installation with Laplacian kernel

import torch
from torchkernels.kernels.radial import laplacian, LaplacianKernel

n, p, d = 300, 200, 100 # number of samples, centers, dimensions

DEV = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")    

X = torch.randn(n, d, device=DEV)
Z = torch.randn(p, d, device=DEV)

kernel_matrix1 = laplacian(X, Z, length_scale=1.)

K = LaplacianKernel(length_scale=1.)
kernel_matrix2 = K(X, Z)

torch.testing.assert_close(kernel_matrix1, kernel_matrix2, msg='Laplacian test failed')
print('Laplacian test complete!')

Example: Differentiating the kernel function

import torch
from torchkernels.kernels.radial import laplacian
from torch.func import vmap, grad, jacrev

n, p, d, c = 300, 200, 100, 3 # number of samples, centers, dimensions, outputs

DEV = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")    

X = torch.randn(n, d, device=DEV)
Z = torch.randn(p, d, device=DEV)
a = torch.randn(p, device=DEV)
A = torch.randn(p, c, device=DEV)

f = lambda x: laplacian(x, Z, length_scale=1., in_place=False).squeeze() @ a
print(vmap(grad(f))(X).shape)
# torch.Size([300, 100])

F = lambda x: laplacian(x, Z, length_scale=1., in_place=False).squeeze() @ A
print(vmap(jacrev(F))(X).shape)
# torch.Size([300, 3, 100])

Random features

See an example of Logistic regression with random features of the Laplacian kernel.

Currently supported Kernels

  • Laplacian, Gaussian, Dispersal (Exponential power kernel)
  • Normalized dot-product kernel for arbitrary functions
  • Neural Network Gaussian Process (NNGP) and Tangent Kernel (NTK) with ReLU activations

Other utilities

  • extracting top eigenvectors of a kernel matrix
  • Random feature maps for:
    • Gaussian kernel
    • Laplacian kernel
    • Matern kernel
    • Exponential-power kernel $K(x,z) = \exp(-|x-z|^\gamma)$
  • Differentiable models