Configuration

Anemoi Training uses Hydra for configuration management, allowing for flexible and modular configuration of the training pipeline. This guide explains how to use Hydra effectively in the project.

Hydra Basics

Hydra is a framework for elegantly configuring complex applications. It allows for:

  1. Hierarchical configuration

  2. Configuration composition

  3. Dynamic object instantiation

Object Instantiation with Hydra

Hydra provides powerful tools for instantiating objects directly from configuration files:

  • hydra.utils.instantiate(): Creates object instances

  • hydra.utils.call(): Calls functions with configured parameters

Example: Instantiating an Optimizer

Consider the following Python class:

class Optimizer:
    def __init__(self, algorithm: str, learning_rate: float) -> None:
        self.opt_algorithm = algorithm
        self.lr = learning_rate

Configuration in YAML:

optimizer:
  _target_: my_code.Optimizer
  algorithm: SGD
  learning_rate: 0.01

Instantiating in code:

from hydra.utils import instantiate

optimizer = instantiate(config.optimizer)

Configurable Components in Anemoi Training

Anemoi Training uses Hydra’s instantiation feature for various components, including:

  1. Model architectures

  2. Pressure level scalers

  3. Graph definitions

And there are plans to extend these to other areas, such as:

  1. Loss functions

  2. Callbacks

  3. Data loaders

Example: Configuring a Pressure Level Scaler

In config.training.pressure_level_scaler, users can define custom scaling behavior:

pressure_level_scaler:
    _target_: anemoi.training.losses.scalers.ReLUPressureLevelScaler
    min_weight: 0.2

Best Practices for Hydra Configuration

  1. Use configuration groups for logically related settings.

  2. Leverage Hydra’s composition feature to combine configurations.

  3. Use interpolation to reduce redundancy in configurations.

  4. Provide default values for all configurable parameters.

  5. Use type hints in your classes to ensure correct instantiation.

Advanced Hydra Features

1. Config Groups

Organize related configurations into groups for easier management and overriding.

2. Multi-run

Hydra supports running multiple configurations in a single execution:

python train.py --multirun optimizer.learning_rate=0.001,0.01,0.1

3. Sweeps

Define parameter sweeps for hyperparameter tuning, a powerful feature, but usually only required when the model development is relatively mature:

# config.yaml
defaults:
  - override hydra/sweeper: optuna

hydra:
  sweeper:
    sampler:
      _target_: optuna.samplers.TPESampler
    direction: minimize
    n_trials: 20
    params:
      optimizer.learning_rate: range(0.0001, 0.1, log=true)

Run the sweep:

python train.py --multirun

By leveraging these Hydra features, you can create flexible, maintainable, and powerful configurations for Anemoi Training.