optuna.integration

The integration module contains classes used to integrate Optuna with external machine learning frameworks.

Note

Optuna’s integration modules for third-party libraries have started migrating from Optuna itself to a package called optuna-integration. Please check the repository and the documentation.

For most of the ML frameworks supported by Optuna, the corresponding Optuna integration class serves only to implement a callback object and functions, compliant with the framework’s specific callback API, to be called with each intermediate step in the model training. The functionality implemented in these callbacks across the different ML frameworks includes:

1. Reporting intermediate model scores back to the Optuna trial using optuna.trial.Trial.report(),

2. According to the results of optuna.trial.Trial.should_prune(), pruning the current model by raising optuna.TrialPruned(), and

3. Reporting intermediate Optuna data such as the current trial number back to the framework, as done in MLflowCallback.

For scikit-learn, an integrated OptunaSearchCV estimator is available that combines scikit-learn BaseEstimator functionality with access to a class-level Study object.

BoTorch

optuna.integration.BoTorchSampler	A sampler that uses BoTorch, a Bayesian optimization library built on top of PyTorch.
optuna.integration.botorch.ehvi_candidates_func	Expected Hypervolume Improvement (EHVI).
optuna.integration.botorch.logei_candidates_func	Log Expected Improvement (LogEI).
optuna.integration.botorch.qei_candidates_func	Quasi MC-based batch Expected Improvement (qEI).
optuna.integration.botorch.qnei_candidates_func	Quasi MC-based batch Noisy Expected Improvement (qNEI).
optuna.integration.botorch.qehvi_candidates_func	Quasi MC-based batch Expected Hypervolume Improvement (qEHVI).
optuna.integration.botorch.qnehvi_candidates_func	Quasi MC-based batch Noisy Expected Hypervolume Improvement (qNEHVI).
optuna.integration.botorch.qparego_candidates_func	Quasi MC-based extended ParEGO (qParEGO) for constrained multi-objective optimization.

CatBoost

optuna.integration.CatBoostPruningCallback

Callback for catboost to prune unpromising trials.

Dask

optuna.integration.DaskStorage

Dask-compatible storage class.

LightGBM

optuna.integration.LightGBMPruningCallback	Callback for LightGBM to prune unpromising trials.
optuna.integration.lightgbm.train	Wrapper of LightGBM Training API to tune hyperparameters.
optuna.integration.lightgbm.LightGBMTuner	Hyperparameter tuner for LightGBM.
optuna.integration.lightgbm.LightGBMTunerCV	Hyperparameter tuner for LightGBM with cross-validation.

MLflow

optuna.integration.MLflowCallback

Callback to track Optuna trials with MLflow.

Weights & Biases

optuna.integration.WeightsAndBiasesCallback

Callback to track Optuna trials with Weights & Biases.

pycma

optuna.integration.PyCmaSampler	A Sampler using cma library as the backend.
optuna.integration.CmaEsSampler	Wrapper class of PyCmaSampler for backward compatibility.

PyTorch

optuna.integration.PyTorchIgnitePruningHandler	PyTorch Ignite handler to prune unpromising trials.
optuna.integration.PyTorchLightningPruningCallback	PyTorch Lightning callback to prune unpromising trials.
optuna.integration.TorchDistributedTrial	A wrapper of Trial to incorporate Optuna with PyTorch distributed.

scikit-learn

optuna.integration.OptunaSearchCV

Hyperparameter search with cross-validation.

scikit-optimize

optuna.integration.SkoptSampler

Sampler using Scikit-Optimize as the backend.

TensorFlow

optuna.integration.TensorBoardCallback

Callback to track Optuna trials with TensorBoard.

XGBoost

optuna.integration.XGBoostPruningCallback

Callback for XGBoost to prune unpromising trials.

Dependencies of each integration

We summarize the necessary dependencies for each integration.

Integration	Dependencies
AllenNLP	allennlp, torch, psutil, jsonnet
BoTorch	botorch, gpytorch, torch
Catalyst	catalyst
CatBoost	catboost
ChainerMN	chainermn
Chainer	chainer
pycma	cma
Dask	distributed
FastAI (v1, v2)	fastai
Keras	keras
LightGBMTuner	lightgbm, scikit-learn
LightGBMPruningCallback	lightgbm
MLflow	mlflow
MXNet	mxnet
PyTorch Distributed	torch
PyTorch (Ignite)	pytorch-ignite
PyTorch (Lightning)	pytorch-lightning
SHAP	scikit-learn, shap
Scikit-learn	pandas, scipy, scikit-learn
Scikit-optimize	scikit-optimize
SKorch	skorch
TensorBoard	tensorboard, tensorflow
TensorFlow	tensorflow, tensorflow-estimator
TensorFlow + Keras	tensorflow
Weights & Biases	wandb
XGBoost	xgboost