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

tests/trainer/test_train_strategy.py

@@ -1,10 +1,9 @@
-import torch
 import numpy as np
-import pytest
+import torch
 
 from astrai.config import *
-from astrai.trainer.schedule import *
 from astrai.data.dataset import *
+from astrai.trainer.schedule import *
 
 
 def test_schedule_factory_random_configs():
@@ -16,41 +15,57 @@ def test_schedule_factory_random_configs():
 
     # Test multiple random configurations
     for _ in range(5):  # Test 5 random configurations
-        schedule_configs = [
-            CosineScheduleConfig(
-                warmup_steps=np.random.randint(50, 200),
-                total_steps=np.random.randint(1000, 5000),
-                min_rate=np.random.uniform(0.01, 0.1),
-            ),
-            SGDRScheduleConfig(
-                warmup_steps=np.random.randint(50, 200),
-                cycle_length=np.random.randint(500, 2000),
-                t_mult=np.random.randint(1, 3),
-                min_rate=np.random.uniform(0.01, 0.1),
-            ),
-        ]
-
-        for config in schedule_configs:
-            # Validate configuration
-            config.validate()
-
-            # Create scheduler using factory
-            scheduler = SchedulerFactory.load(optimizer, config)
-
-            # Verify scheduler type
-            if isinstance(config, CosineScheduleConfig):
+        # Test multiple random configurations
+        cosine_params = {
+            "schedule_type": "cosine",
+            "warmup_steps": np.random.randint(50, 200),
+            "total_steps": np.random.randint(1000, 5000),
+            "min_rate": np.random.uniform(0.01, 0.1),
+        }
+        sgdr_params = {
+            "schedule_type": "sgdr",
+            "warmup_steps": np.random.randint(50, 200),
+            "cycle_length": np.random.randint(500, 2000),
+            "t_mult": np.random.randint(1, 3),
+            "min_rate": np.random.uniform(0.01, 0.1),
+        }
+        for params in [cosine_params, sgdr_params]:
+            schedule_type = params["schedule_type"]
+            # Convert parameters for scheduler constructor
+            if schedule_type == "cosine":
+                warmup_steps = params["warmup_steps"]
+                total_steps = params["total_steps"]
+                min_rate = params["min_rate"]
+                lr_decay_steps = total_steps - warmup_steps
+                scheduler = SchedulerFactory.create(
+                    optimizer,
+                    schedule_type,
+                    warmup_steps=warmup_steps,
+                    lr_decay_steps=lr_decay_steps,
+                    min_rate=min_rate,
+                )
                 assert isinstance(scheduler, CosineScheduler)
