[KAN学习日记](1)构建基本的网络框架以及网格细化

教程地址:https://github.com/KindXiaoming/pykan/blob/master/README.md
本系列文章用于记录自己学习过程

一.如何利用网格细化最大化KAN拟合函数的能力

(1)初始化模型和创建数据集

from kan import *

# 初始化KAN，设置网格为3
model = KAN(width=[2, 1, 1], grid=3, k=3)

# 创建数据集
f = lambda x: torch.exp(torch.sin(torch.pi * x[:, [0]]) + x[:, [1]] ** 2)
dataset = create_dataset(f, n_var=2)

(2)使用不同网格的KAN模型对比
训练KAN（网格=3）

model.train(dataset, opt="LBFGS", steps=20)

现在初始化一个网格大小为10的KAN模型，并从之前的模型进行初始化：

# initialize a more fine-grained KAN with G=10
model2 = KAN(width=[2,1,1], grid=10, k=3)
# initialize model2 from model
model2.initialize_from_another_model(model, dataset['train_input']);

用更细粒度的KAN模型训练数据集：(网格=10)

model2.train(dataset, opt="LBFGS", steps=20);

两个不同细粒度KAN模型loss对比

损失变低，说明效果更好.现在可以迭代地使网格更加细化。
(3)迭代网格细化
我们可以通过逐步增加网格大小，观察损失变化情况：

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

train_losses = []
test_losses = []
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 += results['train_loss']
    test_losses += results['test_loss']

训练过程中的损失变化如下：

(4)可视化展现训练损失和测试损失的动态变化

plt.plot(train_losses)
plt.plot(test_losses)
plt.legend(['train', 'test'])
plt.ylabel('RMSE')
plt.xlabel('step')
plt.yscale('log')
plt.show()

最后，我们可以观察参数数量与损失的关系：

n_params = 3 * grids
train_vs_G = train_losses[(steps-1)::steps]
test_vs_G = test_losses[(steps-1)::steps]

plt.plot(n_params, train_vs_G, marker="o")
plt.plot(n_params, test_vs_G, marker="o")
plt.plot(n_params, 100 * n_params**(-4.), ls="--", color="black")
plt.xscale('log')
plt.yscale('log')
plt.legend(['train', 'test', r'$N^{-4}$'])
plt.xlabel('number of params')
plt.ylabel('RMSE')
plt.show()

二.如何使用深层KAN模型来解决复杂任务

Kolmogorov-Arnold（KA）定理原本表示仅需要两层函数组合（内函数和外函数），但这些函数可能是非光滑甚至是分形的。我们将KA表示法推广到任意深度。一个使用光滑激活函数的两层KAN模型无法处理的例子需要至少三层KAN模型才能完成。
例如:

在本教程中，我们将展示如何使用深层KAN模型来解决复杂任务，并比较两层和三层KAN模型的性能。

(1)创建和训练三层KAN模型
创建模型：创建一个2D输入、1D输出，5个隐藏神经元的三层KAN模型，使用三次样条（k=3）和5个网格间隔（grid=5）。

model = KAN(width=[4, 2, 1, 1], grid=3, k=3, seed=0)
f = lambda x: torch.exp((torch.sin(torch.pi * (x[:, [0]]**2 + x[:, [1]]**2)) + torch.sin(torch.pi * (x[:, [2]]**2 + x[:, [3]]**2))) / 2)
dataset = create_dataset(f, n_var=4, train_num=3000)

训练模型

model.train(dataset, opt="LBFGS", steps=20, lamb=0.001, lamb_entropy=2.)

模型可视化:

model.plot(beta=10)
plt.show()

调整网格并再次训练：

grids = [3, 5, 10, 20, 50]
train_rmse = []
test_rmse = []

for grid in grids:
    model = KAN(width=[4, 2, 1, 1], grid=grid, k=3, seed=0).initialize_from_another_model(model, dataset['train_input'])
    results = model.train(dataset, opt="LBFGS", steps=50, stop_grid_update_step=30)
    train_rmse.append(results['train_loss'][-1].item())
    test_rmse.append(results['test_loss'][-1].item())

可视化结果：

n_params = np.array(grids) * (4*2 + 2*1 + 1*1)
plt.plot(n_params, train_rmse, marker="o")
plt.plot(n_params, test_rmse, marker="o")
plt.plot(n_params, 10000*n_params**(-4.), color="black", ls="--")
plt.legend(['train', 'test', r'$N^{-4}$'], loc="lower left")
plt.xscale('log')
plt.yscale('log')
plt.show()

print(train_rmse)
print(test_rmse)

(2)两层KAN模型示例
创建和训练两层KAN模型
创建模型：创建一个2D输入、1D输出，9个隐藏神经元的两层KAN模型，使用三次样条（k=3）和5个网格间隔（grid=5）

model = KAN(width=[4, 9, 1], grid=3, k=3, seed=0)
f = lambda x: torch.exp((torch.sin(torch.pi * (x[:, [0]]**2 + x[:, [1]]**2)) + torch.sin(torch.pi * (x[:, [2]]**2 + x[:, [3]]**2))) / 2)
dataset = create_dataset(f, n_var=4, train_num=3000)

训练模型并可视化

model.train(dataset, opt="LBFGS", steps=20, lamb=0.001, lamb_entropy=2.)
model.plot(beta=10)
plt.show()

调整网格并再次训练：

grids = [3, 5, 10, 20, 50]
train_rmse = []
test_rmse = []

for grid in grids:
    model = KAN(width=[4, 9, 1], grid=grid, k=3, seed=0).initialize_from_another_model(model, dataset['train_input'])
    results = model.train(dataset, opt="LBFGS", steps=50, stop_grid_update_step=30)
    train_rmse.append(results['train_loss'][-1].item())
    test_rmse.append(results['test_loss'][-1].item())

可视化结果：

n_params = np.array(grids) * (4*9 + 9*1)
plt.plot(n_params, train_rmse, marker="o")
plt.plot(n_params, test_rmse, marker="o")
plt.plot(n_params, 300*n_params**(-2.), color="black", ls="--")
plt.legend(['train', 'test', r'$N^{-4}$'], loc="lower left")
plt.xscale('log')
plt.yscale('log')
plt.show()

print(train_rmse)
print(test_rmse)

通过上述示例可以看出，三层KAN模型在处理复杂任务时比两层KAN模型表现更好。这证明了增加模型深度能够显著提升模型的表达能力和性能。

通过不断调整网格参数和模型架构，我们可以进一步优化模型的表现。希望本教程能够帮助你更好地理解和应用深层KAN模型