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

style: 使用ruff 工具优化代码风格

Browse Source
This commit is contained in:
ViperEkura
2026-03-30 23:32:28 +08:00
parent 345fd2f091
commit 426af2d75f
52 changed files with 1836 additions and 1493 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/trainer/test_callbacks.py
+20 -29
View File
@@ -3,57 +3,48 @@ 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
)
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)
train_config = TrainConfig(
model=base_test_env["model"],
strategy='seq',
strategy="seq",
dataset=random_dataset,
optimizer_fn=optimizer_fn,
scheduler_fn=scheduler_fn,
checkpoint_dir=base_test_env["test_dir"],
ckpt_dir=base_test_env["test_dir"],
n_epoch=1,
batch_size=2,
checkpoint_interval=3,
ckpt_interval=3,
accumulation_steps=1,
max_grad_norm=1.0,
random_seed=42,
device_type=base_test_env["device"]
device_type=base_test_env["device"],
)
# Create custom callbacks to track calls
callback_calls = []
class TrackingCallback(TrainCallback):
def on_train_begin(self, context):
callback_calls.append('on_train_begin')
callback_calls.append("on_train_begin")
def on_batch_end(self, context):
callback_calls.append('on_batch_end')
callback_calls.append("on_batch_end")
def on_epoch_end(self, context):
callback_calls.append('on_epoch_end')
callback_calls.append("on_epoch_end")
trainer = Trainer(train_config, callbacks=[TrackingCallback()])
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
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
+11 -10
View File
@@ -5,31 +5,32 @@ from khaosz.config import *
from khaosz.trainer import *
from khaosz.data.serialization 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_fn = lambda model: torch.optim.AdamW(model.parameters())
scheduler_fn = lambda optim: SchedulerFactory.load(optim, schedule_config)
train_config = TrainConfig(
strategy="seq",
optimizer_fn=optimizer_fn,
scheduler_fn=scheduler_fn,
model=base_test_env["model"],
dataset=early_stopping_dataset,
checkpoint_dir=base_test_env["test_dir"],
ckpt_dir=base_test_env["test_dir"],
n_epoch=2,
batch_size=2,
checkpoint_interval=1,
ckpt_interval=1,
accumulation_steps=2,
random_seed=np.random.randint(1e4),
device_type=base_test_env["device"]
device_type=base_test_env["device"],
)
trainer = Trainer(train_config)
# Should handle early stopping gracefully
checkpoint = None
try:
@@ -37,11 +38,11 @@ def test_early_stopping_simulation(base_test_env, early_stopping_dataset):
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
assert checkpoint.iteration == 10
tests/trainer/test_train_strategy.py
+29 -27
View File
@@ -9,39 +9,41 @@ 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)
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)
)
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.lr_decay_steps == config.total_steps - config.warmup_steps
)
assert scheduler.min_rate == config.min_rate
elif isinstance(config, SGDRScheduleConfig):
assert isinstance(scheduler, SGDRScheduler)
@@ -49,17 +51,17 @@ def test_schedule_factory_random_configs():
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
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
@@ -67,10 +69,10 @@ def test_schedule_factory_random_configs():
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
@@ -80,12 +82,12 @@ def test_schedule_factory_edge_cases():
# 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()
@@ -93,7 +95,7 @@ def test_schedule_factory_edge_cases():
def test_schedule_factory_invalid_configs():
"""Test scheduler factory with invalid configurations"""
# Test invalid configurations that should raise errors
invalid_configs = [
# Negative warmup steps
@@ -104,7 +106,7 @@ def test_schedule_factory_invalid_configs():
{"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)
@@ -113,24 +115,24 @@ def test_schedule_factory_invalid_configs():
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()
assert scheduler.get_last_lr() == new_scheduler.get_last_lr()
tests/trainer/test_trainer.py
+27 -33
View File
@@ -6,100 +6,94 @@ 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
)
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)
train_config = TrainConfig(
strategy="seq",
model=base_test_env["model"],
dataset=random_dataset,
optimizer_fn=optimizer_fn,
scheduler_fn=scheduler_fn,
checkpoint_dir=base_test_env["test_dir"],
ckpt_dir=base_test_env["test_dir"],
n_epoch=1,
batch_size=batch_size,
checkpoint_interval=5,
ckpt_interval=5,
accumulation_steps=1,
max_grad_norm=1.0,
random_seed=np.random.randint(1000),
device_type=base_test_env["device"]
device_type=base_test_env["device"],
)
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
)
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)
train_config = TrainConfig(
strategy="seq",
model=base_test_env["model"],
optimizer_fn=optimizer_fn,
scheduler_fn=scheduler_fn,
dataset=random_dataset,
checkpoint_dir=base_test_env["test_dir"],
ckpt_dir=base_test_env["test_dir"],
n_epoch=1,
batch_size=2,
checkpoint_interval=10,
ckpt_interval=10,
accumulation_steps=accumulation_steps,
max_grad_norm=1.0,
random_seed=42,
device_type=base_test_env["device"]
device_type=base_test_env["device"],
)
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}
{"batch_size": 4, "accumulation_steps": 2},
]
for config in small_batch_configs:
schedule_config = CosineScheduleConfig(
warmup_steps=10,
total_steps=20
)
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)
train_config = TrainConfig(
strategy="seq",
model=base_test_env["model"],
dataset=random_dataset,
optimizer_fn=optimizer_fn,
scheduler_fn=scheduler_fn,
checkpoint_dir=base_test_env["test_dir"],
ckpt_dir=base_test_env["test_dir"],
n_epoch=1,
batch_size=config["batch_size"],
checkpoint_interval=5,
ckpt_interval=5,
accumulation_steps=config["accumulation_steps"],
max_grad_norm=1.0,
random_seed=42,
device_type=base_test_env["device"]
device_type=base_test_env["device"],
)
assert train_config.accumulation_steps == config["accumulation_steps"]
assert train_config.accumulation_steps == config["accumulation_steps"]