ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
ViperEkura
2026-04-03 22:06:32 +08:00
parent 3a7d98a950
commit 0852b852f8
51 changed files with 299 additions and 434 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/trainer/conftest.py
+7 -6
View File
@@ -1,6 +1,9 @@
import pytest
import torch
from torch.utils.data import Dataset
import pytest
from astrai.config import TrainConfig
from astrai.trainer.schedule import SchedulerFactory
class TrainerDataset(Dataset):
@@ -54,13 +57,11 @@ def create_train_config(
Returns:
TrainConfig instance configured for testing
"""
from astrai.config import TrainConfig
from astrai.config.schedule_config import CosineScheduleConfig
from astrai.trainer.schedule import SchedulerFactory
schedule_config = CosineScheduleConfig(warmup_steps=10, total_steps=20)
optimizer_fn = lambda m: torch.optim.AdamW(m.parameters(), lr=0.001)
scheduler_fn = lambda optim: SchedulerFactory.load(optim, schedule_config)
scheduler_fn = lambda optim: SchedulerFactory.create(
optim, "cosine", warmup_steps=10, lr_decay_steps=10, min_rate=0.05
)
return TrainConfig(
strategy=strategy,
tests/trainer/test_callbacks.py
+3 -3
View File
@@ -6,10 +6,10 @@ from astrai.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_fn = lambda model: torch.optim.AdamW(model.parameters())
scheduler_fn = lambda optim: SchedulerFactory.load(optim, schedule_config)
scheduler_fn = lambda optim: SchedulerFactory.create(
optim, "cosine", warmup_steps=10, lr_decay_steps=10, min_rate=0.05
)
train_config = TrainConfig(
model=base_test_env["model"],
tests/trainer/test_early_stopping.py
+7 -5
View File
@@ -1,18 +1,20 @@
import os
import torch
import numpy as np
import torch
from astrai.config import *
from astrai.trainer import *
from astrai.data.serialization import Checkpoint
from astrai.trainer import *
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_fn = lambda model: torch.optim.AdamW(model.parameters())
scheduler_fn = lambda optim: SchedulerFactory.load(optim, schedule_config)
scheduler_fn = lambda optim: SchedulerFactory.create(
optim, "cosine", warmup_steps=10, lr_decay_steps=10, min_rate=0.05
)
train_config = TrainConfig(
strategy="seq",
tests/trainer/test_train_strategy.py
+85 -65
View File
@@ -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
tests/trainer/test_trainer.py
-2
View File
@@ -1,5 +1,3 @@
import torch
from astrai.data.dataset import *
from astrai.trainer import Trainer