AstrAI/khaosz/trainer/__init__.py

19 lines
412 B
Python

from khaosz.trainer.dataset import DatasetLoader
from khaosz.trainer.trainer import Trainer
from khaosz.trainer.strategy import (
TrainConfig,
CosineScheduleConfig,
SgdrScheduleConfig,
StrategyFactory,
SchedulerFactory
)
__all__ = [
"DatasetLoader",
"Trainer",
"TrainConfig",
"CosineScheduleConfig",
"SgdrScheduleConfig",
"StrategyFactory",
"SchedulerFactory"
]