style: 修改为显式导入

This commit is contained in:
2026-04-04 16:02:49 +08:00
parent 3346c75584
commit b531232a9b
12 changed files with 38 additions and 34 deletions
+10 -6
View File
@@ -3,18 +3,22 @@ import os
import numpy as np
import torch
from astrai.config import *
from astrai.config.train_config import TrainConfig
from astrai.data.serialization import Checkpoint
from astrai.trainer import *
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"""
optimizer_fn = lambda model: torch.optim.AdamW(model.parameters())
scheduler_fn = lambda optim: SchedulerFactory.create(
optim, "cosine", warmup_steps=10, lr_decay_steps=10, min_rate=0.05
)
def optimizer_fn(model):
return torch.optim.AdamW(model.parameters())
def scheduler_fn(optim):
return SchedulerFactory.create(
optim, "cosine", warmup_steps=10, lr_decay_steps=10, min_rate=0.05
)
train_config = TrainConfig(
strategy="seq",