Hierarchical Configs

Structures (typically dataclasses) can be nested to build hierarchical configuration objects. This helps with modularity and grouping in larger projects.

import dataclasses
import enum
import pathlib

import tyro


class OptimizerType(enum.Enum):
    ADAM = enum.auto()
    SGD = enum.auto()


@dataclasses.dataclass
class OptimizerConfig:
    # Gradient-based optimizer to use.
    algorithm: OptimizerType = OptimizerType.ADAM

    # Learning rate to use.
    learning_rate: float = 3e-4

    # Coefficient for L2 regularization.
    weight_decay: float = 1e-2


@dataclasses.dataclass
class ExperimentConfig:
    # Various configurable options for our optimizer.
    optimizer: OptimizerConfig

    # Batch size.
    batch_size: int = 32

    # Total number of training steps.
    train_steps: int = 100_000

    # Random seed. This is helpful for making sure that our experiments are all
    # reproducible!
    seed: int = 0


def train(
    out_dir: pathlib.Path,
    config: ExperimentConfig,
    restore_checkpoint: bool = False,
    checkpoint_interval: int = 1000,
) -> None:
    """Train a model.

    Args:
        out_dir: Where to save logs and checkpoints.
        config: Experiment configuration.
        restore_checkpoint: Set to restore an existing checkpoint.
        checkpoint_interval: Training steps between each checkpoint save.
    """
    print(f"{out_dir=}, {restore_checkpoint=}, {checkpoint_interval=}")
    print()
    print(f"{config=}")


if __name__ == "__main__":
    tyro.cli(train)

---

python 02_nesting/01_nesting.py --help

usage: 01_nesting.py [-h] [OPTIONS]

Train a model.

╭─ options ──────────────────────────────────────────────────────────────────╮
│ -h, --help              show this help message and exit                    │
│ --out-dir PATH          Where to save logs and checkpoints. (required)     │
│ --restore-checkpoint, --no-restore-checkpoint                              │
│                         Set to restore an existing checkpoint. (default:   │
│                         False)                                             │
│ --checkpoint-interval INT                                                  │
│                         Training steps between each checkpoint save.       │
│                         (default: 1000)                                    │
╰────────────────────────────────────────────────────────────────────────────╯
╭─ config options ───────────────────────────────────────────────────────────╮
│ Experiment configuration.                                                  │
│ ────────────────────────────────────────────────────────────────────────── │
│ --config.batch-size INT                                                    │
│                         Batch size. (default: 32)                          │
│ --config.train-steps INT                                                   │
│                         Total number of training steps. (default: 100000)  │
│ --config.seed INT       Random seed. This is helpful for making sure that  │
│                         our experiments are all reproducible! (default: 0) │
╰────────────────────────────────────────────────────────────────────────────╯
╭─ config.optimizer options ─────────────────────────────────────────────────╮
│ --config.optimizer.algorithm {ADAM,SGD}                                    │
│                         Gradient-based optimizer to use. (default: ADAM)   │
│ --config.optimizer.learning-rate FLOAT                                     │
│                         Learning rate to use. (default: 0.0003)            │
│ --config.optimizer.weight-decay FLOAT                                      │
│                         Coefficient for L2 regularization. (default: 0.01) │
╰────────────────────────────────────────────────────────────────────────────╯

---

python 02_nesting/01_nesting.py --out-dir . --config.optimizer.algorithm SGD

out_dir=PosixPath('.'), restore_checkpoint=False, checkpoint_interval=1000

config=ExperimentConfig(optimizer=OptimizerConfig(algorithm=<OptimizerType.SGD: 2>, learning_rate=0.0003, weight_decay=0.01), batch_size=32, train_steps=100000, seed=0)

---

python 02_nesting/01_nesting.py --out-dir . --restore-checkpoint

out_dir=PosixPath('.'), restore_checkpoint=True, checkpoint_interval=1000

config=ExperimentConfig(optimizer=OptimizerConfig(algorithm=<OptimizerType.ADAM: 1>, learning_rate=0.0003, weight_decay=0.01), batch_size=32, train_steps=100000, seed=0)