63 lines
1.4 KiB
Python
63 lines
1.4 KiB
Python
|
"""Typed training configuration system."""
|
||
|
|
|
||
|
|
from .data_config import DataConfig, LeRobotDataConfig
|
||
|
|
from .hyperparams_config import (
|
||
|
|
AdamWConfig,
|
||
|
|
ConstantSchedulerConfig,
|
||
|
|
CosineSchedulerConfig,
|
||
|
|
DMuonConfig,
|
||
|
|
OptimizerConfig,
|
||
|
|
SchedulerConfig,
|
||
|
|
StepSchedulerConfig,
|
||
|
|
TrainHyperParams,
|
||
|
|
)
|
||
|
|
from .infra_config import (
|
||
|
|
CheckpointConfig,
|
||
|
|
DebugConfig,
|
||
|
|
DistributedConfig,
|
||
|
|
LoggingConfig,
|
||
|
|
)
|
||
|
|
from .loader import load_config, save_config
|
||
|
|
from .model_config import ModelConfig, QActModelConfig
|
||
|
|
from .registry import (
|
||
|
|
register_data_config,
|
||
|
|
register_model_config,
|
||
|
|
register_optimizer_config,
|
||
|
|
register_scheduler_config,
|
||
|
|
)
|
||
|
|
from .task_config import TaskConfig
|
||
|
|
from .train_config import TrainConfig
|
||
|
|
|
||
|
|
__all__ = [
|
||
|
|
# Core
|
||
|
|
"TrainConfig",
|
||
|
|
"TaskConfig",
|
||
|
|
# Model configs
|
||
|
|
"ModelConfig",
|
||
|
|
"QActModelConfig",
|
||
|
|
# Data configs
|
||
|
|
"DataConfig",
|
||
|
|
"LeRobotDataConfig",
|
||
|
|
# Hyperparams
|
||
|
|
"TrainHyperParams",
|
||
|
|
"OptimizerConfig",
|
||
|
|
"AdamWConfig",
|
||
|
|
"DMuonConfig",
|
||
|
|
"SchedulerConfig",
|
||
|
|
"CosineSchedulerConfig",
|
||
|
|
"ConstantSchedulerConfig",
|
||
|
|
"StepSchedulerConfig",
|
||
|
|
# Infra configs
|
||
|
|
"DistributedConfig",
|
||
|
|
"LoggingConfig",
|
||
|
|
"CheckpointConfig",
|
||
|
|
"DebugConfig",
|
||
|
|
# Functions
|
||
|
|
"load_config",
|
||
|
|
"save_config",
|
||
|
|
"register_data_config",
|
||
|
|
"register_model_config",
|
||
|
|
"register_optimizer_config",
|
||
|
|
"register_scheduler_config",
|
||
|
|
]
|