junshan-kit 2.4.0__py2.py3-none-any.whl → 2.4.1__py2.py3-none-any.whl

junshan_kit/ExperimentHub.py

@@ -1,13 +1,9 @@
 import sys, os, torch, random
 import numpy as np
 import torch.nn as nn
-from torch.utils.data import Subset
-
-# # 获取当前脚本所在目录
-# script_dir = os.path.dirname(os.path.abspath(__file__))
-# sys.path.append(os.path.join(script_dir, 'src'))
-
-from junshan_kit import datahub, Models, TrainingParas
+import torch.utils.data as Data
+from torch.utils.data import Subset, random_split
+from junshan_kit import ComOptimizers, datahub, Models, TrainingParas, SPBM
 
 # -------------------------------------
 def set_seed(seed=42):
@@ -191,4 +187,142 @@ class Train_Steps:
 
         return train_dataset, test_dataset
 
+    # <validation> : Step 5.2   --> step.py
+    def set_val_set(self, data_name, train_dataset, Paras):
+        if Paras["validation"][data_name]:
+            size_ = len(train_dataset)
+            val_size = int(size_ * Paras["validation_rate"][data_name])
+            train_size = size_ - val_size
+
+            train_dataset, val_dataset = random_split(
+                train_dataset,
+                [train_size, val_size],
+                generator=torch.Generator().manual_seed(Paras["seed"]),
+            )
+
+        else:
+            val_dataset = Subset(train_dataset, [])
+
+        return train_dataset, val_dataset
+    # <validation>
+
+
+    # <get_dataloader> Step 5.3 -->step.py
+    def get_dataloader(self, data_name, train_dataset, test_dataset, Paras):
+        set_seed(Paras["seed"])
+        g = torch.Generator()
+        g.manual_seed(Paras["seed"])
+
+        # Create training DataLoader
+        
+        train_loader = Data.DataLoader(
+            dataset=train_dataset,
+            shuffle=True,
+            batch_size=self.args.bs,
+            generator=g,
+            num_workers=4,
+        )
+
+        # test loader
+        test_loader = Data.DataLoader(
+            dataset=test_dataset,
+            shuffle=False,
+            batch_size=self.args.bs,
+            generator=g,
+            num_workers=4,
+        )
+    
+        return train_loader, test_loader
+    # <get_dataloader>
+
+    def hyperparas_and_path(
+        self,
+        model_name,
+        data_name,
+        optimizer_name,
+        Paras,
+    ):
+        params_gird = Paras["optimizer_dict"][optimizer_name]["params"]
+        keys, values = list(params_gird.keys()), list(params_gird.values())
+
+        # Set the path for saving results
+        folder_path = f'./{Paras["results_folder_name"]}/seed_{Paras["seed"]}/{model_name}/{data_name}/{optimizer_name}/train_{Paras["train_data_num"]}_test_{Paras["test_data_num"]}/Batch_size_{self.args.bs}/epoch_{self.args.e}/{Paras["time_str"]}'
+        os.makedirs(folder_path, exist_ok=True)
+
+        return keys, values, folder_path
+
+    
+    # <Reloading> Step 7.3 --> step.py
+    def reloading_model_dataloader(self,
+        base_model_fn,
+        initial_state_dict,
+        data_name,
+        train_dataset,
+        test_dataset,
+        Paras,
+    ):
+        set_seed(Paras["seed"])
+        model = base_model_fn()
+        model.load_state_dict(initial_state_dict)
+        model.to(Paras["device"])
+        train_loader, test_loader = self.get_dataloader(
+            data_name, train_dataset, test_dataset, Paras
+        )
+
+        return model, train_loader, test_loader
+    # <Reloading>
+
+    def chosen_optimizer(self, optimizer_name, model, hyperparams, Paras):
+        if optimizer_name == "SGD":
+            optimizer = torch.optim.SGD(model.parameters(), lr=hyperparams["alpha"])
+
+        elif optimizer_name == "ADAM":
+            optimizer = torch.optim.Adam(
+                model.parameters(),
+                lr=hyperparams["alpha"],
+                betas=(hyperparams["beta1"], hyperparams["beta2"]),
+                eps=hyperparams["epsilon"],
+            )
+
+        elif optimizer_name in ["SPBM-TR"]:
+            optimizer = SPBM.TR(model.parameters(), model, hyperparams, Paras)
+
+        elif optimizer_name in ["SPBM-TR-NoneLower"]:
+            optimizer = SPBM.TR_NoneLower(model.parameters(), model, hyperparams, Paras)
+        
+        elif optimizer_name in ["SPBM-TR-NoneSpecial"]:
+            optimizer = SPBM.TR_NoneSpecial(model.parameters(), model, hyperparams, Paras)
+        
+        elif optimizer_name in ["SPBM-TR-NoneCut"]:
+            optimizer = SPBM.TR_NoneCut(model.parameters(), model, hyperparams, Paras)
+
+        elif optimizer_name in ["SPBM-PF-NoneLower"]:
+            optimizer = SPBM.PF_NoneLower(model.parameters(), model, hyperparams, Paras)
+
+        elif optimizer_name in ["SPBM-PF"]:
+            optimizer = SPBM.PF(model.parameters(), model, hyperparams, Paras)
+        
+        elif optimizer_name in ["SPBM-PF-NoneCut"]:
+            optimizer = SPBM.PF_NoneCut(model.parameters(), model, hyperparams, Paras)
+
+        elif optimizer_name in ["SPSmax"]:
+            optimizer = ComOptimizers.SPSmax(
+                model.parameters(), model, hyperparams, Paras
+            )
 
