AstrAI/khaosz/trainer/__init__.py

30 lines
631 B
Python

from khaosz.trainer.trainer import Trainer
from khaosz.trainer.strategy import (
CosineScheduleConfig,
SgdrScheduleConfig,
StrategyFactory,
SchedulerFactory
)
from khaosz.trainer.train_callback import (
TrainCallback,
ProgressBarCallback,
CheckpointCallback,
TrainCallback,
SchedulerCallback,
StepMonitorCallback
)
__all__ = [
"Trainer",
"StrategyFactory",
"CosineScheduleConfig",
"SgdrScheduleConfig",
"SchedulerFactory",
# callback
"TrainCallback",
"ProgressBarCallback",
"CheckpointCallback",
"SchedulerCallback",
"StepMonitorCallback"
]