refactor(tests): 重构测试文件目录结构

tests/trainer/test_callbacks.py

@@ -0,0 +1,58 @@
+import torch
+
+from khaosz.config import *
+from khaosz.trainer import *
+
+def test_callback_integration(base_test_env, random_dataset):
+    """Test that all callbacks are properly integrated"""
+    schedule_config = CosineScheduleConfig(
+        warmup_steps=10,
+        total_steps=20
+    )
+    
+    optimizer = torch.optim.AdamW(base_test_env["model"].parameters())
+    scheduler = SchedulerFactory.load(optimizer, schedule_config)
+    
+    train_config = TrainConfig(
+        model=base_test_env["model"],
+        strategy='seq',
+        dataset=random_dataset,
+        optimizer=optimizer,
+        scheduler=scheduler,
+        checkpoint_dir=base_test_env["test_dir"],
+        n_epoch=1,
+        batch_size=2,
+        checkpoint_interval=3,
+        accumulation_steps=1,
+        max_grad_norm=1.0,
+        random_seed=42
+    )
+    
+
+    
+    # Create custom callbacks to track calls
+    callback_calls = []
+    
+    class TrackingCallback(TrainCallback):
+        def on_train_begin(self, context):
+            callback_calls.append('on_train_begin')
+        
+        def on_batch_end(self, context):
+            callback_calls.append('on_batch_end')
+        
+        def on_epoch_end(self, context):
+            callback_calls.append('on_epoch_end')
+    
+
+    
+    trainer = Trainer(
+        train_config, 
+        callbacks=[TrackingCallback()]
+    )
+    
+    trainer.train()
+    
+    # Verify callbacks were called
+    assert 'on_train_begin' in callback_calls
+    assert 'on_batch_end' in callback_calls
+    assert 'on_epoch_end' in callback_calls

tests/trainer/test_early_stopping.py

@@ -0,0 +1,45 @@
+import os
+import torch
+import numpy as np
+from khaosz.config import *
+from khaosz.trainer import *
+from khaosz.data.checkpoint import Checkpoint
+
+def test_early_stopping_simulation(base_test_env, early_stopping_dataset):
+    """Simulate early stopping behavior"""
+    
+    schedule_config = CosineScheduleConfig(warmup_steps=10, total_steps=20)
+    optimizer = torch.optim.AdamW(base_test_env["model"].parameters())
+    scheduler = SchedulerFactory.load(optimizer, schedule_config)
+    
+    train_config = TrainConfig(
+        strategy="seq",
+        scheduler=scheduler,
+        model=base_test_env["model"],
+        dataset=early_stopping_dataset,
+        optimizer=optimizer,
+        checkpoint_dir=base_test_env["test_dir"],
+        n_epoch=2,
+        batch_size=2,
+        checkpoint_interval=1,
+        accumulation_steps=2,
+        random_seed=np.random.randint(1e4),
+    )
+
+    trainer = Trainer(train_config)
+    
+    # Should handle early stopping gracefully
+    checkpoint = None
+    try:
+        checkpoint = trainer.train()
+    except Exception:
+        # Handle any exceptions
+        pass
+    
+    load_dir = os.path.join(base_test_env["test_dir"], "epoch_0_iter_2")
+    checkpoint = Checkpoint.load(load_dir)
+    trainer.train(checkpoint)
+    
+    load_dir = os.path.join(base_test_env["test_dir"], "epoch_1_iter_10")
+    checkpoint = Checkpoint.load(load_dir)
+    assert checkpoint.iteration == 10

tests/trainer/test_train_strategy.py

@@ -0,0 +1,136 @@
+import torch
+import numpy as np
+import pytest
+
+from khaosz.config import *
+from khaosz.trainer.schedule import *
+from khaosz.data.dataset import *
+
+
+def test_schedule_factory_random_configs():
+    """Test scheduler factory with random configurations"""
+    
+    # Create a simple model and optimizer for testing
+    model = torch.nn.Linear(10, 2)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=0.001)
+    
+    # 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):
+                assert isinstance(scheduler, CosineScheduler)
+                assert scheduler.warmup_steps == config.warmup_steps
+                assert scheduler.lr_decay_steps == config.total_steps - config.warmup_steps
+                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
+            
+            # Test scheduler state dict functionality
+            state_dict = scheduler.state_dict()
+            assert 'warmup_steps' in state_dict
+            assert 'min_rate' in state_dict
+            
+            # Test scheduler step functionality
+            initial_lr = scheduler.get_last_lr()
+            scheduler.step()
+            new_lr = scheduler.get_last_lr()
+            
+            # Learning rate should change after step, or if it's the first step,
+            # the epoch counter should increment
+            assert initial_lr != new_lr or scheduler.last_epoch > -1
+
+
+def test_schedule_factory_edge_cases():
+    """Test scheduler factory with edge cases and boundary conditions"""
+    
+    model = torch.nn.Linear(10, 2)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=0.001)
+    
+    # Test edge cases for CosineScheduleConfig
+    edge_cases = [
+        # Minimal warmup and steps
+        CosineScheduleConfig(warmup_steps=1, total_steps=10, min_rate=0.01),
+        # Large values
+        CosineScheduleConfig(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),
+    ]
+    
+    for config in edge_cases:
+        config.validate()
+        scheduler = SchedulerFactory.load(optimizer, config)
+        assert scheduler is not None
+        
+        # Test multiple steps
+        for _ in range(10):
+            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)
+    
+    # Take a few steps
+    for _ in range(5):
+        scheduler.step()
+    
+    # Save state
+    state_dict = scheduler.state_dict()
+    
+    # Create new scheduler and load state
+    new_scheduler = SchedulerFactory.load(optimizer, config)
+    new_scheduler.load_state_dict(state_dict)
+    
+    # Verify states match
+    assert scheduler.last_epoch == new_scheduler.last_epoch
+    assert scheduler.get_last_lr() == new_scheduler.get_last_lr()

