Kolmogorov-Arnold Networks¶

from pyKAN import *
kan_model = KAN(
    width=[2,1,1], 
    grid=5, 
    k=3, 
    seed=2, 
    device=device
)

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

# 假设已经训练好一个KAN了
grids = [10,50,100]
for grid in grids:
    model.refine(grid)
    moldel.fit(dataset,steps=20)

model = model.prune(edge_th=1e-2)

model.fix_symbolic(0,0,0,'x')
model.fix_symbolic(0,0,1,'x')
model.fix_symbolic(0,1,0,'x')

add_symbolic('J0', torch.special.bessel_j0, c=1)

model.suggest_symbolic(0,0,0,a_range=(-40,40))

def batch_jacobian(func, model, x, create_graph=False):
    # x in shape (Batch, Length)
    def _func_sum(x):  
        return func(model, x).sum(dim=0)
    return autograd.functional.jacobian(_func_sum, x, create_graph=create_graph).permute(1,0,2)

def new_output(model, x):
    """计算模型的输出"""
    r = torch.sqrt(torch.sum(x**2, dim=1, keepdim=True))
    return model(x) * torch.exp(-r)

def laplacian(model, x):
    """计算拉普拉斯算子"""
    def sol_D1_fun(model,x):
        return batch_jacobian(new_output, model, x, create_graph=True)[:,0,:]

    sol_D2 = batch_jacobian(sol_D1_fun, model, x, create_graph=True)[:,:,:] 
    lap = torch.sum(torch.diagonal(sol_D2, dim1=1, dim2=2), dim=1, keepdim=True)
    return lap

class PhysicsInformedLoss(nn.Module):
    def __init__(self, lambda_data=1.0, lambda_pde=0.1, lambda_reg=0.001, device='cuda'):
        super().__init__()
        self.lambda_data = lambda_data
        self.lambda_pde = lambda_pde
        self.reg = lambda_reg
        self.mse = nn.MSELoss()
        self.device = device

    def forward(self, pred, target, x=None, model=None):
        # 数据驱动损失
        data_loss = self.mse(pred, target)

        # 物理约束损失
        pde_loss = torch.mean(schrodinger_residual(model, x)**2)

        # 正则化损失
        reg_term = model.reg('edge_forward_spline_n', lamb_l1=1.0, lamb_entropy=2.0, lamb_coef=0.1, lamb_coefdiff=0.1)

        # 总损失
        total_loss = self.lambda_data * data_loss + self.lambda_pde * pde_loss +self.reg * reg_term

        return total_loss

