refactor: 优化参数传递，清理导入样式

astrai/trainer/__init__.py

@@ -1,15 +1,14 @@
-from astrai.trainer.trainer import Trainer
-from astrai.trainer.strategy import StrategyFactory, BaseStrategy
-from astrai.trainer.schedule import SchedulerFactory, BaseScheduler
-
+from astrai.trainer.schedule import BaseScheduler, SchedulerFactory
+from astrai.trainer.strategy import BaseStrategy, StrategyFactory
 from astrai.trainer.train_callback import (
-    TrainCallback,
-    GradientClippingCallback,
-    SchedulerCallback,
     CheckpointCallback,
-    ProgressBarCallback,
+    GradientClippingCallback,
     MetricLoggerCallback,
+    ProgressBarCallback,
+    SchedulerCallback,
+    TrainCallback,
 )
+from astrai.trainer.trainer import Trainer
 
 __all__ = [
     # Main trainer

astrai/trainer/metric_util.py

@@ -1,6 +1,7 @@
-import torch.nn as nn
 from typing import Dict
 
+import torch.nn as nn
+
 
 def grad_norm(model: nn.Module, norm_type: int = 2) -> Dict[str, float]:
     """Compute gradient norm for each parameter in the model."""

astrai/trainer/schedule.py

@@ -1,10 +1,10 @@
 """Learning rate scheduler implementations with factory pattern."""
 
 import math
-from abc import abstractmethod, ABC
+from abc import ABC, abstractmethod
 from typing import Any, Dict, List, Type
+
 from torch.optim.lr_scheduler import LRScheduler
-from astrai.config.schedule_config import ScheduleConfig
 
 
 class BaseScheduler(LRScheduler, ABC):
@@ -37,10 +37,6 @@ class SchedulerFactory:
             ...
 
         scheduler = SchedulerFactory.create(optimizer, "custom", **kwargs)
-
-        # Or from config
-        config = CosineScheduleConfig(total_steps=10000)
-        scheduler = SchedulerFactory.load(optimizer, config)
     """
 
     SCHEDULER_MAP: Dict[str, Type[BaseScheduler]] = {}
@@ -67,7 +63,7 @@ class SchedulerFactory:
         return decorator
 
     @classmethod
-    def create(cls, optimizer, schedule_type: str, **kwargs) -> BaseScheduler:
+    def create(cls, optimizer, schedule_type: str = "none", **kwargs) -> BaseScheduler:
         """Create a scheduler instance by type name.
 
         Args:
@@ -90,29 +86,13 @@ class SchedulerFactory:
         scheduler_cls = cls.SCHEDULER_MAP[schedule_type]
         return scheduler_cls(optimizer, **kwargs)
 
-    @staticmethod
-    def load(optimizer, schedule_config: ScheduleConfig) -> BaseScheduler:
-        """Create a scheduler from a ScheduleConfig object.
-
-        Args:
-            optimizer: PyTorch optimizer
-            schedule_config: ScheduleConfig instance
-
-        Returns:
-            Scheduler instance
-        """
-        kwargs = schedule_config.get_kwargs()
-        schedule_type = kwargs.pop("schedule_type")
-        return SchedulerFactory.create(optimizer, schedule_type, **kwargs)
-
     @classmethod
     def available_types(cls) -> list:
         """Return list of registered scheduler type names."""
         return list(cls.SCHEDULER_MAP.keys())
 
 
-# ============== Scheduler Classes ==============
-# All scheduler classes are registered at class definition time using the decorator
+# ----------- Scheduler implementations -----------
 
 
 @SchedulerFactory.register("cosine")

astrai/trainer/strategy.py

@@ -1,14 +1,14 @@
 """Training strategy implementations with factory pattern."""
 
 import copy
+from abc import ABC, abstractmethod
+from typing import Any, Callable, Dict, Union
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch.nn.parallel import DistributedDataParallel as DDP
-
 from torch import Tensor
-from typing import Any, Callable, Dict, Union
-from abc import ABC, abstractmethod
+from torch.nn.parallel import DistributedDataParallel as DDP
 
 
 def unwrap_model(model: nn.Module) -> nn.Module:

astrai/trainer/train_callback.py

@@ -1,25 +1,25 @@
-import os
 import json
+import os
 import time
-import torch.nn as nn
-
 from pathlib import Path
-from tqdm import tqdm
-from torch.nn.utils import clip_grad_norm_
 from typing import Callable, List, Optional, Protocol
 
+import torch.nn as nn
+from torch.nn.utils import clip_grad_norm_
+from tqdm import tqdm
+
+from astrai.data.serialization import Checkpoint
 from astrai.parallel import only_on_rank
 from astrai.trainer.metric_util import (
+    ctx_get_grad_max,
+    ctx_get_grad_mean,
+    ctx_get_grad_min,
+    ctx_get_grad_nan_num,
+    ctx_get_grad_norm,
+    ctx_get_grad_std,
     ctx_get_loss,
     ctx_get_lr,
-    ctx_get_grad_max,
-    ctx_get_grad_min,
-    ctx_get_grad_norm,
-    ctx_get_grad_mean,
-    ctx_get_grad_std,
-    ctx_get_grad_nan_num,
 )
-from astrai.data.serialization import Checkpoint
 from astrai.trainer.train_context import TrainContext
 
 

astrai/trainer/train_context.py

@@ -1,16 +1,16 @@
+from dataclasses import dataclass, field
+from typing import Optional, Self
+
 import torch.nn as nn
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 from torch.utils.data import DataLoader
 
+from astrai.config.train_config import TrainConfig
 from astrai.data import ResumableDistributedSampler
 from astrai.data.serialization import Checkpoint
-from astrai.trainer.strategy import StrategyFactory, BaseStrategy
-from astrai.config.train_config import TrainConfig
-from astrai.parallel.setup import get_current_device, get_world_size, get_rank
-
-from dataclasses import dataclass, field
-from typing import Optional, Self
+from astrai.parallel.setup import get_current_device, get_rank, get_world_size
+from astrai.trainer.strategy import BaseStrategy, StrategyFactory
 
 
 @dataclass

astrai/trainer/trainer.py

@@ -1,17 +1,18 @@
 import logging
-from typing import Optional, List
+from typing import List, Optional
+
 from astrai.config import TrainConfig
-from astrai.trainer.train_callback import (
-    TrainCallback,
-    ProgressBarCallback,
-    CheckpointCallback,
-    MetricLoggerCallback,
-    GradientClippingCallback,
-    SchedulerCallback,
-)
-from astrai.trainer.train_context import TrainContext, TrainContextBuilder
 from astrai.data.serialization import Checkpoint
 from astrai.parallel.setup import spawn_parallel_fn
+from astrai.trainer.train_callback import (
+    CheckpointCallback,
+    GradientClippingCallback,
+    MetricLoggerCallback,
+    ProgressBarCallback,
+    SchedulerCallback,
+    TrainCallback,
+)
+from astrai.trainer.train_context import TrainContext, TrainContextBuilder
 
 logger = logging.getLogger(__name__)