Open In Colab to execute or view/download this notebook on GitHub

Hyperparameter Optimization

Many of the speech processing tasks implemented as part of the SpeechBrain project rely on the careful selection of hyperparameters, such as:

  • The number of layers

  • Normalization

  • Hidden layer dimensions

  • Weights within cost functions

  • etc

Selecting such hyperparameters by hand can be tedious. This tutorial will show how to use the automated hyperparameter optimization techniques implemented as part of the Oríon project to automatically fit hyperparameters in a systematic way.

Prerequisites

Imports

import os

Install SpeechBrain

SpeechBrain can be downloaded from the GitHub repository listed below.

%%capture
# Installing SpeechBrain via pip
BRANCH = 'develop'
!python -m pip install git+https://github.com/speechbrain/speechbrain.git@$BRANCH

Dependency Fixes

PyYAML 6.0 is not backwards-compatible, a 5.x version is needed to support HyperPyYAML

%%capture
!pip install pyyaml==5.4.1

Install Oríon

Oríon can be installed using pip or conda

%%capture
!pip install orion[profet]

from speechbrain.utils import hpopt as hp

Update the Recipe to Support Hyperparameter Optimization

SpeechBrain comes with a convenience wrapper called hpopt, which is capable of reporting objective values to Orion or to other tools.

For a complete example on how to implement it,

  1. Add the following import statement to the top of your recipe:

from speechbrain.utils import hpopt as hp

  1. Wrap the main code of your recipe in a hyperparameter optimization context. Set objective_key to the metric that Orion will optimize.

    Before:

hparams_file, run_opts, overrides = sb.parse_arguments(sys.argv[1:])

with open(hparams_file) as fin:
    hparams = load_hyperpyyaml(fin, overrides)

## ...

spk_id_brain = SpkIdBrain(
    modules=hparams["modules"],
    opt_class=hparams["opt_class"],
    hparams=hparams,
    run_opts=run_opts,
    checkpointer=hparams["checkpointer"],
)

# The `fit()` method iterates the training loop, calling the methods
# necessary to update the parameters of the model. Since all objects
# with changing state are managed by the Checkpointer, training can be
# stopped at any point, and will be resumed on next call.
spk_id_brain.fit(
    epoch_counter=spk_id_brain.hparams.epoch_counter,
    train_set=datasets["train"],
    valid_set=datasets["valid"],
    train_loader_kwargs=hparams["dataloader_options"],
    valid_loader_kwargs=hparams["dataloader_options"],
)

After:

```python
with hp.hyperparameter_optimization(objective_key="error") as hp_ctx: # <-- Initialize the context
    hparams_file, run_opts, overrides = hp_ctx.parse_arguments(sys.argv[1:]) # <-- Replace sb with hp_ctx

    with open(hparams_file) as fin:
        hparams = load_hyperpyyaml(fin, overrides)

    ## ...

        spk_id_brain = SpkIdBrain(
            modules=hparams["modules"],
            opt_class=hparams["opt_class"],
            hparams=hparams,
            run_opts=run_opts,
            checkpointer=hparams["checkpointer"],
        )

        # The `fit()` method iterates the training loop, calling the methods
        # necessary to update the parameters of the model. Since all objects
        # with changing state are managed by the Checkpointer, training can be
        # stopped at any point, and will be resumed on next call.
        spk_id_brain.fit(
            epoch_counter=spk_id_brain.hparams.epoch_counter,
            train_set=datasets["train"],
            valid_set=datasets["valid"],
            train_loader_kwargs=hparams["dataloader_options"],
            valid_loader_kwargs=hparams["dataloader_options"],
        )
```

  1. Add code to report the stats

e.g. in on_stage_end when stage == sb.Stage.VALID

hp.report_result(stage_stats)

The last result reported through this function will be reported for hyperparameter optimization.

The key specified in objective_key parameter needs to be present in the dictionary passed to report_result.

  1. Add the following lines in your main hyperparameter file train.yaml:

hpopt_mode: null
hpopt: null

  1. Optional: Create a separate YAML file overriding any hyperparameters to be used during hyperparameter optimization that are different from the ones used during regular training other than the ones being fitted. A typical approach would reduce the number of epochs and the number of training samples.

This step can be omitted if the number of parameters being overridden is small. In this case, they can be passed on the command line instead.

Example:

hpopt.yaml:

number_of_epochs: 1
ckpt_enable: false

  1. ❗ Important: Most recipes use a checkpointer to save snapshots of the model after each epoch (or on a custom schedule) to ensure that training can be resumed if it is interrupted. During hyperparameter optimization, this can cause issues because if the model’s architecture (e.g. the number of layers, neurons per layer, etc) changes from one set of hyperparamter values to the next, an attempt to restore a checkpoint will fail.

One possible solution is to make the run of the checkpointer conditional and to disable it in hpopt.yaml

Before:

