Easy Hyperparameter Management with Hydra, MLflow, and Optuna

Mar 31 · 6 min read

Two major methods can be considered for hyperparameter management in machine learning.

  • Configuring hyperparameters from the command line using argparse
Hyperparameter management using argparse
  • Creating a file with a list of hyperparameters for every experiment
Hyperparameter management via configuration files

An Example of a Typical Hyperparameter Management

When determining effective hyperparameter values, the process of trial-and-error with multiple hyperparameters is not only cumbersome, but also causes modifications to the configuration file according to the number of potential candidate hyperparameters, making it difficult to save and compare a vast quantity of results.

These drawbacks can be solved with Hydra, MLflow(tm), and Optuna.

Hyperparameter Management using Hydra+MLflow+Optuna


Basic Usage

  • Reference the specified hyperparameters inside the program by passing a decorator to the function.

Changing and Executing Values from the Command Line

Hyperparameter Grid Search


MLflow Tracking

Basic Usage

You can use the following command to compare the hyperparameters that have been recorded on the local server.

Code Example of Hydra and MLflow

Similarly, it is possible to specify multiple hyperparameters directly on the command line to perform a grid search.

After executing the above command, the following screen will be displayed on the local server. The results of the pre-registered metrics, such as loss and accuracy, can be easily compared for all hyperparameter combinations.

Values recorded by MLflow can be checked on the GUI.

Integration with Optuna

What is Optuna?


Modifications for using Optuna

In addition, the return value of the function that is decorated is the objective variable to be minimized or maximized. The following code is an example where the return value of a function is the accuracy to be maximized.

Only the above changes are required to use Optuna. Following these modifications, the variables to be optimized and their search ranges can be specified directly from the command line. The following code demonstrates an example of the learning rate of the optimizer and the number of nodes of the model that maximize the accuracy along with the search conditions from the command line. Choices are converted to categorical variables, thus the four learning rates of the optimizer can be searched. In the following example, model.node1 is searched within the range of [10, 500] as an integer. Though various other distributions are also compatible with this function, they are not discussed here.

Visualization on MLflow

Examples of hyperparameters searched by Optuna





A hyperparameter optimization framework

Medium is an open platform where 170 million readers come to find insightful and dynamic thinking. Here, expert and undiscovered voices alike dive into the heart of any topic and bring new ideas to the surface. Learn more

Follow the writers, publications, and topics that matter to you, and you’ll see them on your homepage and in your inbox. Explore

If you have a story to tell, knowledge to share, or a perspective to offer — welcome home. It’s easy and free to post your thinking on any topic. Write on Medium

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store