Scribd
Upload
132 views

Kolmogorov-Arnold-Networks in Python

Uploaded by

Neel24787bose
Copyright
© © All Rights Reserved
Available Formats
Download as PDF, TXT or read online on Scribd
132 views

Kolmogorov-Arnold-Networks in Python

Uploaded by

Neel24787bose
Copyright
© © All Rights Reserved
Available Formats
Download as PDF, TXT or read online on Scribd
You are on page 1/ 8

KAN

June 23, 2024

1 KAN: Kolmogorov-Arnold Networks


By Cristian Leo
Link to Article

[1]: import numpy as np # linear algebra

import torch # PyTorch


import torch.nn as nn # Neural network module
from torchsummary import summary # Model summary
from kan import KAN, create_dataset # KAN model and dataset

import matplotlib.pyplot as plt # Plotting

import warnings
warnings.filterwarnings("ignore")

/Users/cristianleo/Documents/Documents - Cristian’s Laptop/GitHub/models-from-


scratch-python/venv/lib/python3.10/site-packages/threadpoolctl.py:1214:

1
RuntimeWarning:
Found Intel OpenMP ('libiomp') and LLVM OpenMP ('libomp') loaded at
the same time. Both libraries are known to be incompatible and this
can cause random crashes or deadlocks on Linux when loaded in the
same Python program.
Using threadpoolctl may cause crashes or deadlocks. For more
information and possible workarounds, please see
https://github.com/joblib/threadpoolctl/blob/master/multiple_openmp.md

warnings.warn(msg, RuntimeWarning)

[2]: # create dataset


f = lambda x: torch.exp(torch.sin(torch.pi*x[:,[0]]) + x[:,[1]]**2)
dataset = create_dataset(f, n_var=2)

[3]: # Create a tensor of points


x = torch.linspace(-2, 2, 100)

# Compute the expression


y = torch.exp(torch.sin(torch.pi*x) + x**2)

# Convert the result back to a numpy array


y = y.detach().numpy()

# Create the plot


plt.figure(figsize=(10, 6)) # Increase the size of the figure
plt.plot(x, y, 'r-', linewidth=2) # Change the line color to red and increase␣
↪the line width

# Add title and labels


plt.title('Plot of the Expression', fontsize=16)
plt.xlabel('x', fontsize=14)
plt.ylabel('y', fontsize=14)

# Add a grid
plt.grid(True)

# Show the plot


plt.show()

2
[4]: grids = np.array([5,10,20,50,100])

train_losses_kan = []
test_losses_kan = []
steps = 50
k = 3

for i in range(grids.shape[0]):
if i == 0:
model = KAN(width=[2,1,1], grid=grids[i], k=k)
if i != 0:
model = KAN(width=[2,1,1], grid=grids[i], k=k).
↪initialize_from_another_model(model, dataset['train_input'])

results = model.train(dataset, opt="LBFGS", steps=steps,␣


↪stop_grid_update_step=30)

train_losses_kan += results['train_loss']
test_losses_kan += results['test_loss']

