import torch import torch.nn as nn import numpy as np from .spline import * class KANLayer(nn.Module): """ KANLayer class Attributes: ----------- in_dim: int input dimension out_dim: int output dimension size: int the number of splines = input dimension * output dimension k: int the piecewise polynomial order of splines grid: 2D torch.float grid points noises: 2D torch.float injected noises to splines at initialization (to break degeneracy) coef: 2D torch.tensor coefficients of B-spline bases scale_base: 1D torch.float magnitude of the residual function b(x) scale_sp: 1D torch.float mangitude of the spline function spline(x) base_fun: fun residual function b(x) mask: 1D torch.float mask of spline functions. setting some element of the mask to zero means setting the corresponding activation to zero function. grid_eps: float in [0,1] a hyperparameter used in update_grid_from_samples. When grid_eps = 0, the grid is uniform; when grid_eps = 1, the grid is partitioned using percentiles of samples. 0 < grid_eps < 1 interpolates between the two extremes. weight_sharing: 1D tensor int allow spline activations to share parameters lock_counter: int counter how many activation functions are locked (weight sharing) lock_id: 1D torch.int the id of activation functions that are locked device: str device Methods: -------- __init__(): initialize a KANLayer forward(): forward update_grid_from_samples(): update grids based on samples' incoming activations initialize_grid_from_parent(): initialize grids from another model get_subset(): get subset of the KANLayer (used for pruning) lock(): lock several activation functions to share parameters unlock(): unlock already locked activation functions """ def __init__(self, in_dim=3, out_dim=2, num=5, k=3, noise_scale=0.1, scale_base=1.0, scale_sp=1.0, base_fun=torch.nn.SiLU(), grid_eps=0.02, grid_range=[-1, 1], sp_trainable=True, sb_trainable=True, save_plot_data = True, device='cpu'): '''' initialize a KANLayer Args: ----- in_dim : int input dimension. Default: 2. out_dim : int output dimension. Default: 3. num : int the number of grid intervals = G. Default: 5. k : int the order of piecewise polynomial. Default: 3. noise_scale : float the scale of noise injected at initialization. Default: 0.1. scale_base : float the scale of the residual function b(x). Default: 1.0. scale_sp : float the scale of the base function spline(x). Default: 1.0. base_fun : function residual function b(x). Default: torch.nn.SiLU() grid_eps : float When grid_eps = 0, the grid is uniform; when grid_eps = 1, the grid is partitioned using percentiles of samples. 0 < grid_eps < 1 interpolates between the two extremes. Default: 0.02. grid_range : list/np.array of shape (2,) setting the range of grids. Default: [-1,1]. sp_trainable : bool If true, scale_sp is trainable. Default: True. sb_trainable : bool If true, scale_base is trainable. Default: True. device : str device Returns: -------- self Example ------- >>> model = KANLayer(in_dim=3, out_dim=5) >>> (model.in_dim, model.out_dim) (3, 5) ''' super(KANLayer, self).__init__() # size self.out_dim = out_dim self.in_dim = in_dim self.num = num self.k = k # shape: (size, num) ### grid size: (batch, in_dim, out_dim, G + 1) => (batch, in_dim, G + 2*k + 1) grid = torch.linspace(grid_range[0], grid_range[1], steps=num + 1)[None,:].expand(self.in_dim, num+1) grid = extend_grid(grid, k_extend=k) self.grid = torch.nn.Parameter(grid).requires_grad_(False) noises = (torch.rand(self.num+1, self.in_dim, self.out_dim) - 1 / 2) * noise_scale / num noises = noises.to(device) # shape: (size, coef) self.coef = torch.nn.Parameter(curve2coef(self.grid[:,k:-k].permute(1,0), noises, self.grid, k, device)) if isinstance(scale_base, float): self.scale_base = torch.nn.Parameter(torch.ones(in_dim, out_dim, device=device) * scale_base).requires_grad_(sb_trainable) # make scale trainable else: self.scale_base = torch.nn.Parameter(scale_base.to(device)).requires_grad_(sb_trainable) self.scale_sp = torch.nn.Parameter(torch.ones(in_dim, out_dim, device=device) * scale_sp).requires_grad_(sp_trainable) # make scale trainable self.base_fun = base_fun self.mask = torch.nn.Parameter(torch.ones(in_dim, out_dim, device=device)).requires_grad_(False) self.grid_eps = grid_eps ### remove weight_sharing & lock parts #self.weight_sharing = torch.arange(out_dim*in_dim).reshape(out_dim, in_dim) #self.lock_counter = 0 #self.lock_id = torch.zeros(out_dim*in_dim).reshape(out_dim, in_dim) self.device = device def forward(self, x): ''' KANLayer forward given input x Args: ----- x : 2D torch.float inputs, shape (number of samples, input dimension) Returns: -------- y : 2D torch.float outputs, shape (number of samples, output dimension) preacts : 3D torch.float fan out x into activations, shape (number of sampels, output dimension, input dimension) postacts : 3D torch.float the outputs of activation functions with preacts as inputs postspline : 3D torch.float the outputs of spline functions with preacts as inputs Example ------- >>> model = KANLayer(in_dim=3, out_dim=5) >>> x = torch.normal(0,1,size=(100,3)) >>> y, preacts, postacts, postspline = model(x) >>> y.shape, preacts.shape, postacts.shape, postspline.shape (torch.Size([100, 5]), torch.Size([100, 5, 3]), torch.Size([100, 5, 3]), torch.Size([100, 5, 3])) ''' batch = x.shape[0] # x: shape (batch, in_dim) => shape (size, batch) (size = out_dim * in_dim) #x = torch.einsum('ij,k->ikj', x, torch.ones(self.out_dim, device=self.device)).reshape(batch, self.size).permute(1, 0) preacts = x[:,None,:].clone().expand(batch, self.out_dim, self.in_dim) base = self.base_fun(x) # (batch, in_dim) y = coef2curve(x_eval=x, grid=self.grid, coef=self.coef, k=self.k, device=self.device) # y shape: (batch, in_dim, out_dim) postspline = y.clone().permute(0,2,1) # postspline shape: (batch, out_dim, in_dim) y = self.scale_base[None,:,:] * base[:,:,None] + self.scale_sp[None,:,:] * y y = self.mask[None,:,:] * y postacts = y.clone().permute(0,2,1) y = torch.sum(y, dim=1) # shape (batch, out_dim) return y, preacts, postacts, postspline def update_grid_from_samples(self, x): ''' update grid from samples Args: ----- x : 2D torch.float inputs, shape (number of samples, input dimension) Returns: -------- None Example ------- >>> model = KANLayer(in_dim=1, out_dim=1, num=5, k=3) >>> print(model.grid.data) >>> x = torch.linspace(-3,3,steps=100)[:,None] >>> model.update_grid_from_samples(x) >>> print(model.grid.data) tensor([[-1.0000, -0.6000, -0.2000, 0.2000, 0.6000, 1.0000]]) tensor([[-3.0002, -1.7882, -0.5763, 0.6357, 1.8476, 3.0002]]) ''' batch = x.shape[0] #x = torch.einsum('ij,k->ikj', x, torch.ones(self.out_dim, ).to(self.device)).reshape(batch, self.size).permute(1, 0) x_pos = torch.sort(x, dim=0)[0] y_eval = coef2curve(x_pos, self.grid, self.coef, self.k, device=self.device) num_interval = self.grid.shape[1] - 1 - 2*self.k ids = [int(batch / num_interval * i) for i in range(num_interval)] + [-1] grid_adaptive = x_pos[ids, :].permute(1,0) margin = 0.01 h = (grid_adaptive[:,[-1]] - grid_adaptive[:,[0]])/num_interval grid_uniform = grid_adaptive[:,[0]] + h * torch.arange(num_interval+1,)[None, :] grid = self.grid_eps * grid_uniform + (1 - self.grid_eps) * grid_adaptive self.grid.data = extend_grid(grid, k_extend=self.k) self.coef.data = curve2coef(x_pos, y_eval, self.grid, self.k, device=self.device) def initialize_grid_from_parent(self, parent, x): ''' update grid from a parent KANLayer & samples Args: ----- parent : KANLayer a parent KANLayer (whose grid is usually coarser than the current model) x : 2D torch.float inputs, shape (number of samples, input dimension) Returns: -------- None Example ------- >>> batch = 100 >>> parent_model = KANLayer(in_dim=1, out_dim=1, num=5, k=3) >>> print(parent_model.grid.data) >>> model = KANLayer(in_dim=1, out_dim=1, num=10, k=3) >>> x = torch.normal(0,1,size=(batch, 1)) >>> model.initialize_grid_from_parent(parent_model, x) >>> print(model.grid.data) tensor([[-1.0000, -0.6000, -0.2000, 0.2000, 0.6000, 1.0000]]) tensor([[-1.0000, -0.8000, -0.6000, -0.4000, -0.2000, 0.0000, 0.2000, 0.4000, 0.6000, 0.8000, 1.0000]]) ''' batch = x.shape[0] # preacts: shape (batch, in_dim) => shape (size, batch) (size = out_dim * in_dim) #x_eval = torch.einsum('ij,k->ikj', x, torch.ones(self.out_dim, ).to(self.device)).reshape(batch, self.size).permute(1, 0) x_eval = x pgrid = parent.grid # (in_dim, G+2*k+1) pk = parent.k y_eval = coef2curve(x_eval, pgrid, parent.coef, pk, device=self.device) '''print(x_pos.shape) sp2 = KANLayer(in_dim=1, out_dim=self.in_dim, k=1, num=x_pos.shape[1] - 2*self.k - 1, scale_base=0., device=self.device) print(sp2.grid[:,sp2.k:-sp2.k].shape, x_pos[:,self.k:-self.k].shape, sp2.grid.shape) sp2.coef.data = curve2coef(sp2.grid[:,sp2.k:-sp2.k], x_pos[:,self.k:-self.k], sp2.grid, k=1, device=self.device) y_eval = coef2curve(x_eval, parent.grid, parent.coef, parent.k, device=self.device) percentile = torch.linspace(-1, 1, self.num + 1).to(self.device) self.grid.data = sp2(percentile.unsqueeze(dim=1))[0].permute(1, 0)''' h = (pgrid[:,[-pk]] - pgrid[:,[pk]])/self.num grid = pgrid[:,[pk]] + torch.arange(self.num+1,) * h grid = extend_grid(grid, k_extend=self.k) self.grid.data = grid self.coef.data = curve2coef(x_eval, y_eval, self.grid, self.k, self.device) def get_subset(self, in_id, out_id): ''' get a smaller KANLayer from a larger KANLayer (used for pruning) Args: ----- in_id : list id of selected input neurons out_id : list id of selected output neurons Returns: -------- spb : KANLayer Example ------- >>> kanlayer_large = KANLayer(in_dim=10, out_dim=10, num=5, k=3) >>> kanlayer_small = kanlayer_large.get_subset([0,9],[1,2,3]) >>> kanlayer_small.in_dim, kanlayer_small.out_dim (2, 3) ''' spb = KANLayer(len(in_id), len(out_id), self.num, self.k, base_fun=self.base_fun, device=self.device) spb.grid.data = self.grid[in_id] spb.coef.data = self.coef[in_id][:,out_id] spb.scale_base.data = self.scale_base[in_id][:,out_id] spb.scale_sp.data = self.scale_sp[in_id][:,out_id] spb.mask.data = self.mask[in_id][:,out_id] spb.in_dim = len(in_id) spb.out_dim = len(out_id) return spb