refactor: 修改参数传递方案

khaosz/config/train_config.py

@@ -22,15 +22,14 @@ class TrainConfig:
         default=None,
         metadata={"help": "Dataset for training."}
     )
-    optimizer: Optimizer = field(
+    optimizer_fn: Callable[[nn.Module], Optimizer] = field(
         default=None,
-        metadata={"help": "Optimizer for training."}
+        metadata={"help": "Optimizer factory for training."}
     )
-    scheduler: LRScheduler = field(
+    scheduler_fn: Callable[[Optimizer], LRScheduler] = field(
         default=None,
-        metadata={"help": "Scheduler for training."}
+        metadata={"help": "Scheduler factory for training."}
     )
-    
     n_epoch: int = field(
         default=1,
         metadata={"help": "Number of epochs for training."}
@@ -105,19 +104,10 @@ class TrainConfig:
         default=None,
         metadata={"help": "Parallel function for training."}
     )
-    state_dict_wrapper: Optional[Callable] = field(
+    state_dict_fn: Optional[Callable] = field(
          default=None,
          metadata={"help": "Parallel function for state  dict saving."}
     )
-    
-    optimizer_factory: Optional[Callable[[nn.Module], Optimizer]] = field(
-        default=None,
-        metadata={"help": "Optimizer factory for training."}
-    )
-    scheduler_factory: Optional[Callable[[Optimizer], LRScheduler]] = field(
-        default=None,
-        metadata={"help": "Scheduler factory for training."}
-    )
 
     # others
     device_ids: Optional[List[int]] = field(
@@ -137,19 +127,10 @@ class TrainConfig:
         self.validate()
     
     def validate(self):
-        required_fields = ["model", "strategy", "dataset"]
+        required_fields = ["model", "strategy", "dataset", "optimizer_fn", "scheduler_fn"]
         
         for field_name in required_fields:
             if getattr(self, field_name) is None:
                 raise ValueError(f"{field_name} is required.")
         
-        factory_case = all([self.optimizer_factory, self.scheduler_factory])
-        argument_case = all([self.optimizer, self.scheduler])
-        self.nprocs = max(self.nprocs, 1)
-        
-        if self.nprocs > 1 and not factory_case:
-            raise ValueError("Distributed training requires optimizer and scheduler factories.")
-        elif self.nprocs == 1 and not argument_case:
-            raise ValueError("Single process training requires optimizer and scheduler arguments.")
-