Update modeling_super_linear.py

modeling_super_linear.py

@@ -390,37 +390,37 @@ 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:
-        #     for i in range(self.n_experts):
-        #         print(f"creating expert {i}")
-        #         self.experts[str(i)] = self.layer_dict[self.layer_type](self.seq_len, self.pred_len)
+            self.n_experts = len(self.experts)
+        else:
+            for i in range(self.n_experts):
+                print(f"creating expert {i}")
+                self.experts[str(i)] = self.layer_dict[self.layer_type](self.seq_len, self.pred_len)
 
 
         if configs.misc_moe>0: