Neptune-Optuna Integration

What will you get with this integration?

Optuna is an open source hyperparameter optimization framework to automate hyperparameter search. With Neptune integration, you can:

  • see the experiment as it is running,

  • see charts of logged run scores,

  • log the parameters tried at every run,

  • log the best parameters after training,

  • log and explore interactive optuna visualizations like plot_contour, plot_slice, plot_parallel_coordinate, and optimization_history,

  • save the optuna.study object itself.

Note

This integration is tested with optuna==2.3.0, neptune-client==0.4.125, and neptune-contrib==0.24.7.

Where to start?

To get started with this integration, follow the quickstart below. You can also skip the basics and take a look at how to log Optuna visualizations and the optuna.study object during or after the sweep in the advanced options section.

If you want to try things out and focus only on the code you can either:

  1. Open the Colab notebook (badge-link below) with quickstart code and run it as an anonymous user “neptuner” - zero setup, it just works,

  2. View quickstart code as a plain Python script on GitHub.

Quickstart

This quickstart will show you how to:

  • Install the necessary neptune packages

  • Connect Neptune to your Optuna hyperparameter tuning code and create the first experiment

  • Log metrics, figures, and artifacts from your first Optuna sweep to Neptune, and

  • Explore them in the Neptune UI.

Before you start

You have Python 3.x and following libraries installed:

Note

The integration should work on newer versions of neptune-client and neptune-contrib too.

You also need minimal familiarity with Optuna. Have a look at the Optuna tutorial guide to get started.

pip install --quiet optuna neptune-client neptune-contrib['monitoring']

Step 1: Initialize Neptune

Run the code below:

import neptune

neptune.init(api_token='ANONYMOUS', project_qualified_name='shared/optuna-integration')

Tip

You can also use your personal API token. Read more about how to securely set the Neptune API token.

Step 2: Create an Experiment

Run the code below to create a Neptune experiment:

neptune.create_experiment('optuna-sweep')

This also creates a link to the experiment. Open the link in a new tab. The charts will currently be empty, but keep the window open. You will be able to see live metrics once logging starts.

Step 3: Create the Neptune Callback

import neptunecontrib.monitoring.optuna as opt_utils

neptune_callback = opt_utils.NeptuneCallback()

Step 4: Run Optuna with the Neptune callback

Pass the neptune_callback as a callback to study.optimize() to monitor the metrics and parameters checked at each run.

study = optuna.create_study(direction='maximize')
study.optimize(objective, n_trials=100, callbacks=[neptune_callback])

Step 5: Monitor your Optuna training in Neptune

Now you can switch to the Neptune tab which you had opened previously to watch the optimization live!

Check out this example experiment.

Advanced Options

Log charts and study object during sweep

While creating the Neptune Callback, you can set log_study=True and log_charts=True to log interactive charts from optuna.visualization and the study object itself after every iteration.

neptune_callback = opt_utils.NeptuneCallback(log_study=True, log_charts=True)

Warning

Depending on the size of the optuna.study object and the charts, this might add some overhead to the sweep. To avoid this, you can log the study object and charts after the sweep.

Log charts and study object after sweep

You can log the optuna.study object and charts after the sweep has completed by running:

opt_utils.log_study_info(study)

Check out this example experiment with advanced logging.

What’s next

Now that you know how to integrate Neptune with Optuna, you can check: