AstrAI/khaosz/trainer/__init__.py

32 lines
691 B
Python

from khaosz.trainer.data_util import DatasetLoader
from khaosz.trainer.trainer import Trainer
from khaosz.trainer.strategy import (
TrainConfig,
CosineScheduleConfig,
SgdrScheduleConfig,
StrategyFactory,
SchedulerFactory
)
from khaosz.trainer.trainer_callback import (
ProgressBarCallback,
CheckpointCallback,
TrainerCallback,
SchedulerCallback
)
__all__ = [
# strategy
"DatasetLoader",
"Trainer",
"TrainConfig",
"CosineScheduleConfig",
"SgdrScheduleConfig",
"StrategyFactory",
"SchedulerFactory",
# callback
"ProgressBarCallback",
"CheckpointCallback",
"TrainerCallback",
"SchedulerCallback",
]