Update modeling_super_linear.py

modeling_super_linear.py

@@ -367,18 +367,18 @@ class superLinear(nn.Module):
         else:
             self.freq_experts = configs.freq_experts.split('_')
 
-        #print("self.freq_experts:", self.freq_experts)
+        
 
         self.moe_loss = None
         self.top_k_experts = configs.top_k_experts
-       # self.noisy_gating = configs.noisy_gating
         self.n_experts = configs.moe_n_experts
         self.freeze_experts = configs.freeze_experts
         self.layer_type = configs.layer_type
         self.model_name = "SuperLinear"
 
-        #print("self.layer_type", self.layer_type)
+
         self.layer_dict = {'DLinear': DLinear, 'Linear': Linear, 'NLinear': NLinear, 'RLinear': RLinear}
+        
         path = configs.linear_checkpoints_path + configs.linear_checkpoints_dir 
         dirs = os.listdir(path)
         checkpoints_paths = [path + "/" + d + "/" + "checkpoint.pth" for d in dirs]
@@ -390,31 +390,31 @@ class superLinear(nn.Module):
                     cycle = cp.split("/")
 
         self.experts = {}
-        if self.freq_experts is not None:
-            for expert_freq in self.freq_experts:
-                if expert_freq == "naive" or expert_freq == "Naive":
-                    self.experts[expert_freq] = Naive(self.seq_len, self.pred_len)
-                elif expert_freq == "mean" or expert_freq == "Mean":    
-                    self.experts[expert_freq] = Mean(self.seq_len, self.pred_len)
-                else:
-                    self.experts[expert_freq] = self.layer_dict[self.layer_type](self.seq_len, self.pred_len)
-                    if configs.load_linear:
-                        cycle = self.map_to_cycle(expert_freq)
-                        cycle_str = f'cycle_{cycle}/'
-                        cycle_checkpoint_path = [cp for cp in checkpoints_paths if (cycle_str in cp and self.layer_type in cp)]
-                        if len(cycle_checkpoint_path) > 0:
-                            print()
-                            print(cycle_str)
-                            cycle_checkpoint_path = cycle_checkpoint_path[0]
-                            #print(f'loading checkpoint with layer type: {self.layer_type} and cycle: {cycle_str}')
-                            print(cycle_checkpoint_path)
-                            self.experts[expert_freq].load_state_dict(torch.load(cycle_checkpoint_path))
-                        else:
-                            print(f"Checkpoint for {cycle_str} not found in {path}")
-                            raise ValueError(f"Checkpoint for {cycle_str} not found in {path}")
-                        if configs.freeze_experts:
-                            for param in self.experts[expert_freq].parameters():
-                                param.requires_grad = False
+        # if self.freq_experts is not None:
+        #     for expert_freq in self.freq_experts:
+        #         if expert_freq == "naive" or expert_freq == "Naive":
+        #             self.experts[expert_freq] = Naive(self.seq_len, self.pred_len)
+        #         elif expert_freq == "mean" or expert_freq == "Mean":    
+        #             self.experts[expert_freq] = Mean(self.seq_len, self.pred_len)
+        #         else:
+        #             self.experts[expert_freq] = self.layer_dict[self.layer_type](self.seq_len, self.pred_len)
+        #             if configs.load_linear:
+        #                 cycle = self.map_to_cycle(expert_freq)
+        #                 cycle_str = f'cycle_{cycle}/'
+        #                 cycle_checkpoint_path = [cp for cp in checkpoints_paths if (cycle_str in cp and self.layer_type in cp)]
+        #                 if len(cycle_checkpoint_path) > 0:
+        #                     print()
+        #                     print(cycle_str)
+        #                     cycle_checkpoint_path = cycle_checkpoint_path[0]
+        #                     #print(f'loading checkpoint with layer type: {self.layer_type} and cycle: {cycle_str}')
+        #                     print(cycle_checkpoint_path)
+        #                     self.experts[expert_freq].load_state_dict(torch.load(cycle_checkpoint_path))
+        #                 else:
+        #                     print(f"Checkpoint for {cycle_str} not found in {path}")
+        #                     raise ValueError(f"Checkpoint for {cycle_str} not found in {path}")
+        #                 if configs.freeze_experts:
+        #                     for param in self.experts[expert_freq].parameters():
+        #                         param.requires_grad = False
                         
             self.n_experts = len(self.experts)
         else: