An annotated example of using Hydra

This code base in meant to show-case how to use Hydra to manage configuration in a project.

Installing dependencies

If you have uv installed:

uv sync

If not, run (after activating an appropriate python env):

pip install -r pylock.toml

Running the code

To run the code with default configs, simply do:

python main.py

The output should be something like this:

Starting a run with seed 42 on CPU.
Model architecture:
Sequential(
  (0): Sequential(
    (0): Linear(in_features=28, out_features=28, bias=True)
    (1): LayerNorm((28,), eps=1e-05, elementwise_affine=True)
    (2): GELU(approximate='none')
  )
  (1): Linear(in_features=28, out_features=10, bias=True)
)

We can, for example, set the seed with:

python main.py seed=123

To find out which arguments are available, run:

python main.py --help

It should look like this:

Help output (quite long)

main is powered by Hydra.

== Configuration groups ==
Compose your configuration from those groups (group=option)

dm: celeba, celeba_male_blond, celeba_male_smiling, celeba_male_smiling_small, cmnist
model: cnn, fcn, single_linear_layer
opt: weight_decay_adamw


== Config ==
Override anything in the config (foo.bar=value)

dm:
  root: !!python/object/apply:pathlib.PosixPath
  - /srv/galene0/shared/data
  default_res: 28
  label_map: null
  colors: null
  num_colors: 10
  val_prop: 0.2
model:
  num_hidden: 1
  hidden_dim: null
  activation: gelu
  norm: LN
  dropout_prob: 0.0
  final_bias: true
  input_norm: false
opt:
  lr: 5.0e-05
  weight_decay: 0.0
  optimizer_cls: torch.optim.AdamW
  optimizer_kwargs: null
wandb:
  name: null
  dir: ./local_logging
  id: null
  project: hydra-example
  group: null
  entity: predictive-analytics-lab
  tags: null
  job_type: null
  mode: disabled
  resume: null
seed: 42
gpu: -1

We can see that seed and gpu are "top-level" config values, but there are also a lot of values in subconfigs. For example, we can change the num_hidden value in the model config:

python main.py model.num_hidden=2

This should print a different model architecture.

But we can also see that there are different "schemas" available in the "model" group:

model: cnn, fcn, single_linear_layer

This allows us to completely change the model architecture:

python main.py model=cnn

With the "cnn" architecture, the config key model.hidden_dim is not available anymore, so the following should raise an error:

python main.py model=cnn model.hidden_dim=10

Output:

Could not override 'model.hidden_dim'.
To append to your config use +model.hidden_dim=10
Key 'hidden_dim' not in 'SimpleCNNFactory'
    full_key: model.hidden_dim
    reference_type=ModelFactory
    object_type=SimpleCNNFactory

To see the available options for the "cnn" architecture, you can run:

python main.py model=cnn --help

Now the "model" part of the output should look like this:

model:
  kernel_size: 5
  pool_stride: 2
  activation: relu

This tells us that we can set the kernel_size, pool_stride, and activation values for the "cnn" architecture. For example:

python main.py model=cnn model.kernel_size=3 model.pool_stride=1 model.activation=selu

W&B config

To enable W&B logging, you can set the wandb.mode to "online" or "offline" (default is "disabled"):

python main.py wandb.mode=offline

Basics of config files

Let's say you are often working with a configuration of the CelebA dataset which looks like this:

python main.py dm=celeba dm.superclass=Smiling dm.subclass=Male

(Once again, you can set dm=celeba and use --help to see the available options for the CelebA dataset.)

But after a while, it gets tiring to always type this out. Instead, you can create a config file in the configs/dm/ directory to store often-used datamodule configurations. For example, you can create a file configs/dm/celeba_male_smiling.yaml with the following content:

defaults:
  - celeba

superclass: Smiling
subclass: Male

You can think of the defaults entry as specifying inheritance.

Subsequently, you can use the config file like this:

python main.py dm=celeba_male_smiling

You can still override any values on the command line, just as before:

python main.py dm=celeba_male_smiling dm.superclass=Blond_Hair dm.default_res=64

And you can create config files which inherit from other config files you created. For example, you can create a file configs/dm/celeba_male_smiling_small.yaml with the following content:

defaults:
  - celeba_male_smiling

default_res: 64

Multiruns

Often, we want to run the same experiment with only slightly different configurations. This can be done with the --multirun flag. For example, to run the code with seeds 42 and 43, we can do:

python main.py --multirun seed=42,43

The output should look like this:

