Easy Hyperparameter Management with Hydra, MLflow, and Optuna
Two major methods can be considered for hyperparameter management in machine learning.
- Configuring hyperparameters from the command line using argparse
- Creating a file with a list of hyperparameters for every experiment
An Example of a Typical Hyperparameter Management
When using argparse for managing hyperparameters, it is convenient to change them directly from the command line, but the number of hyperparameters to be configured may be huge. Moreover, when managing hyperparameters by writing them in a configuration file, the configuration file has to be modified every time a hyperparameter is changed.
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
Hyperparameter management using Hydra+MLflow+Optuna allows users to modify and execute the configured hyperparameters without directly editing the configuration files from the command line. Hydra can perform a grid search of variables from the command line, and there will be no need to worry about the continuous increase of configuration files. Moreover, the Optuna plug-in for Hydra allows us to harness Optuna’s powerful hyperparameter search function from the command line, which makes hyperparameter search and management more effective. In this article, hyperparameter management using Hydra+MLflow will be presented, and the introduction of Optuna using the Hydra plug-in will also be illustrated.
Hydra
Hydra is a hyperparameter management tool released by Facebook AI Research, which allows users to define hyperparameters using YAML files in a hierarchically structured order. Also, it is possible to change and execute the configuration values directly from the command line and to perform a grid search on hyperparameters with a single command.
Basic Usage
- Write the hyperparameters that need to be managed in a config file in YAML format.
- Reference the specified hyperparameters inside the program by passing a decorator to the function.
Changing and Executing Values from the Command Line
In Hydra, when you want to adjust the value of a hyperparameter and re-execute the program, you can change the value of the hyperparameter directly from the command line and execute the program. In the preceding example, if you wish to change the node1 value from 128 to 64 and to re-execute the program, you can directly specify the value on the command line and change and execute the program.
Hyperparameter Grid Search
If you wish to execute several hyperparameters in order, you can specify them directly from the command line without rewriting the config file. In the previous example, if you want to change node1 at {128, 256} and node2 at {16, 32}, you can specify them from the command line as shown below. The command will be implemented according to the number of hyperparameter combinations that have been specified.
MLflow
MLflow is an open source platform that simplifies the machine learning lifecycle. MLflow, when combined with Hydra, enables management, grid search, storage, and comparison of hyperparameters. In this article, I will introduc
MLflow Tracking
MLflow Tracking provides logging of hyperparameters when running machine learning code, metrics such as loss and accuracy, and an API to assist managing output files and the like. MLflow Tracking can be installed via pip.
Basic Usage
The logging functions that are provided by MLflow enable you to record hyperparameter metrics. More concretely, you can issue a runID with start_run(), register hyperparameters with log_param(), record metrics with log_metric(), and record the output file and others with log_artifact().
You can use the following command to compare the hyperparameters that have been recorded on the local server.
Code Example of Hydra and MLflow
With the combined use of Hydra and MLflow, it is possible to manage and modify hyperparameters without significantly altering the original file. The following code demonstrates an example of implementing Hydra and MLflow in a machine learning model written in PyTorch.
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.
Integration with Optuna
In this article, I have presented the hyperparameter management and grid search using Hydra+MLflow and the comparison of metrics for all the hyperparameter combinations. Optuna can realize not only the grid search of hyperparameters by Hydra but also the optimization of hyperparameters. In addition, the use of the Hydra plug-in makes using Optuna significantly easier. From here, the combined usage of Hydra+MLflow+Optuna will be explained with concrete examples.
What is Optuna?
Optuna is an open source hyperparameter automatic optimization framework. Optuna automates the trial-and-error process of tuning hyperparameter values for best performance.
Hydra+MLflow+Optuna
With the use of Optuna Sweeper plug-in for Hydra, hyperparameter optimization can be performed by simply describing the variables and search conditions to be optimized in the configuration file used in Hydra. Furthermore, it allows the variables and search conditions to be optimized directly from the command line. The Optuna Sweeper plug-in can be installed by pip.
Modifications for using Optuna
The Optuna configuration can be completed by writing the Optuna settings in the configuration file used by the Hydra decorator as illustrated below.
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
As a matter of course, the values of the hyperparameters and the objective function searched by Optuna can be visualized on a local server built with MLflow.
Summary
Hydra+MLflow
By combining Hydra for managing hyperparameters and doing grid search with MLflow for recording hyperparameters and visualizing results, hyperparameters can be easily changed, managed, and compared on the GUI each time the experiment is run.
Hydra+MLflow+Optuna!
Adding Optuna to the above combination, expands the hyperparameter tuning from only grid search to more detailed hyperparameter search algorithms. The searched hyperparameters and the objective variables values can be directly managed and compared by MLflow.