+        elif optimizer_name in ["ALR-SMAG"]:
+            optimizer = ComOptimizers.ALR_SMAG(
+                model.parameters(), model, hyperparams, Paras
+            )
+
+        elif optimizer_name in ["Bundle"]:
+            optimizer = ComOptimizers.Bundle(
+                model.parameters(), model, hyperparams, Paras
+            )
+
+        else:
+            raise NotImplementedError(f"{optimizer_name} is not supported.")
+
+        return optimizer
+
+    

junshan_kit/Print_Info.py

@@ -9,17 +9,17 @@ def training_group(training_group):
     print(f"-------------------------------------------------------")
 
 
-def training_info(use_color, data_name, optimizer_name, folder_path, hyperparams, Paras, model_name):
+def training_info(args, use_color, data_name, optimizer_name, folder_path, hyperparams, Paras, model_name):
     if use_color:
         print("\033[90m" + "-" * 115 + "\033[0m")
         print(
-            f"\033[32m✅ \033[34mDataset:\033[32m {data_name}, \t\033[34mBatch-size:\033[32m {Paras['batch_size'][data_name]}, \t\033[34m(training, val, test) = \033[32m ({Paras['training_samples']}, {Paras['val_samples']}, {Paras['test_samples']}), \t\033[34m device:\033[32m {Paras['device']}"
+            f"\033[32m✅ \033[34mDataset:\033[32m {data_name}, \t\033[34mBatch-size:\033[32m {args.bs}, \t\033[34m(training, test) = \033[32m ({Paras['train_data_num']}, {Paras['test_data_num']}), \t\033[34m device:\033[32m {Paras['device']}"
         )
         print(
             f"\033[32m✅ \033[34mOptimizer:\033[32m {optimizer_name}, \t\033[34mParams:\033[32m {hyperparams}"
         )
         print(
-            f'\033[32m✅ \033[34mmodel:\033[32m {model_name}, \t\033[34mmodel type:\033[32m {Paras["model_type"][model_name]},\t\033[34m loss_fn:\033[32m {Paras["loss_fn"]},\t\033[34m Subset:\033[32m {Paras["select_subset"][data_name]}'
+            f'\033[32m✅ \033[34mmodel:\033[32m {model_name}, \t\033[34mmodel type:\033[32m {Paras["model_type"][model_name]},\t\033[34m loss_fn:\033[32m {Paras["loss_fn"]}'
         )
         print(f"\033[32m✅ \033[34mfolder_path:\033[32m {folder_path}")
         print("\033[90m" + "-" * 115 + "\033[0m")

junshan_kit/SPBM.py

@@ -0,0 +1,350 @@
+from junshan_kit import SPBM, SPBM_func
+import torch, time, os
+from torch.optim.optimizer import Optimizer
+from torch.nn.utils import parameters_to_vector, vector_to_parameters
+
+
+class PF(Optimizer):
+    def __init__(self, params, model, hyperparams, Paras):
+        defaults = dict()
+        super().__init__(params, defaults)
+        self.model = model
+        self.cutting_num = hyperparams['cutting_number']
+        self.M = hyperparams['M']
+        self.delta = hyperparams['delta']
+        self.Paras = Paras
+
+        self.x_his, self.g_his, self.f_his = [], [], []
+
+    def step(self, closure=None):
+        if closure is None:
+            raise RuntimeError("Closure required for CuttingPlaneOptimizer")
+        
+        # 清零梯度并前向计算
+        loss = closure()
+        
+        with torch.no_grad():
+            xk = parameters_to_vector(self.model.parameters())
+            # print(torch.norm(xk))
+            g_k = parameters_to_vector([p.grad if p.grad is not None else torch.zeros_like(p) for p in self.model.parameters()])
+
+            # Add cutting plane
+            x_his, f_his, g_his = SPBM_func.add_cutting(self.x_his, self.f_his, self.g_his,xk.detach().clone(), g_k.detach().clone(), loss.detach().clone(), self.cutting_num)
+
+            ## Cut selection
+            selected_x, selected_f, selected_g = SPBM_func.cut_selection(x_his, f_his, g_his, self.M)
+                    
+            # the coefficient of dual problem
+            Gk, rk, ek = SPBM_func.get_var(selected_x, selected_f, selected_g, self.delta)
+            
+            # SOVER (dual)
+            xk = SPBM_func.subproblem_pf(Gk, ek, xk, self.delta, self.Paras)
+            
+            # print(len(self.f_his))
+            vector_to_parameters(xk, self.model.parameters())
+
+
+        # 暂时返回 loss（tensor 类型）
+        return loss
+    
+# <SPBM-TR>
+class TR(Optimizer):
+    def __init__(self, params, model, hyperparams, Paras):
+        defaults = dict()
+        super().__init__(params, defaults)
+        self.model = model
+        self.cutting_num = hyperparams['cutting_number']
+        self.M = hyperparams['M']
+        self.delta = hyperparams['delta']
+        self.Paras = Paras
+
+        self.x_his, self.g_his, self.f_his = [], [], []
+
+    def step(self, closure=None):
+        if closure is None:
+            raise RuntimeError("Closure required for CuttingPlaneOptimizer")
+        
+        # Reset the gradient and perform forward computation
+        loss = closure()
+        
+        with torch.no_grad():
+            xk = parameters_to_vector(self.model.parameters())
+            # print(torch.norm(xk))
+            g_k = parameters_to_vector([p.grad if p.grad is not None else torch.zeros_like(p) for p in self.model.parameters()])
+
+            # Add cutting plane
+            x_his, f_his, g_his = SPBM_func.add_cutting(self.x_his, self.f_his, self.g_his,xk.detach().clone(), g_k.detach().clone(), loss.detach().clone(), self.cutting_num)
+
+            ## Cut selection
+            selected_x, selected_f, selected_g = SPBM_func.cut_selection(x_his, f_his, g_his, self.M)
+                    
+            # the coefficient of dual problem
+            Gk, rk, ek = SPBM_func.get_var(selected_x, selected_f, selected_g, self.delta)
+            
+            # SOVER (dual)
+            xk = SPBM_func.subproblem_tr_2(Gk, ek, xk, rk, self.Paras)
+            
+            # print(len(self.f_his))
+            vector_to_parameters(xk, self.model.parameters())
+
+        # tensor type
+        return loss
+# <SPBM-TR>  
+
+# <SPBM-TR_NoneSpecial> 
+class TR_NoneSpecial(Optimizer):
+    def __init__(self, params, model, hyperparams, Paras):
+        defaults = dict()
+        super().__init__(params, defaults)
+        self.model = model
+        self.cutting_num = hyperparams['cutting_number']
+        self.M = hyperparams['M']
+        self.delta = hyperparams['delta']
+        self.Paras = Paras
+
+        self.x_his, self.g_his, self.f_his = [], [], []
+
+    def step(self, closure=None):
+        if closure is None:
+            raise RuntimeError("Closure required for CuttingPlaneOptimizer")
+        
+        # Reset the gradient and perform forward computation
+        loss = closure()
+        
+        with torch.no_grad():
+            xk = parameters_to_vector(self.model.parameters())
+            # print(torch.norm(xk))
+            g_k = parameters_to_vector([p.grad if p.grad is not None else torch.zeros_like(p) for p in self.model.parameters()])
+
+            # Add cutting plane
+            x_his, f_his, g_his = SPBM_func.add_cutting(self.x_his, self.f_his, self.g_his,xk.detach().clone(), g_k.detach().clone(), loss.detach().clone(), self.cutting_num)
+
+            ## Cut selection
+            selected_x, selected_f, selected_g = SPBM_func.cut_selection(x_his, f_his, g_his, self.M)
+                    
+            # the coefficient of dual problem
+            Gk, rk, ek = SPBM_func.get_var(selected_x, selected_f, selected_g, self.delta)
+            
+            # SOVER (dual)
+            xk = SPBM_func.subproblem_tr_NoneSpecial(Gk, ek, xk, rk, self.Paras)
+            
+            # print(len(self.f_his))
+            vector_to_parameters(xk, self.model.parameters())
+
+        # tensor type
+        return loss
+# <SPBM-TR_NoneSpecial>  
+
+class TR_primal(Optimizer):
+    def __init__(self, params, model, hyperparams, Paras):
+        defaults = dict()
+        super().__init__(params, defaults)
+        self.model = model
+        self.cutting_num = hyperparams['cutting_number']
+        self.M = hyperparams['M']
+        self.delta = hyperparams['delta']
+        self.Paras = Paras
+
+        self.x_his, self.g_his, self.f_his = [], [], []
+
+    def step(self, closure=None):
+        if closure is None:
+            raise RuntimeError("Closure required for CuttingPlaneOptimizer")
+        
+        # Reset the gradient and perform forward computation
+        loss = closure()
+        
+        with torch.no_grad():
+            xk = parameters_to_vector(self.model.parameters())
+            # print(torch.norm(xk))
+            g_k = parameters_to_vector([p.grad if p.grad is not None else torch.zeros_like(p) for p in self.model.parameters()])
+
+            # Add cutting plane
+            x_his, f_his, g_his = SPBM_func.add_cutting(self.x_his, self.f_his, self.g_his,xk.detach().clone(), g_k.detach().clone(), loss.detach().clone(), self.cutting_num)
+
+            ## Cut selection
+            selected_x, selected_f, selected_g = SPBM_func.cut_selection(x_his, f_his, g_his, self.M)
+                    
+            # the coefficient of dual problem
+            Gk, rk, ek = SPBM_func.get_var(selected_x, selected_f, selected_g, self.delta)
+            
+            # SOVER (dual)
+            xk = SPBM_func.subproblem_tr_primal(Gk, ek, xk, rk, self.Paras)
+            
+            # print(len(self.f_his))
+            vector_to_parameters(xk, self.model.parameters())
+
+        # tensor type
+        return loss
+    
+
+class TR_NoneLower(Optimizer):
+    def __init__(self, params, model, hyperparams, Paras):
+        defaults = dict()
+        super().__init__(params, defaults)
+        self.model = model
+        self.cutting_num = hyperparams['cutting_number']
+        self.M = hyperparams['M']
+        self.delta = hyperparams['delta']
+        self.Paras = Paras
+
+        self.x_his, self.g_his, self.f_his = [], [], []
+
+    def step(self, closure=None):
+        if closure is None:
+            raise RuntimeError("Closure required for CuttingPlaneOptimizer")
+        
+        # Reset the gradient and perform forward computation
+        loss = closure()
+        
+        with torch.no_grad():
+            xk = parameters_to_vector(self.model.parameters())
+            # print(torch.norm(xk))
+            g_k = parameters_to_vector([p.grad if p.grad is not None else torch.zeros_like(p) for p in self.model.parameters()])
+
+            # Add cutting plane
+            x_his, f_his, g_his = SPBM_func.add_cutting(self.x_his, self.f_his, self.g_his,xk.detach().clone(), g_k.detach().clone(), loss.detach().clone(), self.cutting_num)
+
+            ## Cut selection
+            selected_x, selected_f, selected_g = SPBM_func.cut_selection(x_his, f_his, g_his, self.M)
+                    
+            # the coefficient of dual problem
+            Gk, rk, ek = SPBM_func.get_var(selected_x, selected_f, selected_g, self.delta)
+            
+            # SOVER (dual)
+            xk = SPBM_func.subproblem_tr_NoneLower(Gk, ek, xk, rk, self.Paras)
+            
+            # print(len(self.f_his))
+            vector_to_parameters(xk, self.model.parameters())
+
+        # tensor type
+        return loss
+
+class TR_NoneCut(Optimizer):
+    def __init__(self, params, model, hyperparams, Paras):
+        defaults = dict()
+        super().__init__(params, defaults)
+        self.model = model
+        self.cutting_num = hyperparams['cutting_number']
+        self.M = hyperparams['M']
+        self.delta = hyperparams['delta']
+        self.Paras = Paras
+
+        self.x_his, self.g_his, self.f_his = [], [], []
+
+    def step(self, closure=None):
+        if closure is None:
+            raise RuntimeError("Closure required for CuttingPlaneOptimizer")
+        
+        # Reset the gradient and perform forward computation
+        loss = closure()
+        
+        with torch.no_grad():
+            xk = parameters_to_vector(self.model.parameters())
+            # print(torch.norm(xk))
+            g_k = parameters_to_vector([p.grad if p.grad is not None else torch.zeros_like(p) for p in self.model.parameters()])
+
+            # Add cutting plane
+            x_his, f_his, g_his = SPBM_func.add_cutting(self.x_his, self.f_his, self.g_his,xk.detach().clone(), g_k.detach().clone(), loss.detach().clone(), self.cutting_num)
+
+            # ## Cut selection
+            # selected_x, selected_f, selected_g = SPBM_func.cut_selection(x_his, f_his, g_his, self.M)
+                    
+            # the coefficient of dual problem
+            Gk, rk, ek = SPBM_func.get_var(x_his, f_his, g_his, self.delta)
+            
+            # SOVER (dual)
+            xk = SPBM_func.subproblem_tr_NoneLower(Gk, ek, xk, rk, self.Paras)
+            
+            # print(len(self.f_his))
+            vector_to_parameters(xk, self.model.parameters())
+
+        # tensor type
+        return loss
+    
+# ************************** SPBM-PF **************************
+class PF_NoneLower(Optimizer):
+    def __init__(self, params, model, hyperparams, Paras):
+        defaults = dict()
+        super().__init__(params, defaults)
+        self.model = model
+        self.cutting_num = hyperparams['cutting_number']
+        self.M = hyperparams['M']
+        self.delta = hyperparams['delta']
+        self.Paras = Paras
+
+        self.x_his, self.g_his, self.f_his = [], [], []
+
+    def step(self, closure=None):
+        if closure is None:
+            raise RuntimeError("Closure required for CuttingPlaneOptimizer")
+        
+        # Reset the gradient and perform forward computation
+        loss = closure()
+        
+        with torch.no_grad():
+            xk = parameters_to_vector(self.model.parameters())
+            # print(torch.norm(xk))
+            g_k = parameters_to_vector([p.grad if p.grad is not None else torch.zeros_like(p) for p in self.model.parameters()])
+
+            # Add cutting plane
+            x_his, f_his, g_his = SPBM_func.add_cutting(self.x_his, self.f_his, self.g_his,xk.detach().clone(), g_k.detach().clone(), loss.detach().clone(), self.cutting_num)
+
+            ## Cut selection
+            selected_x, selected_f, selected_g = SPBM_func.cut_selection(x_his, f_his, g_his, self.M)
+                    
+            # the coefficient of dual problem
+            Gk, rk, ek = SPBM_func.get_var(selected_x, selected_f, selected_g, self.delta)
+            
+            # SOVER (dual)
+            xk = SPBM_func.subproblem_pf_NoneLower(Gk, ek, xk, self.delta, self.Paras)
+            
+            # print(len(self.f_his))
+            vector_to_parameters(xk, self.model.parameters())
+
+        # tensor type
+        return loss
+    
+
+class PF_NoneCut(Optimizer):
+    def __init__(self, params, model, hyperparams, Paras):
+        defaults = dict()
+        super().__init__(params, defaults)
+        self.model = model
+        self.cutting_num = hyperparams['cutting_number']
+        self.M = hyperparams['M']
+        self.delta = hyperparams['delta']
+        self.Paras = Paras
+
+        self.x_his, self.g_his, self.f_his = [], [], []
+
+    def step(self, closure=None):
+        if closure is None:
+            raise RuntimeError("Closure required for CuttingPlaneOptimizer")
+        
+        # Reset the gradient and perform forward computation
+        loss = closure()
+        
+        with torch.no_grad():
+            xk = parameters_to_vector(self.model.parameters())
+            # print(torch.norm(xk))
+            g_k = parameters_to_vector([p.grad if p.grad is not None else torch.zeros_like(p) for p in self.model.parameters()])
+
+            # Add cutting plane
+            x_his, f_his, g_his = SPBM_func.add_cutting(self.x_his, self.f_his, self.g_his,xk.detach().clone(), g_k.detach().clone(), loss.detach().clone(), self.cutting_num)
+
+            # ## Cut selection
+            # selected_x, selected_f, selected_g = SPBM_func.cut_selection(x_his, f_his, g_his, self.M)
+                    
+            # the coefficient of dual problem
+            Gk, rk, ek = SPBM_func.get_var(x_his, f_his, g_his, self.delta)
+            
+            # SOVER (dual)
+            xk = SPBM_func.subproblem_pf_NoneLower(Gk, ek, xk, self.delta, self.Paras)
+            
+            # print(len(self.f_his))
+            vector_to_parameters(xk, self.model.parameters())
+
+        # tensor type
+        return loss
+