print(f"Train RMSE: {results['train_loss'][-1]:.8f} | Test RMSE:␣


↪{results['test_loss'][-1]:.8f}")

train loss: 6.41e-03 | test loss: 6.56e-03 | reg: 2.78e+00 : 100%|��| 50/50
[00:05<00:00, 8.37it/s]
train loss: 2.87e-04 | test loss: 3.27e-04 | reg: 2.79e+00 : 100%|��| 50/50

3
[00:05<00:00, 9.92it/s]
train loss: 2.03e-05 | test loss: 2.38e-05 | reg: 2.79e+00 : 100%|��| 50/50
[00:04<00:00, 10.60it/s]
train loss: 4.74e-06 | test loss: 1.42e-05 | reg: 2.79e+00 : 100%|��| 50/50
[00:05<00:00, 9.87it/s]
train loss: 2.60e-06 | test loss: 1.37e-05 | reg: 2.79e+00 : 100%|��| 50/50
[00:04<00:00, 10.53it/s]
Train RMSE: 0.00000260 | Test RMSE: 0.00001374

[5]: plt.figure(figsize=(10, 6)) # Increase the size of the figure


plt.plot(train_losses_kan, 'r-', linewidth=2) # Change the line color to red␣
↪and increase the line width

plt.plot(test_losses_kan, 'b--', linewidth=2) # Change the line color to blue,␣


↪use dashed line style and increase the line width

plt.legend(['Train', 'Test'], fontsize=14)


plt.title('Train and Test Losses', fontsize=16)
plt.ylabel('RMSE', fontsize=14)
plt.xlabel('Step', fontsize=14)
plt.yscale('log')
plt.grid(True) # Add a grid
plt.show()

4
[6]: # Define the MLP
class MLP(nn.Module):
def __init__(self):
super(MLP, self).__init__()
self.layers = nn.Sequential(
nn.Linear(dataset['train_input'].shape[1], 64),
nn.ReLU(),
nn.Linear(64, 64),
nn.ReLU(),
nn.Linear(64, 1)
)

def forward(self, x):


return self.layers(x)

# Instantiate the model


model = MLP()

summary(model, input_size=(dataset['train_input'].shape[1],))

----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Linear-1 [-1, 64] 192
ReLU-2 [-1, 64] 0
Linear-3 [-1, 64] 4,160
ReLU-4 [-1, 64] 0
Linear-5 [-1, 1] 65
================================================================
Total params: 4,417
Trainable params: 4,417
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.00
Forward/backward pass size (MB): 0.00
Params size (MB): 0.02
Estimated Total Size (MB): 0.02
----------------------------------------------------------------

[7]: # Define the loss function and the optimizer


criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)

train_loss_mlp = []
test_loss_mlp = []

# Train the model

5
epochs = 250
for epoch in range(epochs):
optimizer.zero_grad()
output = model(dataset['train_input']).squeeze()
loss = criterion(output, dataset['train_label'])
loss.backward()
optimizer.step()
train_loss_mlp.append(loss.item()**0.5)

# Test the model


model.eval()
with torch.no_grad():
output = model(dataset['test_input']).squeeze()
loss = criterion(output, dataset['test_label'])
test_loss_mlp.append(loss.item()**0.5)

print(f'Epoch {epoch+1}/{epochs}, Train Loss: {train_loss_mlp[-1]:.2f},␣


↪Test Loss: {test_loss_mlp[-1]:.2f}', end='\r')

Epoch 250/250, Train Loss: 1.36, Test Loss: 1.39

[8]: plt.figure(figsize=(10, 6)) # Increase the size of the figure


plt.plot(train_loss_mlp, 'r-', linewidth=2) # Change the line color to red and␣
↪increase the line width

plt.plot(test_loss_mlp, 'b--', linewidth=2) # Change the line color to blue,␣


↪use dashed line style and increase the line width

plt.legend(['Train', 'Test'], fontsize=14)


plt.title('MLP RMSE Loss', fontsize=16)
plt.ylabel('RMSE', fontsize=14)
plt.xlabel('Step', fontsize=14)
plt.yscale('log')
plt.grid(True) # Add a grid
plt.show()

6
[9]: fig, ax = plt.subplots(1, 2, figsize=(15, 6))

# Plot the KAN model vs the MLP model on the training set
ax[0].plot(train_losses_kan, 'b-')
ax[0].plot(train_loss_mlp, 'r--')
ax[0].set_title('Training Set')
ax[0].set_xlabel('Step')
ax[0].set_ylabel('RMSE')
ax[0].legend(['KAN', 'MLP'])

# Plot the KAN model vs the MLP model on the test set
ax[1].plot(test_losses_kan, 'b-')
ax[1].plot(test_loss_mlp, 'r--')
ax[1].set_title('Test Set')
ax[1].set_xlabel('True Value')
ax[1].set_ylabel('Predicted Value')

plt.show()

7
8

You might also like