style: 修改为显式导入
This commit is contained in:
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user