def train_kan_pinn(model, dataset, adam_steps=10, lbfgs_steps=5, batch_size=32): # 添加 batch_size 参数
    model.train()

    x_train = dataset['train_input']
    y_train = dataset['train_label']
    x_test = dataset['test_input']
    y_test = dataset['test_label']

    # --- 为 Adam 创建 DataLoader ---
    train_dataset = torch.utils.data.TensorDataset(x_train, y_train)
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

    # 创建损失函数
    criterion = PhysicsInformedLoss(lambda_data=1, lambda_pde=1,lambda_reg=0.00001)


    best_loss = float('inf')
    best_model_state = None

    total_steps = adam_steps + lbfgs_steps

    # 进度条
    from tqdm import tqdm

    # --- Adam 优化阶段 (批量训练) ---
    optimizer_adam = torch.optim.Adam(model.parameters(), lr=0.1) 
    scheduler_adam = lr_scheduler.StepLR(optimizer_adam, step_size=1, gamma=0.9)
    print("--- 开始 Adam 批量训练 ---")
    # 定义梯度裁剪阈值
    clip_value = 0.0000000001 # 可以调整这个值

    for step in range(adam_steps):

        model.train()
        epoch_loss = 0.0
        num_batches = 0
        # 使用 tqdm 包装 train_loader 以显示批次进度
        batch_pbar = tqdm(train_loader, desc=f'Adam 训练 (Epoch {step+1}/{adam_steps})', leave=False, ncols=100)
        for x_batch, y_batch in batch_pbar:
            optimizer_adam.zero_grad()
            x_batch.requires_grad_(True)
            r_batch = torch.norm(x_batch, dim=1)
            outputs = model(x_batch)*torch.exp(-r_batch)
            loss = criterion(outputs, y_batch, x_batch, model)
            loss.backward()

            # --- 添加梯度裁剪 ---
            torch.nn.utils.clip_grad_norm_(model.parameters(), clip_value)
            # 或者使用 clip_grad_value_
            # torch.nn.utils.clip_grad_value_(model.parameters(), clip_value)

            optimizer_adam.step()
            epoch_loss += loss.item()
            num_batches += 1
            # 获取当前学习率并显示
            current_lr = optimizer_adam.param_groups[0]['lr']
            batch_pbar.set_postfix({'Batch Loss': f"{loss.item():.6f}", 'LR': f"{current_lr:.6f}"}) # 显示当前批次损失和学习率

        # 在每个 Adam epoch (step) 结束后更新学习率
        scheduler_adam.step()

        avg_epoch_loss = epoch_loss / num_batches # 计算平均 epoch 损失 (可选，主要用于观察)

        # --- 每个 Adam epoch 后进行评估 ---
        model.eval()
        # 训练集损失 (在整个训练集上评估)
        x_train_eval = x_train.detach().requires_grad_(True) # 评估需要梯度以计算 PDE loss
        r_train_eval = torch.norm(x_train_eval, dim=1)
        train_pred = model(x_train_eval)*torch.exp(-r_train_eval)
        current_train_loss = criterion(train_pred, y_train, x_train_eval, model).item()
        train_losses.append(current_train_loss)

        # 测试集损失 (在整个测试集上评估)
        x_test_eval = x_test.detach().requires_grad_(True) # 评估需要梯度以计算 PDE loss
        r_test_eval = torch.norm(x_test_eval, dim=1)
        test_pred = model(x_test_eval)*torch.exp(-r_test_eval)
        current_test_loss = criterion(test_pred, y_test, x_test_eval, model).item()
        test_losses.append(current_test_loss)

        # 测试集损失评估后，更新最佳模型状态
        if current_test_loss < best_loss:
            best_loss = current_test_loss
            best_model_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
        print(
            f"Adam 训练 (Epoch {step+1}/{adam_steps}) - 训练损失: {current_train_loss:.6f}, 测试损失: {current_test_loss:.6f}"
        )

    print("--- 开始 LBFGS 微调 ---")
    for grid in [5,10,20]:
        if grid == 5:
            pass
        else:
            model.refine(grid)
        optimizer_lbfgs = LBFGS(model.parameters(), lr=1, history_size=10, line_search_fn="strong_wolfe", tolerance_grad=1e-32, tolerance_change=1e-32, tolerance_ys=1e-32)
        lbfgs_pbar = tqdm(range(lbfgs_steps), desc='LBFGS 微调', ncols=100) # LBFGS 的进度条
        for step in range(lbfgs_steps):
            current_total_step = adam_steps + step
            lbfgs_pbar.set_description(f'LBFGS 微调 (Step {step+1}/{lbfgs_steps})')
            model.train()

            def closure():
                optimizer_lbfgs.zero_grad()
                x_train.requires_grad_(True)
                r_train = torch.norm(x_train, dim=1)
                outputs = model(x_train)*torch.exp(-r_train)
                loss = criterion(outputs, y_train, x_train, model)
                loss.backward()
                return loss

            optimizer_lbfgs.step(closure)

            # --- 每个 LBFGS step 后进行评估 ---
            model.eval()

            # 训练集损失
            x_train_eval = x_train.detach().requires_grad_(True)
            r_train_eval = torch.norm(x_train_eval, dim=1)
            train_pred = model(x_train_eval)*torch.exp(-r_train_eval)
            current_train_loss = criterion(train_pred, y_train, x_train_eval, model).item()
            train_losses.append(current_train_loss)

            # 测试集损失
            x_test_eval = x_test.detach().requires_grad_(True)
            r_test_eval = torch.norm(x_test_eval, dim=1)
            test_pred = model(x_test_eval)*torch.exp(-r_test_eval)
            current_test_loss = criterion(test_pred, y_test, x_test_eval, model).item()
            test_losses.append(current_test_loss)

            # 测试集损失评估后，更新最佳模型状态
            if current_test_loss < best_loss:
                best_loss = current_test_loss
                best_model_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}

            if step%1==0:
                model.update_grid_from_samples(x_train_eval)

            # 更新 LBFGS 进度条
            lbfgs_pbar.set_postfix({
                '训练损失': f"{current_train_loss:.6f}",
                '测试损失': f"{current_test_loss:.6f}",
            })
            lbfgs_pbar.update(1)
        lbfgs_pbar.close()
        print("--- LBFGS 微调结束 ---")



    # --- 保存最佳模型 ---
    if best_model_state is not None:
        save_path = 'your_path'
        torch.save(best_model_state, save_path)


    # 返回训练结果
    results = {
        'train_loss': train_losses,
        'test_loss': test_losses,
        'best_loss': best_loss if best_model_state is not None else float('inf'), 
    }

    return results

if step % 1==0:
    model.update_grid_from_samples(x_train_eval)