313 lines
13 KiB
Python
313 lines
13 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import numpy as np
|
|
from .spline import *
|
|
from .utils import sparse_mask
|
|
|
|
|
|
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', sparse_init=False):
|
|
''''
|
|
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):
|
|
if sparse_init:
|
|
mask = sparse_mask(in_dim, out_dim)
|
|
else:
|
|
mask = 1.
|
|
self.scale_base = torch.nn.Parameter(torch.ones(in_dim, out_dim, device=device) * scale_base * mask).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 * mask).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
|