[2024-05-16 18:40:47,580][HYDRA] Launching 2 jobs locally
[2024-05-16 18:40:47,581][HYDRA]        #0 : seed=42
Starting a run with seed 42 on CPU.
Model architecture:
Sequential(
  (0): Sequential(
    (0): Linear(in_features=28, out_features=28, bias=True)
    (1): LayerNorm((28,), eps=1e-05, elementwise_affine=True)
    (2): GELU(approximate='none')
  )
  (1): Linear(in_features=28, out_features=10, bias=True)
)
[2024-05-16 18:40:47,672][HYDRA]        #1 : seed=43
Starting a run with seed 43 on CPU.
Model architecture:
Sequential(
  (0): Sequential(
    (0): Linear(in_features=28, out_features=28, bias=True)
    (1): LayerNorm((28,), eps=1e-05, elementwise_affine=True)
    (2): GELU(approximate='none')
  )
  (1): Linear(in_features=28, out_features=10, bias=True)
)

By default, Hydra simply runs the code with the different config values sequentially. But if you have access to a SLURM cluster, you can also run the jobs in parallel. To do this, we need to set hydra/launcher=... to the appropriate value. In this code base, there is a config file at configs/hydra/launcher/slurm_kyiv.yaml that has the correct settings for the Kyiv machine. To use it, we can run:

python main.py --multirun seed=42,43 hydra/launcher=slurm_kyiv

Hydra doesn't limit you to iterating over just one parameter. You can also iterate over multiple parameters. For example, to run the code with seeds 42 and 43 and model.num_hidden set to 1, 2 and 3, you can do:

python main.py --multirun seed=42,43 model.hidden_dim=10 model.num_hidden=1,2,3

This will start 6 jobs: 2 seeds x 3 model.num_hidden values.

Hyperparameter search

If you don't want to just blindly try out model parameters, you can also use Hydra to perform hyperparameter search with libraries like Optuna, Nevergrad, and Ax.

In order to use any of these libraries, you need to return a value from the function decorated with @hydra.main(...) (located in main.py). The value should be a float and should represent the metric you want to optimize. For example, if you want to optimize the accuracy of the model, you can return the validation accuracy. If you want to minimize the loss, you can return the validation loss.

To use Optuna, for example, create a config file in configs/hydra/sweeper/. The name doesn't matter, but we'll choose fcn_params.yaml and fill the content like this:

defaults:
  # Select here which hyperparameter optimization library to use.
  - optuna

sampler:
  warn_independent_sampling: true

study_name: fcn_params
n_trials: 25
n_jobs: 3
direction:
  # The direction depends on what you are returning from your main() function.
  - maximize

params:
  opt.lr: tag(log, interval(1.e-5, 1.e-3))
  model.num_hidden: range(1, 4)
  model.dropout_prob: choice(0.0, 0.1)

Let's break down what's going on here:

• First, we select the right defaults value. In this case, we want to use Optuna, so we set optuna.
• Then, some settings for Optuna follow. To see all available settings, check out the docs: https://hydra.cc/docs/plugins/optuna_sweeper/
• Finally, we specify params, which is the heart of this config file. Here, we specify the parameters we want to optimize. For each parameter, we have to specify the search space. For example, opt.lr: tag(log, interval(1.e-5, 1.e-3)) means that we want to optimize the opt.lr parameter and we want to search in the log space between 1.e-5 and 1.e-3. The model.num_hidden: range(1, 4) we want to search in the range from 1 to 3 (inclusive). The model.dropout_prob: choice(0.0, 0.1) means that we want to choose between 0.0 and 0.1.

Global experiment configs

We already saw that you can create config files to, for example, save a common data module configuration – such a config file is then put in configs/dm/. However, such config files can only configure one subcomponent of the whole configuration. This means you might end up with long commands like this one:

python main.py model=single_linear_layer model.activation=selu dm=celeba_male_smiling_small dm.download=false seed=1 opt.lr=0.001 opt.weight_decay=0.001

which you have to type over and over again.

Luckily, we can define a "primary config" in the top-level configs/ directory, that act as global configurations. For the above command, we can create a file at configs/good_run.yaml with the following content:

---
defaults:
  - base_config
  - model: single_linear_layer
  - dm: celeba_male_smiling_small

seed: 1

model:
  activation: selu

dm:
  download: false

opt:
  lr: 0.001
  weight_decay: 0.001

Note that we need to first inherit from base_config so that hydra knows which configuration fields are allowed. The name “base_config” is arbitrary but it needs to match the name that is passed to register_hydra_config() in main.py.

And then we can run the code with these config values by running:

python main.py --config-name good_run

Instead of --config-name you can also use the shorthand -cn.

You can still override values:

python main.py -cn good_run model.final_bias=false

Note: It's very easy to fall into the trap of defining everything in these global configs, but that leads to lots of duplication. Try to put as much configuration as possible into the component-specific config files (i.e., those in configs/dm/ and configs/model and so on), because those configs are easy to reuse across different experiments.

