AstrAI/khaosz/config/__init__.py

25 lines
582 B
Python

from khaosz.config.model_config import ModelConfig
from khaosz.config.param_config import BaseModelIO, ModelParameter
from khaosz.config.schedule_config import (
ScheduleConfig,
CosineScheduleConfig,
SGDRScheduleConfig,
ScheduleConfigFactory,
)
from khaosz.config.train_config import TrainConfig
__all__ = [
# Base I/O
"BaseModelIO",
"ModelParameter",
# Model configuration
"ModelConfig",
"TrainConfig",
# Schedule configuration
"ScheduleConfig",
"CosineScheduleConfig",
"SGDRScheduleConfig",
"ScheduleConfigFactory",
]