Getting started with Optuna

If you are looking to get started with Hyperparameter tuning with tools like Ray Tune, Huggingface HPO or the like, one option that is available is Optuna.

The demo video on the home page is a nice introduction, but since that is from 2019, it is a bit out of date. Here are the cleaned up instructions.

First install Optuna and the dashboard (which is a nice reason to look at it in comparision to its competitors):

$ pip install optuna optuna-dashboard

The basic example on the webpage is nice, but the one in the video is a bit more non-trivial. The code, without any Optuna integration looks like the following:

from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.neural_network import MLPClassifier

mnist = load_digits()
x_train, x_test, y_train, y_test = train_test_split(mnist.data, mnist.target)

layers = [128, 64, 64, 128]
clf = MLPClassifier(hidden_layer_sizes=tuple(layers))
clf.fit(x_train, y_train)

print(clf.score(x_test, y_test))

This is the same thing from the video, but updated to todays convention of sklearn. If you run this, it will train a small 4 layer Neural network in a few seconds and report the test accuracy to be something like …

Now, let us update the same code to use Optuna to search for the paramters, in this case the number of layers and the number of perceptrons in each layer. The updated code is below, reproduced fully.

from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.neural_network import MLPClassifier

import optuna

mnist = load_digits()
x_train, x_test, y_train, y_test = train_test_split(mnist.data, mnist.target)

def train(trial):
    layers = []

    n_layers = trial.suggest_int('n_layers', 1, 4)
    for i in range(n_layers):
        layers.append(trial.suggest_int(str(i), 1, 128))

    clf = MLPClassifier(hidden_layer_sizes=tuple(layers))
    clf.fit(x_train, y_train)
    return clf.score(x_test, y_test)

storage = "sqlite:///simple.db"
study = optuna.create_study(study_name="simple", storage = storage, direction="maximize")
study.optimize(train, n_trials=100)

The logic of the additions should be easy enough to follow. Points to note include:

1. Wrap your main logic into a function that takes a "trial" object which can be used to generate params, in this case integers in a range.
2. Return the objective out of this function (you can have multiple objectives, see this tutorial)
3. Specify a storage database in the study creation - allows use of dashboard later.
4. Specify the optimization direction - maximize in this case, since we want to increace the accuracy. It is set to minimize by default.
5. Run the optimize method, here setting the number of trials to 100.

The code runs for a few minutes, trying out version one after another. It will report version, line after line.

Finally, you can explore this study using the dashboard, which has some cool views:

$ optuna-dashboard sqlite:///simple.db