Optuna¶

Autoware-ML integrates Optuna for automated hyperparameter optimization. Optuna intelligently searches the hyperparameter space, finding optimal configurations faster than grid search or random search.

How It Works¶

Optuna sits above Hydra in the configuration hierarchy:

Optuna (suggests hyperparameters)
    ↓
Hydra (configures training)
    ↓
Lightning (runs training)
    ↓
MLflow (logs results)

Each trial:

Optuna suggests hyperparameter values
Hydra builds the config with those values
Training runs and reports the objective metric
Optuna uses the result to inform the next trial

Running a Hyperparameter Search¶

autoware-ml train --config-name my_task/my_model \
    --multirun \
    hydra/sweeper=optuna \
    hydra.sweeper.n_trials=50 \
    hydra.sweeper.direction=minimize

This launches 50 trials, minimizing the objective metric (default: validation loss).

Defining Search Spaces¶

Define hyperparameter ranges in your config:

configs/my_task/my_model_optuna.yaml

# @package _global_
defaults:
  - /my_task/my_model_base
  - _self_

hydra:
  sweeper:
    params:
      model.optimizer.lr: interval(0.0001, 0.01)
      model.optimizer.weight_decay: interval(0.001, 0.1)
      datamodule.train_dataloader_cfg.batch_size: choice(2, 4, 8, 16)
      trainer.max_epochs: range(10, 50, step=10)

Search Space Types¶

Type	Syntax	Example
Continuous	`interval(low, high)`	`interval(0.0001, 0.01)`
Log-scale	`interval(low, high, log=true)`	`interval(1e-5, 1e-2, log=true)`
Integer range	`range(start, end, step)`	`range(10, 100, step=10)`
Categorical	`choice(a, b, c)`	`choice(adam, sgd, adamw)`

Log-Scale for Learning Rates¶

Learning rates are typically searched on a log scale:

hydra:
  sweeper:
    params:
      model.optimizer.lr: interval(1e-5, 1e-2, log=true)

Configuring the Sweeper¶

Configure the Optuna sweeper in your config:

hydra:
  sweeper:
    _target_: hydra_plugins.hydra_optuna_sweeper.OptunaSweeper
    study_name: my_optimization
    direction: minimize
    n_trials: 100

    sampler:
      _target_: optuna.samplers.TPESampler
      seed: 42

Common samplers: TPESampler (default), CmaEsSampler, RandomSampler, GridSampler.

Objective Metrics¶

By default, Optuna optimizes val/loss. To optimize a different metric:

hydra:
  sweeper:
    direction: maximize

Parallel Trials¶

Run multiple trials in parallel:

autoware-ml train --config-name my_task/my_model_optuna \
    --multirun \
    hydra.sweeper.n_jobs=4

GPU Considerations

Each parallel trial needs GPU memory. Ensure sufficient VRAM or use multiple GPUs.

Multiple workers can connect to the same study.

Viewing Results¶

MLflow Integration¶

All trials are logged to MLflow. Use the MLflow UI to compare trials as regular training runs.

Best Practices¶

Start with wide ranges, then narrow based on results
Use pruning to avoid wasting compute on bad configurations
Fix unimportant parameters to reduce search space

Example: Full Optimization Config¶

configs/my_task/my_model_optuna.yaml

# @package _global_
defaults:
  - /my_task/my_model_base
  - override /hydra/sweeper: optuna
  - _self_

hydra:
  sweeper:
    _target_: hydra_plugins.hydra_optuna_sweeper.OptunaSweeper
    study_name: my_optimization
    direction: minimize
    n_trials: 100

    sampler:
      _target_: optuna.samplers.TPESampler
      seed: 42

    params:
      model.optimizer.lr: interval(1e-5, 1e-2, log=true)
      model.optimizer.weight_decay: interval(1e-4, 1e-1, log=true)
      datamodule.train_dataloader_cfg.batch_size: choice(2, 4, 8)

trainer:
  max_epochs: 20  # Shorter for faster trials

Run it:

autoware-ml train --config-name my_task/my_model_optuna --multirun

Learn More¶

For detailed Optuna documentation, see the official Optuna documentation.