| """``softasp-train`` — launch a training run (supports torchrun for DDP). |
| |
| Usage:: |
| |
| # Single-GPU |
| uv run --locked softasp-train instantiation=sudoku model.pretrained_from=backbone.pt |
| |
| # Multi-GPU via torchrun |
| uv run --locked torchrun --nproc_per_node=4 -m softasp.cli.train \\ |
| instantiation=my_distributed_task paths=local |
| """ |
|
|
| from __future__ import annotations |
|
|
| import logging |
| from pathlib import Path |
|
|
| import hydra |
| from omegaconf import DictConfig, OmegaConf |
| from softasp_core.training.schedule import CORRECTED_PROTOCOL, LEGACY_PROTOCOL |
| from softasp_core.training.seed import seed_everything |
| from softasp_core.training.trainer import Trainer |
| from softasp_core.utils.logging import setup_logging |
|
|
| from softasp.configuration import config_directory |
| from softasp.instantiations.checkpoint import read_checkpoint, validate_checkpoint |
| from softasp.instantiations.runtime import provider |
|
|
| logger = logging.getLogger(__name__) |
|
|
| _CONFIGS = config_directory() |
|
|
|
|
| @hydra.main(version_base="1.3", config_path=_CONFIGS, config_name="config") |
| def main(cfg: DictConfig) -> None: |
| run_training(cfg) |
|
|
|
|
| def run_training(cfg: DictConfig) -> Path: |
| """Train from an explicit configuration and return the exact run directory.""" |
| setup_logging(level=cfg.logging.get("level", "INFO")) |
| logger.info("=== softasp-train ===") |
|
|
| |
| OmegaConf.resolve(cfg) |
|
|
| protocol = cfg.trainer.get("protocol", LEGACY_PROTOCOL) |
| if protocol not in (LEGACY_PROTOCOL, CORRECTED_PROTOCOL): |
| raise ValueError(f"Unknown training protocol: {protocol}") |
| if protocol == CORRECTED_PROTOCOL: |
| |
| seed_everything(cfg.trainer.seed, deterministic=cfg.trainer.deterministic) |
| OmegaConf.update(cfg, "initialization_seeded_before_model", True, force_add=True) |
|
|
| model_provider = provider(cfg, "model") |
| |
| prepared = model_provider.prepare(cfg) |
| train_ds, val_ds = provider(cfg, "dataset").build(cfg) |
| model = prepared if prepared is not None else model_provider.build(cfg) |
| evaluator = provider(cfg, "metrics").build(cfg) |
|
|
| |
| trainer = Trainer( |
| cfg=cfg, model=model, evaluator=evaluator, |
| checkpoint_validator=lambda path: validate_checkpoint(cfg, read_checkpoint(path)), |
| ) |
| run_dir = trainer.fit(train_ds, val_ds) |
|
|
| logger.info("Training complete.") |
| return run_dir |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
