yms-kan 0.0.9__py3-none-any.whl → 0.0.10__py3-none-any.whl

yms_kan/MLP.py

@@ -3,43 +3,44 @@ import torch.nn as nn
 import matplotlib.pyplot as plt
 import numpy as np
 from tqdm import tqdm
+
+from . import plot_tree
 from .LBFGS import LBFGS
 
 seed = 0
 torch.manual_seed(seed)
 
+
 class MLP(nn.Module):
-    
+
     def __init__(self, width, act='silu', save_act=True, seed=0, device='cpu'):
         super(MLP, self).__init__()
-        
+
         torch.manual_seed(seed)
-        
+
         linears = []
         self.width = width
         self.depth = depth = len(width) - 1
         for i in range(depth):
-            linears.append(nn.Linear(width[i], width[i+1]))
+            linears.append(nn.Linear(width[i], width[i + 1]))
         self.linears = nn.ModuleList(linears)
-        
+
         #if activation == 'silu':
         self.act_fun = torch.nn.SiLU()
         self.save_act = save_act
         self.acts = None
-        
+
         self.cache_data = None
-        
+
         self.device = device
         self.to(device)
-        
-        
+
     def to(self, device):
         super(MLP, self).to(device)
         self.device = device
-        
+
         return self
-        
-        
+
     def get_act(self, x=None):
         if isinstance(x, dict):
             x = x['train_input']
@@ -52,23 +53,23 @@ class MLP(nn.Module):
         self.save_act = True
         self.forward(x)
         self.save_act = save_act
-        
+
     @property
     def w(self):
         return [self.linears[l].weight for l in range(self.depth)]
-        
+
     def forward(self, x):
-        
+
         # cache data
         self.cache_data = x
-        
+
         self.acts = []
         self.acts_scale = []
         self.wa_forward = []
         self.a_forward = []
-        
+
         for i in range(self.depth):
-            
+
             if self.save_act:
                 act = x.clone()
                 act_scale = torch.std(x, dim=0)
@@ -77,7 +78,7 @@ class MLP(nn.Module):
                 if i > 0:
                     self.acts_scale.append(act_scale)
                 self.wa_forward.append(wa_forward)
-            
+
             x = self.linears[i](x)
             if i < self.depth - 1:
                 x = self.act_fun(x)
@@ -85,9 +86,9 @@ class MLP(nn.Module):
                 if self.save_act:
                     act_scale = torch.std(x, dim=0)
                     self.acts_scale.append(act_scale)
-                
+
         return x
-    
+
     def attribute(self):
         if self.acts == None:
             self.get_act()
@@ -99,47 +100,47 @@ class MLP(nn.Module):
         node_score = torch.ones(self.width[-1]).requires_grad_(True).to(self.device)
         node_scores.append(node_score)
 
-        for l in range(self.depth,0,-1):
-
-            edge_score = torch.einsum('ij,i->ij', torch.abs(self.wa_forward[l-1]), node_score/(self.acts_scale[l-1]+1e-4))
+        for l in range(self.depth, 0, -1):
+            edge_score = torch.einsum('ij,i->ij', torch.abs(self.wa_forward[l - 1]),
+                                      node_score / (self.acts_scale[l - 1] + 1e-4))
             edge_scores.append(edge_score)
 
             # this might be improper for MLPs (although reasonable for KANs)
-            node_score = torch.sum(edge_score, dim=0)/torch.sqrt(torch.tensor(self.width[l-1], device=self.device))
+            node_score = torch.sum(edge_score, dim=0) / torch.sqrt(torch.tensor(self.width[l - 1], device=self.device))
             #print(self.width[l])
             node_scores.append(node_score)
 
         self.node_scores = list(reversed(node_scores))
         self.edge_scores = list(reversed(edge_scores))
         self.wa_backward = self.edge_scores
-    
+
     def plot(self, beta=3, scale=1., metric='w'):
         # metric = 'w', 'act' or 'fa'
-        
+
         if metric == 'fa':
             self.attribute()
-        
+
         depth = self.depth
         y0 = 0.5
-        fig, ax = plt.subplots(figsize=(3*scale,3*y0*depth*scale))
+        fig, ax = plt.subplots(figsize=(3 * scale, 3 * y0 * depth * scale))
         shp = self.width
-        
-        min_spacing = 1/max(self.width)
+
+        min_spacing = 1 / max(self.width)
         for j in range(len(shp)):
             N = shp[j]
             for i in range(N):
                 plt.scatter(1 / (2 * N) + i / N, j * y0, s=min_spacing ** 2 * 5000 * scale ** 2, color='black')
-                
-        plt.ylim(-0.1*y0,y0*depth+0.1*y0)
-        plt.xlim(-0.02,1.02)
+
+        plt.ylim(-0.1 * y0, y0 * depth + 0.1 * y0)
+        plt.xlim(-0.02, 1.02)
 
         linears = self.linears
-        
+
         for ii in range(len(linears)):
             linear = linears[ii]
             p = linear.weight
             p_shp = p.shape
-            
+
             if metric == 'w':
                 pass
             elif metric == 'act':
@@ -150,12 +151,15 @@ class MLP(nn.Module):
                 raise Exception('metric = \'{}\' not recognized. Choices are \'w\', \'act\', \'fa\'.'.format(metric))
             for i in range(p_shp[0]):
                 for j in range(p_shp[1]):