self.checkpointer.save_and_keep_only(meta=stats, min_keys=["error"])

After:

if self.hparams.ckpt_enable:
    self.checkpointer.save_and_keep_only(meta=stats, min_keys=["error"])

An alternative strategy is to reconfigure the checkpointer to save each run in a separate directory. For this scenario, the hyperparameter optimization wrapper can supply a variable named trial_id, which can be interpolated into the output path.

Given below is an example of this strategy:

hpopt.yaml:

number_of_epochs: 1
ckpt_enable: False
trial_id: hpopt
output_folder: !ref ./results/speaker_id/<trial_id>

train.yaml:

# ...
save_folder: !ref <output_folder>/save
# ...
checkpointer: !new:speechbrain.utils.checkpoints.Checkpointer
    checkpoints_dir: !ref <save_folder> #<-- will contain trial_id
    recoverables:
        embedding_model: !ref <embedding_model>
        classifier: !ref <classifier>
        normalizer: !ref <mean_var_norm>
        counter: !ref <epoch_counter>

Inspecting Results

Use the orion info command to inspect the results of hyperparameter fitting.

The tool will output basic statistics about the hyperparameter fitting experiment, including the number of runs completed, the objective value for the best trial and the hyperparameter values corresponding to that run.

In the example below, the best objective achieved value is shown under evaluation, and the corresponding hyperparameter values are shown under params.

Stats
=====
completed: False
trials completed: 4
best trial:
  id: c1a71e0988d70005302ab655d7e391d3
  evaluation: 0.2384105920791626
  params:
    /emb_dim: 128
    /tdnn_channels: 128
start time: 2021-11-14 21:01:12.760704
finish time: 2021-11-14 21:13:25.043336
duration: 0:12:12.282632

!orion info --name speaker-id

Hyperparameter Optimization at Scale

Multiple GPUs

Since Orion simply wraps the execution of the training script and launches it for each set of hyperparameters using the OS shell, training scripts that support Data-Parallel (DP) or Distributed Data Parallel (DDP) execution can be used with hyperparameter fitting without modification.

For information on how to set up DP/DDP experiments, refer to the SpeechBrain documentation and the Multi-GPU Considerations tutorial.

Parallel or Distributed Oríon

Oríon itself provide support for parallel and distributed hyperparameter fitting.

To use multiple parallel workers on a single node, pass the --n-workers parameter to the Oríon CLI.

The example below will start the experiment with three workers:

orion hunt -n $HPOPT_EXPERIMENT_NAME -c $HOPT_CONFIG_FILE --n-workers 3 python train.py hparams/$HPARAMS \
    --hpopt hpopt.yaml \
    --hpopt_mode orion \
    --emb_dim~"choices([128,256,512,768,1024])" \
    --tdnn_channels~"choices([128,256,512,768,1024])"

For more advanced scenarios, including distributed hyperparameter fittig on multiple nodes, refer to the Parallel Workers page in Oríon’s official documentation.

Citing SpeechBrain

If you use SpeechBrain in your research or business, please cite it using the following BibTeX entry:

@misc{speechbrainV1,
  title={Open-Source Conversational AI with {SpeechBrain} 1.0},
  author={Mirco Ravanelli and Titouan Parcollet and Adel Moumen and Sylvain de Langen and Cem Subakan and Peter Plantinga and Yingzhi Wang and Pooneh Mousavi and Luca Della Libera and Artem Ploujnikov and Francesco Paissan and Davide Borra and Salah Zaiem and Zeyu Zhao and Shucong Zhang and Georgios Karakasidis and Sung-Lin Yeh and Pierre Champion and Aku Rouhe and Rudolf Braun and Florian Mai and Juan Zuluaga-Gomez and Seyed Mahed Mousavi and Andreas Nautsch and Xuechen Liu and Sangeet Sagar and Jarod Duret and Salima Mdhaffar and Gaelle Laperriere and Mickael Rouvier and Renato De Mori and Yannick Esteve},
  year={2024},
  eprint={2407.00463},
  archivePrefix={arXiv},
  primaryClass={cs.LG},
  url={https://arxiv.org/abs/2407.00463},
}
@misc{speechbrain,
  title={{SpeechBrain}: A General-Purpose Speech Toolkit},
  author={Mirco Ravanelli and Titouan Parcollet and Peter Plantinga and Aku Rouhe and Samuele Cornell and Loren Lugosch and Cem Subakan and Nauman Dawalatabad and Abdelwahab Heba and Jianyuan Zhong and Ju-Chieh Chou and Sung-Lin Yeh and Szu-Wei Fu and Chien-Feng Liao and Elena Rastorgueva and François Grondin and William Aris and Hwidong Na and Yan Gao and Renato De Mori and Yoshua Bengio},
  year={2021},
  eprint={2106.04624},
  archivePrefix={arXiv},
  primaryClass={eess.AS},
  note={arXiv:2106.04624}
}