ViperEkura
6ae1828449
- 删除 Registry 中未使用的 category/priority 字段，_entries 简化为直接存储类引用 - 修正 __init_subclass__ 避免叶子类（AutoRegressiveLM 等）创建空注册表 - 删除 5 个工厂的薄 create() 覆写，统一使用 BaseFactory.create(name, *args, **kwargs) - 删除 3 处零调用的 available_types/available_strategies 别名死代码 - 删除零调用的 BaseModelConfig.to_file 死代码 - 将 BaseConfig.from_json/to_json 重命名为 from_file/to_file，消除与子类重复 - 移除两个 inference builder 中总是被覆写的 prompt_tokens=0
58 lines
1.6 KiB
Python
58 lines
1.6 KiB
Python
import os
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from astrai.config.train_config import TrainConfig
|
|
from astrai.trainer.schedule import SchedulerFactory
|
|
from astrai.trainer.trainer import Trainer
|
|
|
|
|
|
def test_early_stopping_simulation(base_test_env, early_stopping_dataset):
|
|
"""Simulate early stopping behavior"""
|
|
|
|
def optimizer_fn(model):
|
|
return torch.optim.AdamW(model.parameters())
|
|
|
|
def scheduler_fn(optim):
|
|
return SchedulerFactory.create(
|
|
"cosine", optim, warmup_steps=10, lr_decay_steps=10, min_rate=0.05
|
|
)
|
|
|
|
train_config = TrainConfig(
|
|
strategy="seq",
|
|
optimizer_fn=optimizer_fn,
|
|
scheduler_fn=scheduler_fn,
|
|
model_fn=lambda: base_test_env["model"],
|
|
dataset=early_stopping_dataset,
|
|
ckpt_dir=base_test_env["test_dir"],
|
|
log_dir=os.path.join(base_test_env["test_dir"], "logs"),
|
|
n_epoch=2,
|
|
batch_per_device=2,
|
|
ckpt_interval=1,
|
|
grad_accum_steps=2,
|
|
random_seed=np.random.randint(1e4),
|
|
device_type=base_test_env["device"],
|
|
)
|
|
|
|
trainer = Trainer(train_config)
|
|
|
|
# Should handle early stopping gracefully
|
|
try:
|
|
trainer.train()
|
|
except Exception:
|
|
pass
|
|
|
|
# Resume from latest checkpoint
|
|
load_dir = os.path.join(base_test_env["test_dir"], "epoch_0_iter_2")
|
|
trainer = Trainer(train_config)
|
|
trainer.train(resume_dir=load_dir)
|
|
|
|
# Verify checkpoint was saved at expected iteration
|
|
load_dir = os.path.join(base_test_env["test_dir"], "epoch_1_iter_10")
|
|
import json
|
|
|
|
with open(os.path.join(load_dir, "meta.json")) as f:
|
|
meta = json.load(f)
|
|
assert meta["iteration"] == 10
|