-                    plt.plot([1/(2*p_shp[0])+i/p_shp[0], 1/(2*p_shp[1])+j/p_shp[1]], [y0*(ii+1),y0*ii], lw=0.5*scale, alpha=np.tanh(beta*np.abs(p[i,j].cpu().detach().numpy())), color="blue" if p[i,j]>0 else "red")
-                    
+                    plt.plot([1 / (2 * p_shp[0]) + i / p_shp[0], 1 / (2 * p_shp[1]) + j / p_shp[1]],
+                             [y0 * (ii + 1), y0 * ii], lw=0.5 * scale,
+                             alpha=np.tanh(beta * np.abs(p[i, j].cpu().detach().numpy())),
+                             color="blue" if p[i, j] > 0 else "red")
+
         ax.axis('off')
-        
+
     def reg(self, reg_metric, lamb_l1, lamb_entropy):
-        
+
         if reg_metric == 'w':
             acts_scale = self.w
         if reg_metric == 'act':
@@ -164,7 +168,7 @@ class MLP(nn.Module):
             acts_scale = self.wa_backward
         if reg_metric == 'a':
             acts_scale = self.acts_scale
-        
+
         if len(acts_scale[0].shape) == 2:
             reg_ = 0.
 
@@ -178,9 +182,9 @@ class MLP(nn.Module):
                 entropy_row = - torch.mean(torch.sum(p_row * torch.log2(p_row + 1e-4), dim=1))
                 entropy_col = - torch.mean(torch.sum(p_col * torch.log2(p_col + 1e-4), dim=0))
                 reg_ += lamb_l1 * l1 + lamb_entropy * (entropy_row + entropy_col)
-                
+
         elif len(acts_scale[0].shape) == 1:
-            
+
             reg_ = 0.
 
             for i in range(len(acts_scale)):
@@ -193,20 +197,21 @@ class MLP(nn.Module):
                 reg_ += lamb_l1 * l1 + lamb_entropy * entropy
 
         return reg_
-    
+
     def get_reg(self, reg_metric, lamb_l1, lamb_entropy):
         return self.reg(reg_metric, lamb_l1, lamb_entropy)
-        
-    def fit(self, dataset, opt="LBFGS", steps=100, log=1, lamb=0., lamb_l1=1., lamb_entropy=2., loss_fn=None, lr=1., batch=-1,
-              metrics=None, in_vars=None, out_vars=None, beta=3, device='cpu', reg_metric='w', display_metrics=None):
+
+    def fit(self, dataset, opt="LBFGS", steps=100, log=1, lamb=0., lamb_l1=1., lamb_entropy=2., loss_fn=None, lr=1.,
+            batch=-1,
+            metrics=None, device='cpu', reg_metric='w', display_metrics=None):
 
         if lamb > 0. and not self.save_act:
             print('setting lamb=0. If you want to set lamb > 0, set =True')
-            
+
         old_save_act = self.save_act
         if lamb == 0.:
             self.save_act = False
-       
+
         pbar = tqdm(range(steps), desc='description', ncols=100)
 
         if loss_fn == None:
@@ -217,7 +222,8 @@ class MLP(nn.Module):
         if opt == "Adam":
             optimizer = torch.optim.Adam(self.parameters(), lr=lr)
         elif opt == "LBFGS":
-            optimizer = LBFGS(self.parameters(), lr=lr, history_size=10, line_search_fn="strong_wolfe", tolerance_grad=1e-32, tolerance_change=1e-32, tolerance_ys=1e-32)
+            optimizer = LBFGS(self.parameters(), lr=lr, history_size=10, line_search_fn="strong_wolfe",
+                              tolerance_grad=1e-32, tolerance_change=1e-32, tolerance_ys=1e-32)
 
         results = {}
         results['train_loss'] = []
@@ -252,10 +258,10 @@ class MLP(nn.Module):
             return objective
 
         for _ in pbar:
-            
-            if _ == steps-1 and old_save_act:
+
+            if _ == steps - 1 and old_save_act:
                 self.save_act = True
-            
+
             train_id = np.random.choice(dataset['train_input'].shape[0], batch_size, replace=False)
             test_id = np.random.choice(dataset['test_input'].shape[0], batch_size_test, replace=False)
 
@@ -274,9 +280,9 @@ class MLP(nn.Module):
                 loss.backward()
                 optimizer.step()
 
-            test_loss = loss_fn_eval(self.forward(dataset['test_input'][test_id].to(self.device)), dataset['test_label'][test_id].to(self.device))
-            
-            
+            test_loss = loss_fn_eval(self.forward(dataset['test_input'][test_id].to(self.device)),
+                                     dataset['test_label'][test_id].to(self.device))
+
             if metrics != None:
                 for i in range(len(metrics)):
                     results[metrics[i].__name__].append(metrics[i]().item())
@@ -287,7 +293,9 @@ class MLP(nn.Module):
 
             if _ % log == 0:
                 if display_metrics == None:
-                    pbar.set_description("| train_loss: %.2e | test_loss: %.2e | reg: %.2e | " % (torch.sqrt(train_loss).cpu().detach().numpy(), torch.sqrt(test_loss).cpu().detach().numpy(), reg_.cpu().detach().numpy()))
+                    pbar.set_description("| train_loss: %.2e | test_loss: %.2e | reg: %.2e | " % (
+                    torch.sqrt(train_loss).cpu().detach().numpy(), torch.sqrt(test_loss).cpu().detach().numpy(),
+                    reg_.cpu().detach().numpy()))
                 else:
                     string = ''
                     data = ()
@@ -299,9 +307,9 @@ class MLP(nn.Module):
                             raise Exception(f'{metric} not recognized')
                         data += (results[metric][-1],)
                     pbar.set_description(string % data)
-           
+
         return results