Config validation

When you declare a field as, say, gpu: int in the dataclass, hydra will ensure that the value given for gpu is an int. However, often we need to validate configuration values in ways that cannot be expressed as a type. For example, in the code here we require that gpu is either -1, which indicates that no GPU is used, or is equal to a valid GPU index for the current machine. To validate such constraints on the configuration value, we can use __post_init__() in the following way.

@dataclass(kw_only=True)
class Config:
    """Main configuration class for the code base."""

    # ...all the other fields...
    gpu: int = 0  # Set to -1 to use CPU.

    def __post_init__(self) -> None:
        """Post-initialization checks."""
        # Check that the GPU index is valid.
        if self.gpu < -1 or self.gpu >= torch.cuda.device_count():
            raise ValueError(
                f"Invalid GPU index {self.gpu}. Must be between -1 and "
                f"{torch.cuda.device_count() - 1}."
            )

The __post_init__ is run directly after hydra has instantiated the configuration objects. You can try it out by passing an illegal value:

python main.py gpu=-2

The output should look something like this (depending on how many GPUs are available):

Error executing job with overrides: ['gpu=-2']
Traceback (most recent call last):
  (...traceback details...)
ValueError: Invalid GPU index -2. Must be between -1 and 3.

How to make your code easily configurable with Hydra

What we found is the best method to structure your code to make it easy to configure with Hydra is the "builder" or "factory" pattern.

This means you have a dataclass that contains all the configuration values for a particular component of your code (e.g. the model, the data, the optimiser, etc.). And then this class has a build() or init() method that takes additional arguments which are only available at runtime (e.g. the input size of the model, the number of classes in the dataset, etc.), and then instantiates the component.

For example, say your code has such a function for constructing a convolutional neural network:

def build_simple_cnn(
    in_dim: int,
    out_dim: int,
    kernel_size: int = 5,
    pool_stride: int = 2,
    activation: Activation = "relu",
) -> nn.Sequential:
    return nn.Sequential(...)

Then we can pull the configurable parameters into a dataclass and re-write it like this

@dataclass(kw_only=True)
class SimpleCNNFactory(ModelFactory):
    """Factory for a very simple CNN."""

    kernel_size: int = 5
    pool_stride: int = 2
    activation: Activation = "relu"

    @override
    def build(self, in_dim: int, *, out_dim: int) -> nn.Sequential:
        return nn.Sequential(...)  # use `self.kernel_size` etc.

In this code, we also have an abstract base class, which allows us to provide a better type annotation for hydra:

@dataclass  # this needs to be a dataclass even though it has no fields
class ModelFactory(ABC):
    """Interface for model factories."""

    @abstractmethod
    def build(self, in_dim: int, *, out_dim: int) -> nn.Module:
        raise NotImplementedError()

As you can see, SimpleCNNFactory above inherits from this.

Using the current date in config values

Hydra comes with a built-in now resolver which formats the current date and time with a strftime pattern. This is useful, for example, for grouping W&B runs by the day they were started:

python main.py 'wandb.group=${now:%Y-%m-%d}-cnn-dropout'

This sets wandb.group to something like 2026-10-16-cnn-dropout. Note the single quotes: without them, the shell would try to expand ${now:...} itself.

You can use any strftime pattern, e.g. ${now:%Y-%m-%d_%H-%M} to also include the time. The resolver also works in config files:

wandb:
  group: ${now:%Y-%m-%d}-cnn-dropout

Note: within a single run, the value is cached, so it doesn't change. But in a multirun, each job resolves it separately, so if your sweep runs past midnight, the later jobs will get the next day's date.

Structure of the code

• main.py: The main entry point of the code. It sets up Hydra and then calls the main run() function.
• src/
  • run.py: Contains the main Config class that is used to define valid config values. It also contains the run() function that is called by main.py.
  • datasets.py: Contains the DataModule class that is used to load the data.
  • models.py: Contains the ModelFactory class that is used to create the model.
  • optimisation.py: Contains the OptimisationCfg class that is used to build the optimiser that trains the model.
  • logging.py: Contains the WandbCfg class that is used to set up Weights & Biases logging.
• configs/
  • default_config.yaml: The default primary config file for the project. It sets the default values for dm and model.
  • good_run.yaml: An alternative primary config.
  • hydra/
    • launcher/: Contains the SLURM launcher config files.
    • sweeper/: Contains the Optuna sweeper config files.
  • dm/: Contains the config files for the different datasets.
  • model/: Contains the config files for the different model architectures.
  • opt/: Contains the config files for the optimiser.
• tests/test_configs.py: Checks that all config files in configs/ are valid, without running any experiments (run with pytest).