-                assert scheduler.warmup_steps == config.warmup_steps
-                assert (
-                    scheduler.lr_decay_steps == config.total_steps - config.warmup_steps
+                assert scheduler.warmup_steps == warmup_steps
+                assert scheduler.lr_decay_steps == lr_decay_steps
+                assert scheduler.min_rate == min_rate
+            elif schedule_type == "sgdr":
+                warmup_steps = params["warmup_steps"]
+                cycle_length = params["cycle_length"]
+                t_mult = params["t_mult"]
+                min_rate = params["min_rate"]
+                scheduler = SchedulerFactory.create(
+                    optimizer,
+                    schedule_type,
+                    warmup_steps=warmup_steps,
+                    cycle_length=cycle_length,
+                    t_mult=t_mult,
+                    min_rate=min_rate,
                 )
-                assert scheduler.min_rate == config.min_rate
-            elif isinstance(config, SGDRScheduleConfig):
                 assert isinstance(scheduler, SGDRScheduler)
-                assert scheduler.warmup_steps == config.warmup_steps
-                assert scheduler.cycle_length == config.cycle_length
-                assert scheduler.t_mult == config.t_mult
-                assert scheduler.min_rate == config.min_rate
+                assert scheduler.warmup_steps == warmup_steps
+                assert scheduler.cycle_length == cycle_length
+                assert scheduler.t_mult == t_mult
+                assert scheduler.min_rate == min_rate
 
             # Test scheduler state dict functionality
             state_dict = scheduler.state_dict()
@@ -76,16 +91,25 @@ def test_schedule_factory_edge_cases():
     # Test edge cases for CosineScheduleConfig
     edge_cases = [
         # Minimal warmup and steps
-        CosineScheduleConfig(warmup_steps=1, total_steps=10, min_rate=0.01),
+        {"warmup_steps": 1, "total_steps": 10, "min_rate": 0.01},
         # Large values
-        CosineScheduleConfig(warmup_steps=1000, total_steps=10000, min_rate=0.5),
+        {"warmup_steps": 1000, "total_steps": 10000, "min_rate": 0.5},
         # Zero min_rate (edge case)
-        CosineScheduleConfig(warmup_steps=100, total_steps=1000, min_rate=0.0),
+        {"warmup_steps": 100, "total_steps": 1000, "min_rate": 0.0},
     ]
 
-    for config in edge_cases:
-        config.validate()
-        scheduler = SchedulerFactory.load(optimizer, config)
+    for params in edge_cases:
+        warmup_steps = params["warmup_steps"]
+        total_steps = params["total_steps"]
+        min_rate = params["min_rate"]
+        lr_decay_steps = total_steps - warmup_steps
+        scheduler = SchedulerFactory.create(
+            optimizer,
+            "cosine",
+            warmup_steps=warmup_steps,
+            lr_decay_steps=lr_decay_steps,
+            min_rate=min_rate,
+        )
         assert scheduler is not None
 
         # Test multiple steps
@@ -93,34 +117,24 @@ def test_schedule_factory_edge_cases():
             scheduler.step()
 
 
-def test_schedule_factory_invalid_configs():
-    """Test scheduler factory with invalid configurations"""
-
-    # Test invalid configurations that should raise errors
-    invalid_configs = [
-        # Negative warmup steps
-        {"warmup_steps": -10, "total_steps": 1000, "min_rate": 0.1},
-        # Total steps less than warmup steps
-        {"warmup_steps": 500, "total_steps": 400, "min_rate": 0.1},
-        # Invalid min_rate
-        {"warmup_steps": 100, "total_steps": 1000, "min_rate": -0.1},
-        {"warmup_steps": 100, "total_steps": 1000, "min_rate": 1.1},
-    ]
-
-    for kwargs in invalid_configs:
-        with pytest.raises(ValueError):
-            config = CosineScheduleConfig(**kwargs)
-            config.validate()
-
-
 def test_schedule_factory_state_persistence():
     """Test scheduler state persistence (save/load)"""
 
     model = torch.nn.Linear(10, 2)
     optimizer = torch.optim.AdamW(model.parameters(), lr=0.001)
 
-    config = CosineScheduleConfig(warmup_steps=100, total_steps=1000, min_rate=0.1)
-    scheduler = SchedulerFactory.load(optimizer, config)
+    # Create scheduler directly with parameters
+    warmup_steps = 100
+    total_steps = 1000
+    min_rate = 0.1
+    lr_decay_steps = total_steps - warmup_steps
+    scheduler = SchedulerFactory.create(
+        optimizer,
+        "cosine",
+        warmup_steps=warmup_steps,
+        lr_decay_steps=lr_decay_steps,
+        min_rate=min_rate,
+    )
 
     # Take a few steps
     for _ in range(5):
@@ -129,8 +143,14 @@ def test_schedule_factory_state_persistence():
     # Save state
     state_dict = scheduler.state_dict()
 
-    # Create new scheduler and load state
-    new_scheduler = SchedulerFactory.load(optimizer, config)
+    # Create new scheduler with same parameters
+    new_scheduler = SchedulerFactory.create(
+        optimizer,
+        "cosine",
+        warmup_steps=warmup_steps,
+        lr_decay_steps=lr_decay_steps,
+        min_rate=min_rate,
+    )
     new_scheduler.load_state_dict(state_dict)
 
     # Verify states match