tests/trainer/test_trainer.py

@@ -0,0 +1,99 @@
+import torch
+import numpy as np
+
+
+from khaosz.config import *
+from khaosz.trainer import *
+from khaosz.data.dataset import *
+
+def test_different_batch_sizes(base_test_env, random_dataset):
+    """Test training with different batch sizes"""
+    batch_sizes = [1, 2, 4, 8]
+    
+    for batch_size in batch_sizes:
+        schedule_config = CosineScheduleConfig(
+            warmup_steps=10,
+            total_steps=20
+        )
+        optimizer = torch.optim.AdamW(base_test_env["model"].parameters())
+        scheduler = SchedulerFactory.load(optimizer, schedule_config)
+        train_config = TrainConfig(
+            strategy="seq",
+            model=base_test_env["model"],
+            dataset=random_dataset,
+            optimizer=optimizer,
+            scheduler=scheduler,
+            checkpoint_dir=base_test_env["test_dir"],
+            n_epoch=1,
+            batch_size=batch_size,
+            checkpoint_interval=5,
+            accumulation_steps=1,
+            max_grad_norm=1.0,
+            random_seed=np.random.randint(1000)
+        )
+        
+        assert train_config.batch_size == batch_size
+
+def test_gradient_accumulation(base_test_env, random_dataset):
+    """Test training with different gradient accumulation steps"""
+    accumulation_steps_list = [1, 2, 4]
+    
+    for accumulation_steps in accumulation_steps_list:
+        schedule_config = CosineScheduleConfig(
+            warmup_steps=10,
+            total_steps=20
+        )
+        optimizer = torch.optim.AdamW(base_test_env["model"].parameters())
+        scheduler = SchedulerFactory.load(optimizer, schedule_config)
+        train_config = TrainConfig(
+            strategy="seq",
+            model=base_test_env["model"],
+            optimizer=optimizer,
+            scheduler=scheduler,
+            dataset=random_dataset,
+            checkpoint_dir=base_test_env["test_dir"],
+            n_epoch=1,
+            batch_size=2,
+            checkpoint_interval=10,
+            accumulation_steps=accumulation_steps,
+            max_grad_norm=1.0,
+            random_seed=42
+        )
+        
+        trainer = Trainer(train_config)
+        trainer.train()
+        
+        assert train_config.accumulation_steps == accumulation_steps
+
+def test_memory_efficient_training(base_test_env, random_dataset):
+    """Test training with memory-efficient configurations"""
+    # Test with smaller batch sizes and gradient checkpointing
+    small_batch_configs = [
+        {"batch_size": 1, "accumulation_steps": 8},
+        {"batch_size": 2, "accumulation_steps": 4},
+        {"batch_size": 4, "accumulation_steps": 2}
+    ]
+    
+    for config in small_batch_configs:
+        schedule_config = CosineScheduleConfig(
+            warmup_steps=10,
+            total_steps=20
+        )
+        optimizer = torch.optim.AdamW(base_test_env["model"].parameters())
+        scheduler = SchedulerFactory.load(optimizer, schedule_config)
+        train_config = TrainConfig(
+            strategy="seq",
+            model=base_test_env["model"],
+            dataset=random_dataset,
+            optimizer=optimizer,
+            scheduler=scheduler,
+            checkpoint_dir=base_test_env["test_dir"],
+            n_epoch=1,
+            batch_size=config["batch_size"],
+            checkpoint_interval=5,
+            accumulation_steps=config["accumulation_steps"],
+            max_grad_norm=1.0,
+            random_seed=42
+        )
+        
+        assert train_config.accumulation_steps == config["accumulation_steps"]