Base Configurations#
We can integrate tyro.cli() into common configuration patterns: here, we select
one of multiple possible base configurations, create a subcommand for each one, and then
use the CLI to either override (existing) or fill in (missing) values.
1from dataclasses import dataclass
2from typing import Callable, Literal
3
4from torch import nn
5
6import tyro
7
8
9@dataclass(frozen=True)
10class AdamOptimizer:
11 learning_rate: float = 1e-3
12 betas: tuple[float, float] = (0.9, 0.999)
13
14
15@dataclass(frozen=True)
16class ExperimentConfig:
17 # Dataset to run experiment on.
18 dataset: Literal["mnist", "imagenet-50"]
19
20 # Optimizer parameters.
21 optimizer: AdamOptimizer
22
23 # Model size.
24 num_layers: int
25 units: int
26
27 # Batch size.
28 batch_size: int
29
30 # Total number of training steps.
31 train_steps: int
32
33 # Random seed. This is helpful for making sure that our experiments are all
34 # reproducible!
35 seed: int
36
37 # Activation to use. Not specifiable via the commandline.
38 activation: Callable[[], nn.Module]
39
40
41# Note that we could also define this library using separate YAML files (similar to
42# `config_path`/`config_name` in Hydra), but staying in Python enables seamless type
43# checking + IDE support.
44Configs = tyro.extras.subcommand_type_from_defaults(
45 {
46 "small": ExperimentConfig(
47 dataset="mnist",
48 optimizer=AdamOptimizer(),
49 batch_size=2048,
50 num_layers=4,
51 units=64,
52 train_steps=30_000,
53 seed=0,
54 activation=nn.ReLU,
55 ),
56 "big": ExperimentConfig(
57 dataset="imagenet-50",
58 optimizer=AdamOptimizer(),
59 batch_size=32,
60 num_layers=8,
61 units=256,
62 train_steps=100_000,
63 seed=0,
64 activation=nn.GELU,
65 ),
66 }
67)
68
69if __name__ == "__main__":
70 config = tyro.cli(Configs)
71 print(config)
python 03_config_systems/01_base_configs.py --help
usage: 01_base_configs.py [-h] {small,big}
╭─ options ──────────────────────────────────────────╮
│ -h, --help show this help message and exit │
╰────────────────────────────────────────────────────╯
╭─ subcommands ──────────────────────────────────────╮
│ {small,big} │
│ small │
│ big │
╰────────────────────────────────────────────────────╯
python 03_config_systems/01_base_configs.py small --help
usage: 01_base_configs.py small [-h] [SMALL OPTIONS]
╭─ options ──────────────────────────────────────────────────────────────────╮
│ -h, --help show this help message and exit │
│ --dataset {mnist,imagenet-50} │
│ Dataset to run experiment on. (default: mnist) │
│ --num-layers INT Model size. (default: 4) │
│ --units INT Model size. (default: 64) │
│ --batch-size INT Batch size. (default: 2048) │
│ --train-steps INT Total number of training steps. (default: 30000) │
│ --seed INT Random seed. This is helpful for making sure that │
│ our experiments are all reproducible! (default: 0) │
│ --activation {fixed} Activation to use. Not specifiable via the │
│ commandline. (fixed to: <class │
│ 'torch.nn.modules.activation.ReLU'>) │
╰────────────────────────────────────────────────────────────────────────────╯
╭─ optimizer options ────────────────────────────────────────────────────────╮
│ Optimizer parameters. │
│ ──────────────────────────────────────────── │
│ --optimizer.learning-rate FLOAT │
│ (default: 0.001) │
│ --optimizer.betas FLOAT FLOAT │
│ (default: 0.9 0.999) │
╰────────────────────────────────────────────────────────────────────────────╯
python 03_config_systems/01_base_configs.py small --seed 94720
ExperimentConfig(dataset='mnist', optimizer=AdamOptimizer(learning_rate=0.001, betas=(0.9, 0.999)), num_layers=4, units=64, batch_size=2048, train_steps=30000, seed=94720, activation=<class 'torch.nn.modules.activation.ReLU'>)
python 03_config_systems/01_base_configs.py big --help
usage: 01_base_configs.py big [-h] [BIG OPTIONS]
╭─ options ──────────────────────────────────────────────────────────────────╮
│ -h, --help show this help message and exit │
│ --dataset {mnist,imagenet-50} │
│ Dataset to run experiment on. (default: │
│ imagenet-50) │
│ --num-layers INT Model size. (default: 8) │
│ --units INT Model size. (default: 256) │
│ --batch-size INT Batch size. (default: 32) │
│ --train-steps INT Total number of training steps. (default: 100000) │
│ --seed INT Random seed. This is helpful for making sure that │
│ our experiments are all reproducible! (default: 0) │
│ --activation {fixed} Activation to use. Not specifiable via the │
│ commandline. (fixed to: <class │
│ 'torch.nn.modules.activation.GELU'>) │
╰────────────────────────────────────────────────────────────────────────────╯
╭─ optimizer options ────────────────────────────────────────────────────────╮
│ Optimizer parameters. │
│ ──────────────────────────────────────────── │
│ --optimizer.learning-rate FLOAT │
│ (default: 0.001) │
│ --optimizer.betas FLOAT FLOAT │
│ (default: 0.9 0.999) │
╰────────────────────────────────────────────────────────────────────────────╯
python 03_config_systems/01_base_configs.py big --seed 94720
ExperimentConfig(dataset='imagenet-50', optimizer=AdamOptimizer(learning_rate=0.001, betas=(0.9, 0.999)), num_layers=8, units=256, batch_size=32, train_steps=100000, seed=94720, activation=<class 'torch.nn.modules.activation.GELU'>)