-    
+
     @property
     def connection_cost(self):
 
@@ -309,8 +317,9 @@ class MLP(nn.Module):
             cc = 0.
             for linear in self.linears:
                 t = torch.abs(linear.weight)
+
                 def get_coordinate(n):
-                    return torch.linspace(0,1,steps=n+1, device=self.device)[:n] + 1/(2*n)
+                    return torch.linspace(0, 1, steps=n + 1, device=self.device)[:n] + 1 / (2 * n)
 
                 in_dim = t.shape[0]
                 x_in = get_coordinate(in_dim)
@@ -318,44 +327,45 @@ class MLP(nn.Module):
                 out_dim = t.shape[1]
                 x_out = get_coordinate(out_dim)
 
-                dist = torch.abs(x_in[:,None] - x_out[None,:])
+                dist = torch.abs(x_in[:, None] - x_out[None, :])
                 cc += torch.sum(dist * t)
 
         return cc
-    
+
     def swap(self, l, i1, i2):
 
         def swap_row(data, i1, i2):
             data[i1], data[i2] = data[i2].clone(), data[i1].clone()
 
         def swap_col(data, i1, i2):
-            data[:,i1], data[:,i2] = data[:,i2].clone(), data[:,i1].clone()
+            data[:, i1], data[:, i2] = data[:, i2].clone(), data[:, i1].clone()
 
-        swap_row(self.linears[l-1].weight.data, i1, i2)
-        swap_row(self.linears[l-1].bias.data, i1, i2)
+        swap_row(self.linears[l - 1].weight.data, i1, i2)
+        swap_row(self.linears[l - 1].bias.data, i1, i2)
         swap_col(self.linears[l].weight.data, i1, i2)
-    
+
     def auto_swap_l(self, l):
 
         num = self.width[l]
         for i in range(num):
             ccs = []
             for j in range(num):
-                self.swap(l,i,j)
+                self.swap(l, i, j)
                 self.get_act()
                 self.attribute()
                 cc = self.connection_cost.detach().clone()
                 ccs.append(cc)
-                self.swap(l,i,j)
+                self.swap(l, i, j)
             j = torch.argmin(torch.tensor(ccs))
-            self.swap(l,i,j)
+            self.swap(l, i, j)
 
     def auto_swap(self):
         depth = self.depth
         for l in range(1, depth):
             self.auto_swap_l(l)
-            
+
     def tree(self, x=None, in_var=None, style='tree', sym_th=1e-3, sep_th=1e-1, skip_sep_test=False, verbose=False):
         if x == None:
             x = self.cache_data
-        plot_tree(self, x, in_var=in_var, style=style, sym_th=sym_th, sep_th=sep_th, skip_sep_test=skip_sep_test, verbose=verbose)
+        plot_tree(self, x, in_var=in_var, style=style, sym_th=sym_th, sep_th=sep_th, skip_sep_test=skip_sep_test,
+                  verbose=verbose)

yms_kan/MultKAN.py

