heracls 0.4.0

Slayer of Hydra

Heracls - Slayer of Hydra

heracls is a tiny utility package to instantiate typed dataclasses from flexible config sources (dictionary, OmegaConf, YAML, dotlist, CLI, ...). It is designed for projects that want strict, typed config objects while supporting the dynamic overrides commonly used in scripts and experiments.

Installation

The heracls package is available on PyPi and can be installed with pip.

pip install heracls

Alternatively, if you need the latest features, you can install it from source.

pip install git+https://github.com/francois-rozet/heracls

Getting started

The following example demonstrates how to declare a nested dataclass config, instantiate it from command line arguments, serialize it to YAML, and use it in a script. The heracls.ArgumentParser parser infers arguments from the dataclass structure and types.

from dataclasses import dataclass, field
from typing import Literal, Union

import heracls

@dataclass
class ModelConfig:
    name: str = "mlp"
    depth: int = 3
    width: int = 256

@dataclass
class AdamConfig:
    name: Literal["adam"] = "adam"
    betas: tuple[float, float] = (0.95, 0.95)
    learning_rate: float = 1e-3
    weight_decay: float = 0.0

@dataclass
class SGDConfig:
    name: Literal["sgd"] = "sgd"
    momentum: float = 0.0
    learning_rate: float = 1e-3
    weight_decay: float = 0.0
    nesterov: bool = False

@dataclass
class TrainConfig:
    model: ModelConfig = field(default_factory=ModelConfig)
    optimizer: Union[AdamConfig, SGDConfig] = heracls.choice(
        {"adam": AdamConfig, "sgd": SGDConfig},
        default="adam",
    )
    dataset: str = "mnist"
    data_splits: tuple[float, ...] = (0.8, 0.1)
    n_epochs: int = 1024
    n_steps_per_epoch: int = 256
    tasks: list[str] = field(default_factory=list)

def main():
    parser = heracls.ArgumentParser()
    parser.add_argument("--dry", action="store_true", help="dry run")
    parser.add_arguments(TrainConfig, dest="train", root=True)

    args = parser.parse_args()

    print(heracls.to_yaml(args.train))

    if args.dry:
        return

    trainset, validset, testset = load_dataset(args.train.data_splits)

    model = init_model(args.train.model)

    for epoch in range(args.train.n_epochs):
        ...

if __name__ == "__main__":
    main()

$ python train.py --dry --model.depth 5 --optimizer sgd --data_splits 0.7 0.2
model:
  name: mlp
  depth: 5
  width: 256
optimizer:
  name: sgd
  momentum: 0.0
  learning_rate: 0.001
  weight_decay: 0.0
  nesterov: false
dataset: mnist
data_splits:
- 0.7
- 0.2
n_epochs: 1024
n_steps_per_epoch: 256
tasks: []

$ python train.py --help
usage: train.py [-h] [--dry] [--model.name str] [--model.depth int] [--model.width int]
                [--optimizer {adam,sgd}] [--dataset str] [--data_splits float [float ...]]
                [--n_epochs int] [--n_steps_per_epoch int] [--tasks [str ...]]
                [--optimizer.name {adam}] [--optimizer.betas float float]
                [--optimizer.learning_rate float] [--optimizer.weight_decay float]

optional arguments:
  -h, --help                       show this help message and exit
  --dry                            dry run (default: False)

  --optimizer {adam,sgd}           (default: adam)
  --dataset str                    (default: mnist)
  --data_splits float [float ...]  (default: (0.8, 0.1))
  --n_epochs int                   (default: 1024)
  --n_steps_per_epoch int          (default: 256)
  --tasks [str ...]                (default: [])

model:
  --model.name str                 (default: mlp)
  --model.depth int                (default: 3)
  --model.width int                (default: 256)

optimizer:
  --optimizer.name {adam}          (default: adam)
  --optimizer.betas float float    (default: (0.95, 0.95))
  --optimizer.learning_rate float  (default: 0.001)
  --optimizer.weight_decay float   (default: 0.0)