@@ -794,6 +794,8 @@ class MultKAN(nn.Module):
         # >>> print(model(x, singularity_avoiding=True))
         # >>> print(model(x, singularity_avoiding=True, y_th=1.))
         """
+        # x = abs(torch.fft(x, dim=-1,norm='forward'))
+        # _,x = x.chunk(2,dim=-1)
         x = x[:, self.input_id.long()]
         assert x.shape[1] == self.width_in[0]
 
@@ -1063,7 +1065,7 @@ class MultKAN(nn.Module):
                     ha='center', va='center', transform=ax.transData)
 
     def plot(self, folder="./figures", beta=3, metric='backward', scale=0.5, tick=False, sample=False, in_vars=None,
-             out_vars=None, title=None, varscale=1.0):
+             out_vars=None, title=None, varscale=1.0, dpi=100):
         '''
         plot KAN
         
@@ -1164,7 +1166,7 @@ class MultKAN(nn.Module):
                                     s=400 * scale ** 2)
                     plt.gca().spines[:].set_color(color)
 
-                    plt.savefig(f'{folder}/sp_{l}_{i}_{j}.png', bbox_inches="tight", dpi=400)
+                    plt.savefig(f'{folder}/sp_{l}_{i}_{j}.png', bbox_inches="tight", dpi=dpi)
                     plt.close()
 
         def score2alpha(score):
@@ -1647,7 +1649,7 @@ class MultKAN(nn.Module):
 
             if save_fig and _ % save_fig_freq == 0:
                 self.plot(folder=img_folder, in_vars=in_vars, out_vars=out_vars, title="Step {}".format(_), beta=beta)
-                plt.savefig(img_folder + '/' + str(_) + '.jpg', bbox_inches='tight', dpi=100)
+                plt.savefig(img_folder + '/' + str(_) + '.jpg', bbox_inches='tight')
                 plt.close()
                 self.save_act = save_act
 
@@ -1857,7 +1859,7 @@ class MultKAN(nn.Module):
             if save_fig and epoch % save_fig_freq == 0:
                 self.plot(folder=img_folder, in_vars=in_vars, out_vars=out_vars, title="Step {}".format(epoch),
                           beta=beta)
-                plt.savefig(img_folder + '/' + str(epoch) + '.jpg', bbox_inches='tight', dpi=100)
+                plt.savefig(img_folder + '/' + str(epoch) + '.jpg', bbox_inches='tight')
                 plt.close()
                 self.save_act = save_act
 

yms_kan/train_eval_utils.py

@@ -13,11 +13,11 @@ from yms_kan.plotting import plot_confusion_matrix
 from yms_kan.tool import initialize_results_file, append_to_results_file, calculate_metric
 
 
-def train_val(model, dataset: dict, batch_size, batch_size_test, save_path, txt_file=None, opt="LBFGS", epochs=100,
+def train_val(model, dataset: dict, batch_size, batch_size_test, save_path=None, txt_file=None, opt="LBFGS", epochs=100,
               lamb=0.,
               lamb_l1=1., label=None, class_dict=None, lamb_entropy=2., lamb_coef=0.,
               lamb_coefdiff=0., update_grid=True, grid_update_num=10, loss_fn=None, lr=1., start_grid_update_step=-1,
-              stop_grid_update_step=100,
+              stop_grid_update_epoch=100,
               save_fig=False, in_vars=None, out_vars=None, beta=3, save_fig_freq=1, img_folder='./video',
               singularity_avoiding=False, y_th=1000., reg_metric='edge_forward_spline_n'):
     best = -1
@@ -42,7 +42,6 @@ def train_val(model, dataset: dict, batch_size, batch_size_test, save_path, txt_
     else:
         loss_fn = loss_fn
 
-    grid_update_freq = int(stop_grid_update_step / grid_update_num)
 
     if opt == "Adam":
         optimizer = torch.optim.Adam(model.get_params(), lr=lr)
@@ -55,9 +54,11 @@ def train_val(model, dataset: dict, batch_size, batch_size_test, save_path, txt_
     lr_scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=5, min_lr=1e-9)
 
     results = {'train_losses': [], 'val_losses': [], 'accuracies': [], 'precisions': [], 'recalls': [], 'f1-scores': [],
-               'lrs': [], 'all_predictions': [], 'all_labels': [], 'regularize': []}
+               'lrs': [], 'regularize': []}  # , 'all_predictions': [], 'all_labels': []
 
     steps = math.ceil(dataset['train_input'].shape[0] / batch_size)
+    stop_grid_update_step = stop_grid_update_epoch * steps
+    grid_update_freq = int(stop_grid_update_step / grid_update_num)
 
     train_loss = torch.zeros(1).to(model.device)
     reg_ = torch.zeros(1).to(model.device)
@@ -85,7 +86,6 @@ def train_val(model, dataset: dict, batch_size, batch_size_test, save_path, txt_
             os.makedirs(img_folder)
 
     for epoch in range(epochs):
-
         if epoch == epochs - 1 and old_save_act:
             model.save_act = True
 
@@ -172,7 +172,9 @@ def train_val(model, dataset: dict, batch_size, batch_size_test, save_path, txt_
                 best = m["f1-score"]
                 results['all_predictions'] = all_predictions
                 results['all_labels'] = all_labels
-                plot_confusion_matrix(all_labels, all_predictions, class_dict, save_path)
+
+                # if save_path is not None:
+                #     plot_confusion_matrix(all_labels, all_predictions, class_dict, save_path)
                 # if save_path is not None:
                 #     model.saveckpt(path=(os.path.join(save_path, 'save_model') + '/' + 'model'))
             if txt_file is not None:
@@ -196,156 +198,156 @@ def train_val(model, dataset: dict, batch_size, batch_size_test, save_path, txt_
     model.symbolic_enabled = old_symbolic_enabled
     return results
 
-# def train_val(model, dataset: dict, batch_size, batch_size_test, opt="LBFGS", epochs=100, lamb=0.,
-#               lamb_l1=1., label=None, lamb_entropy=2., lamb_coef=0.,
-#               lamb_coefdiff=0., update_grid=True, grid_update_num=10, loss_fn=None, lr=1., start_grid_update_step=-1,
-#               stop_grid_update_step=100,
-#               save_fig=False, in_vars=None, out_vars=None, beta=3, save_fig_freq=1, img_folder='./video',
-#               singularity_avoiding=False, y_th=1000., reg_metric='edge_forward_spline_n'):
-#     # result_info = ['epoch','train_losses', 'val_losses', 'regularize', 'accuracies',
-#     #                'precisions', 'recalls', 'f1-scores']
-#     # initialize_results_file(results_file, result_info)
-#     all_predictions = []
-#     all_labels = []
-#     if lamb > 0. and not model.save_act:
-#         print('setting lamb=0. If you want to set lamb > 0, set model.save_act=True')
-#
-#     old_save_act, old_symbolic_enabled = model.disable_symbolic_in_fit(lamb)
-#     if label is not None:
-#         label = label.to(model.device)
-#
-#     if loss_fn is None:
-#         loss_fn = lambda x, y: torch.mean((x - y) ** 2)
-#     else:
-#         loss_fn = loss_fn
-#
-#     grid_update_freq = int(stop_grid_update_step / grid_update_num)
-#
-#     if opt == "Adam":
-#         optimizer = torch.optim.Adam(model.get_params(), lr=lr)
-#     elif opt == "LBFGS":
-#         optimizer = LBFGS(model.get_params(), lr=lr, history_size=10, line_search_fn="strong_wolfe",
-#                           tolerance_grad=1e-32, tolerance_change=1e-32, tolerance_ys=1e-32)
-#     else:
-#         optimizer = torch.optim.SGD(model.get_params(), lr=lr)
-#
-#     lr_scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=5, min_lr=1e-9)
-#
-#     results = {'train_losses': [], 'val_losses': [], 'regularize': [], 'accuracies': [],
-#                'precisions': [], 'recalls': [], 'f1-scores': []}
-#
-#     steps = math.ceil(dataset['train_input'].shape[0] / batch_size)
-#
-#     train_loss = torch.zeros(1).to(model.device)
-#     reg_ = torch.zeros(1).to(model.device)
-#
-#     def closure():
-#         nonlocal train_loss, reg_
-#         optimizer.zero_grad()
-#         pred = model.forward(batch_train_input, singularity_avoiding=singularity_avoiding, y_th=y_th)
-#         loss = loss_fn(pred, batch_train_label)
-#         if model.save_act:
-#             if reg_metric == 'edge_backward':
-#                 model.attribute()
-#             if reg_metric == 'node_backward':
-#                 model.node_attribute()
-#             reg_ = model.get_reg(reg_metric, lamb_l1, lamb_entropy, lamb_coef, lamb_coefdiff)
-#         else:
-#             reg_ = torch.tensor(0.)
-#         objective = loss + lamb * reg_
-#         train_loss = (train_loss * batch_num + objective.detach()) / (batch_num + 1)
-#         objective.backward()
-#         return objective
-#
-#     if save_fig:
-#         if not os.path.exists(img_folder):
-#             os.makedirs(img_folder)
-#
-#     for epoch in range(epochs):
-#
-#         if epoch == epochs - 1 and old_save_act:
-#             model.save_act = True
-#
-#         if save_fig and epoch % save_fig_freq == 0:
-#             save_act = model.save_act
-#             model.save_act = True
-#
-#         train_indices = np.arange(dataset['train_input'].shape[0])
-#         np.random.shuffle(train_indices)
-#         train_pbar = tqdm(range(steps), desc=f'Epoch {epoch + 1}/{epochs} Training', file=sys.stdout)
-#         for batch_num in train_pbar:
-#             step = epoch * steps + batch_num + 1
-#             i = batch_num * batch_size
-#             batch_train_id = train_indices[i:i + batch_size]
-#             batch_train_input = dataset['train_input'][batch_train_id].to(model.device)
-#             batch_train_label = dataset['train_label'][batch_train_id].to(model.device)
-#
-#             if step % grid_update_freq == 0 and step < stop_grid_update_step and update_grid and step >= start_grid_update_step:
-#                 model.update_grid(batch_train_input)
-#
-#             if opt == "LBFGS":
-#                 optimizer.step(closure)
-#
-#             else:
-#                 optimizer.zero_grad()
-#                 pred = model.forward(batch_train_input, singularity_avoiding=singularity_avoiding,
-#                                      y_th=y_th)
-#                 loss = loss_fn(pred, batch_train_label)
-#                 if model.save_act:
-#                     if reg_metric == 'edge_backward':
-#                         model.attribute()
-#                     if reg_metric == 'node_backward':
-#                         model.node_attribute()
-#                     reg_ = model.get_reg(reg_metric, lamb_l1, lamb_entropy, lamb_coef, lamb_coefdiff)
-#                 else:
-#                     reg_ = torch.tensor(0.)
-#                 loss = loss + lamb * reg_
-#                 train_loss = (train_loss * batch_num + loss.detach()) / (batch_num + 1)
-#                 loss.backward()
-#                 optimizer.step()
-#             train_pbar.set_postfix(loss=train_loss.item())
-#
-#         print(f'{epoch}/{epochs}:train_loss:{train_loss.item()}')
-#         val_loss = torch.zeros(1).to(model.device)
-#         with torch.no_grad():
-#             test_indices = np.arange(dataset['test_input'].shape[0])
-#             np.random.shuffle(test_indices)
-#             test_steps = math.ceil(dataset['test_input'].shape[0] / batch_size_test)
-#             test_pbar = tqdm(range(test_steps), desc=f'Epoch {epoch + 1}/{epochs} Validation', file=sys.stdout)
-#             for batch_num in test_pbar:
-#                 i = batch_num * batch_size_test
-#                 batch_test_id = test_indices[i:i + batch_size_test]
-#                 batch_test_input = dataset['test_input'][batch_test_id].to(model.device)
-#                 batch_test_label = dataset['test_label'][batch_test_id].to(model.device)
-#
-#                 outputs = model.forward(batch_test_input, singularity_avoiding=singularity_avoiding,
-#                                         y_th=y_th)
-#
-#                 loss = loss_fn(outputs, batch_test_label)
-#
-#                 val_loss = (val_loss * batch_num + loss.detach()) / (batch_num + 1)
-#                 test_pbar.set_postfix(loss=loss.item(), val_loss=val_loss.item())
-#                 if label is not None:
-#                     diffs = torch.abs(outputs - label)
-#                     closest_indices = torch.argmin(diffs, dim=1)
-#                     closest_values = label[closest_indices]
-#                     all_predictions.extend(closest_values.detach().cpu().numpy())
-#                     all_labels.extend(batch_test_label.detach().cpu().numpy())
-#
-#             lr_scheduler.step(val_loss)
-#
-#         results['train_losses'].append(train_loss.cpu().item())
-#         results['val_losses'].append(val_loss.cpu().item())
-#         results['regularize'].append(reg_.cpu().item())
-#
-#         if save_fig and epoch % save_fig_freq == 0:
-#             model.plot(folder=img_folder, in_vars=in_vars, out_vars=out_vars, title="Step {}".format(epoch),
-#                        beta=beta)
-#             plt.savefig(img_folder + '/' + str(epoch) + '.jpg', bbox_inches='tight', dpi=100)
-#             plt.close()
-#             model.save_act = save_act
-#
-#     # append_to_results_file(results_file, results, result_info)
-#     model.log_history('fit')
-#     model.symbolic_enabled = old_symbolic_enabled
-#     return results
+
+def fit(model, dataset, batch_size, opt="LBFGS", epochs=100, lamb=0., lamb_l1=1., lamb_entropy=2., loss_fn=None,
+        lr=1., label=None, class_dict=None,
+        txt_file=None,
+        reg_metric='w'):
+    best = -1
+    column_order = ['epoch', 'train_losses', 'val_losses', 'accuracies', 'precisions', 'recalls',
+                    'f1-scores', 'lrs']
+    custom_column_widths = {'epoch': 5, 'train_loss': 12, 'val_loss': 10, 'accuracy': 10, 'precision': 9,
+                            'recall': 7,
+                            'f1-score': 8,
+                            'lr': 3}
+    if txt_file is not None:
+        initialize_results_file(txt_file, column_order)
+    if lamb > 0. and not model.save_act:
+        print('setting lamb=0. If you want to set lamb > 0, set =True')
+
+    old_save_act = model.save_act
+    if lamb == 0.:
+        model.save_act = False
+
+    # pbar = tqdm(range(steps), desc='description', ncols=100)
+
+    if loss_fn == None:
+        loss_fn = lambda x, y: torch.mean((x - y) ** 2)
+    else:
+        loss_fn = loss_fn
+
+    if opt == "Adam":
+        optimizer = torch.optim.Adam(model.parameters(), lr=lr)
+    elif opt == "LBFGS":
+        optimizer = LBFGS(model.parameters(), lr=lr, history_size=10, line_search_fn="strong_wolfe",
+                          tolerance_grad=1e-32, tolerance_change=1e-32, tolerance_ys=1e-32)
+    else:
+        optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9, nesterov=True)
+
+    lr_scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=5, min_lr=1e-9)
+
+    results = {'train_losses': [], 'val_losses': [], 'accuracies': [], 'precisions': [], 'recalls': [], 'f1-scores': [],
+               'lrs': [], 'regularize': []}
+
+    steps = math.ceil(dataset['train_input'].shape[0] / batch_size)
+
+    train_loss = torch.zeros(1).to(model.device)
+    reg_ = torch.zeros(1).to(model.device)
+
+    def closure():
+        nonlocal train_loss, reg_
+        optimizer.zero_grad()
+        pred = model.forward(batch_train_input)
+        loss = loss_fn(pred, batch_train_label)
+        if model.save_act:
+            if reg_metric == 'edge_backward':
+                model.attribute()
+            if reg_metric == 'node_backward':
+                model.node_attribute()
+            reg_ = model.get_reg(reg_metric, lamb_l1, lamb_entropy)
+        else:
+            reg_ = torch.tensor(0.)
+        objective = loss + lamb * reg_
+        train_loss = (train_loss * batch_num + objective.detach()) / (batch_num + 1)
+        objective.backward()
+        return objective
+
+    for epoch in range(epochs):
+
+        if epoch == steps - 1 and old_save_act:
+            model.save_act = True
+
+        train_indices = np.arange(dataset['train_input'].shape[0])
+        np.random.shuffle(train_indices)
+        train_pbar = tqdm(range(steps), desc=f'Epoch {epoch + 1}/{epochs} Training', file=sys.stdout)
+
+        for batch_num in train_pbar:
+            step = epoch * steps + batch_num + 1
+            i = batch_num * batch_size
+            batch_train_id = train_indices[i:i + batch_size]
+            batch_train_input = dataset['train_input'][batch_train_id].to(model.device)
+            batch_train_label = dataset['train_label'][batch_train_id].to(model.device)
+
+            if opt == "LBFGS":
+                optimizer.step(closure)
+
+            if opt == "Adam":
+                optimizer.zero_grad()
+                pred = model.forward(batch_train_input)
+                train_loss = loss_fn(pred, batch_train_input)
+                if model.save_act:
+                    reg_ = model.get_reg(reg_metric, lamb_l1, lamb_entropy)
+                else:
+                    reg_ = torch.tensor(0.)
+                loss = train_loss + lamb * reg_
+                train_loss = (train_loss * batch_num + loss.detach()) / (batch_num + 1)
+
+                loss.backward()
+                optimizer.step()
+
+            train_pbar.set_postfix(loss=train_loss.item())
+
+        val_loss = torch.zeros(1).to(model.device)
+        with torch.no_grad():
+            all_predictions = []
+            all_labels = []
+            test_indices = np.arange(dataset['test_input'].shape[0])
+            np.random.shuffle(test_indices)
+            test_steps = math.ceil(dataset['test_input'].shape[0] / batch_size)
+            test_pbar = tqdm(range(test_steps), desc=f'Epoch {epoch + 1}/{epochs} Validation', file=sys.stdout)
+            for batch_num in test_pbar:
+                i = batch_num * batch_size
+                batch_test_id = test_indices[i:i + batch_size]
+                batch_test_input = dataset['test_input'][batch_test_id].to(model.device)
+                batch_test_label = dataset['test_label'][batch_test_id].to(model.device)
+                outputs = model.forward(batch_test_input)
+                loss = loss_fn(outputs, batch_test_label)
+                val_loss = (val_loss * batch_num + loss.detach()) / (batch_num + 1)
+                test_pbar.set_postfix(val_loss=val_loss.item())
+                if label is not None:
+                    diffs = torch.abs(outputs - label)
+                    closest_indices = torch.argmin(diffs, dim=1)
+                    closest_values = label[closest_indices]
+                    all_predictions.extend(closest_values.detach().cpu().numpy())
+                    all_labels.extend(batch_test_label.detach().cpu().numpy())
+
+            train_lr = lr_scheduler.get_last_lr()[0]
+            lr_scheduler.step(val_loss)
+
+        if label is not None:
+            m = calculate_metric(all_labels, all_predictions, class_dict)
+            print(m)
+            results["accuracies"].append(m["accuracy"])
+            results["precisions"].append(m["precision"])
+            results["recalls"].append(m["recall"])
+            results["f1-scores"].append(m["f1-score"])
+            results["lrs"].append(train_lr)
+            if best < m["f1-score"]:
+                best = m["f1-score"]
+                results['all_predictions'] = all_predictions
+                results['all_labels'] = all_labels
+
+                # if save_path is not None:
+                #     plot_confusion_matrix(all_labels, all_predictions, class_dict, save_path)
+                # if save_path is not None:
+                #     model.saveckpt(path=(os.path.join(save_path, 'save_model') + '/' + 'model'))
+            if txt_file is not None:
+                m.update({'lr': train_lr, 'epoch': epoch, 'train_loss': train_loss.item(), 'val_loss': val_loss.item()})
+                append_to_results_file(txt_file, m, column_order,
+                                       custom_column_widths=custom_column_widths)
+
+        results["train_losses"].append(train_loss.item())
+        results["val_losses"].append(val_loss.item())
+        results["regularize"].append(reg_.item())
+
+    return results

yms_kan/utils.py

@@ -75,9 +75,9 @@ SYMBOLIC_LIB = {'x': (lambda x: x, lambda x: x, 1, lambda x, y_th: ((), x)),
                 '0': (lambda x: x * 0, lambda x: x * 0, 0, lambda x, y_th: ((), x * 0)),
                 'gaussian': (lambda x: torch.exp(-x ** 2), lambda x: sympy.exp(-x ** 2), 3,
                              lambda x, y_th: ((), torch.exp(-x ** 2))),
-                #'cosh': (lambda x: torch.cosh(x), lambda x: sympy.cosh(x), 5),
-                #'sigmoid': (lambda x: torch.sigmoid(x), sympy.Function('sigmoid'), 4),
-                #'relu': (lambda x: torch.relu(x), relu),
+                # 'cosh': (lambda x: torch.cosh(x), lambda x: sympy.cosh(x), 5),
+                # 'sigmoid': (lambda x: torch.sigmoid(x), sympy.Function('sigmoid'), 4),
+                # 'relu': (lambda x: torch.relu(x), relu),
                 }
 
 
@@ -465,6 +465,78 @@ def batch_hessian(model, x, create_graph=False):
     return torch.autograd.functional.jacobian(_jac_sum, x, create_graph=create_graph).permute(1, 0, 2)
 
 
+# def create_from_data(inputs, labels, ratios=[0.8, 0.2], device='cpu'):
+#     # 参数校验
+#     if not (2 <= len(ratios) <= 3):
+#         raise ValueError("ratios必须是长度为2或3的列表")
+#     if not np.isclose(sum(ratios), 1.0):
+#         raise ValueError("ratios元素之和必须为1")
+#     for r in ratios:
+#         if r < 0 or r > 1:
+#             raise ValueError("ratios元素必须在0-1之间")
+#
+#     from collections import defaultdict
+#     class_indices = defaultdict(list)
+#     for idx, label in enumerate(labels):
+#         class_indices[label.item()].append(idx)
+#
+#     # 初始化各数据集索引
+#     split_indices = defaultdict(list)
+#
+#     # 分层抽样
+#     for class_label, indices in class_indices.items():
+#         if not indices:
+#             continue
+#
+#         num_samples = len(indices)
+#         np.random.shuffle(indices)  # 先打乱类别内样本
+#
+#         # 计算各子集分割点
+#         split_points = np.cumsum(ratios).tolist()[:-1]  # 计算分割点（不含最后一个1）
+#         splits = np.split(indices, [int(num_samples * p) for p in split_points])
+#
+#         # 分配到对应数据集
+#         for i, subset in enumerate(splits):
+#             if i == 0:
+#                 split_indices['train'].extend(subset)
+#             elif i == 1:
+#                 split_indices['test'].extend(subset)
+#             elif i == 2:
+#                 split_indices['val'].extend(subset)
+#
+#     # 处理数据集合并（验证集可选）
+#     train_val = {}
+#     if 'test' in split_indices:
+#         # 合并训练+验证并打乱
+#         train_val_idx = np.concatenate([
+#             np.array(split_indices['train']),
+#             np.array(split_indices['val'])
+#         ])
+#         np.random.shuffle(train_val_idx)
+#         train_val = {
+#             'train_input': inputs[train_val_idx[:len(split_indices['train'])]].detach().to(device),
+#             'train_label': labels[train_val_idx[:len(split_indices['train'])]].detach().to(device),
+#             'test_input': inputs[train_val_idx[len(split_indices['train']):]].detach().to(device),
+#             'test_label': labels[train_val_idx[len(split_indices['train']):]].detach().to(device)
+#         }
+#     else:
+#         # 只有训练集
+#         train_idx = np.array(split_indices['train'])
+#         np.random.shuffle(train_idx)
+#         train_val = {
+#             'train_input': inputs[train_idx].detach().to(device),
+#             'train_label': labels[train_idx].detach().to(device)
+#         }
+#
+#     # 处理测试集
+#     test_idx = np.array(split_indices.get('val', []))
+#     np.random.shuffle(test_idx)
+#     test_set = {
+#         'val_input': inputs[test_idx].detach().to(device) if test_idx.size else None,
+#         'val_label': labels[test_idx].detach().to(device) if test_idx.size else None
+#     }
+#
+#     return train_val, test_set
 def create_from_data(inputs, labels, train_ratio=0.8, device='cpu'):
     from collections import defaultdict
     class_indices = defaultdict(list)

yms_kan/version.py

@@ -1 +1 @@
-__version__ = "0.0.9"  # 初始版本
+__version__ = "0.0.10"  # 初始版本

{yms_kan-0.0.9.dist-info → yms_kan-0.0.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: yms_kan
-Version: 0.0.9
+Version: 0.0.10
 Summary: My awesome package
 Author-email: yms <11@qq.com>
 License-Expression: MIT

{yms_kan-0.0.9.dist-info → yms_kan-0.0.10.dist-info}/RECORD

@@ -1,7 +1,7 @@
 yms_kan/KANLayer.py,sha256=-V2Fh5wvPYvfF1tmQVxJKWvvaAHiwo2EiFpd8VDgB1c,14149
 yms_kan/LBFGS.py,sha256=OPeRPDp40jaVH4qPoBDMEub7TPhyvw7pbqwQar3OZ1A,17620
-yms_kan/MLP.py,sha256=ryLzSuBrsGlSHRLwnQZCNj-Ru9BwXJYoHNkAwX14N64,12804
-yms_kan/MultKAN.py,sha256=n58W6tORBDuh_-pUVp-ER-R9KNJKdYjFaisDx8wJpWw,122113
+yms_kan/MLP.py,sha256=JFLogd1EPFVCrBJJtvNMNu68ejdUJ2O6qYc7l4pfFFI,12728
+yms_kan/MultKAN.py,sha256=eFh5jCGRPUzrpeWkvonGGzNWK4w0aODN8Q6M9wZ5IaY,122193
 yms_kan/Symbolic_KANLayer.py,sha256=WhJzC5IMIpXI_K7aYamOrWTK7uckxVdsM9N4oLZMO3I,9897
 yms_kan/__init__.py,sha256=O2c6DIG4PHavXF2v7K9jNqMbJXWr4-gTN3Vs1YSlc64,120
 yms_kan/compiler.py,sha256=7bVwDNX0xmLAjQ8V1FdmkIIIibmy_W5eaeSKBlYL0Vc,18632
@@ -11,13 +11,13 @@ yms_kan/hypothesis.py,sha256=Ec20xadfgOSSWeZHQaGn-h9F2PY7LWFU3iniNI2Zd_4,23165
 yms_kan/plotting.py,sha256=Moi6QTJQxHjutGMgxR9oSsqZSzYY3TP-7WNapdCIqzw,18097
 yms_kan/spline.py,sha256=ZXyGwl2Sc-UrnrcuUXeUQkBOMnetaWcHrbpZaqatCvs,4345
 yms_kan/tool.py,sha256=rkRpqF3EcsAq7a3k1F1zKlxfJ4U9n-FzHyNCJgN4URY,21159
-yms_kan/train_eval_utils.py,sha256=ZY_2GbSjVNAeaGQ24tq1NdjHKrAKKhpTGyamrL98Ap4,16713
-yms_kan/utils.py,sha256=J07L-tgmc1OfU6Tl6mGwHJRizjFN75EJK8BxejaZLUc,23860
-yms_kan/version.py,sha256=BAglq1pSrHfsSqSyAI9RhpMrPmvDRjt0yW5FHNU8gT0,39
+yms_kan/train_eval_utils.py,sha256=3WPtCKLcrotU92s4S0uuIa1rXOAHxyAfDwFJGwUxvy0,16210
+yms_kan/utils.py,sha256=k1fZvv9P6vBBV7LMysoTL2j-bglkBWO0l31dNkWI_Jo,26763
+yms_kan/version.py,sha256=ts9Xi3n2P07g5eVEUSK46avv_nYOOTJ_EeHV2X6IfhM,40
 yms_kan/assets/img/mult_symbol.png,sha256=2f4xUKdweft-qUbHjFI5h9-smnEtc0FWq8hNYZhPAXY,6392
 yms_kan/assets/img/sum_symbol.png,sha256=94QkMUzmEjlCq_yf14nMEQmettaq86FmlGfdl22b4XE,6210
-yms_kan-0.0.9.dist-info/licenses/LICENSE,sha256=BJXDWyF4Groqtnp4Gi9puH4aLg7A2IC3MpHmC-cSxwc,1067
-yms_kan-0.0.9.dist-info/METADATA,sha256=5CcbAeKN87Gh4nrcbkOndaEq734S7Ac4Y4ewkK5D_LU,240
-yms_kan-0.0.9.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-yms_kan-0.0.9.dist-info/top_level.txt,sha256=Z_JDh6yZf-EiW1eKgL6ADsN2yqEMRMspi-o29JZ1WPo,8
-yms_kan-0.0.9.dist-info/RECORD,,
+yms_kan-0.0.10.dist-info/licenses/LICENSE,sha256=BJXDWyF4Groqtnp4Gi9puH4aLg7A2IC3MpHmC-cSxwc,1067
+yms_kan-0.0.10.dist-info/METADATA,sha256=pjZJ0E5OIBv93IoXD10lXpQCPN9ixDIRDuNUwov6Fls,241
+yms_kan-0.0.10.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
+yms_kan-0.0.10.dist-info/top_level.txt,sha256=Z_JDh6yZf-EiW1eKgL6ADsN2yqEMRMspi-o29JZ1WPo,8
+yms_kan-0.0.10.dist-info/RECORD,,

{yms_kan-0.0.9.dist-info → yms_kan-0.0.10.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.1.0)
+Generator: setuptools (79.0.0)
 Root-Is-Purelib: true
 Tag: py3-none-any