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:

Reporting intermediate model scores back to the Optuna trial using optuna.trial.Trial.report(),
According to the results of optuna.trial.Trial.should_prune(), pruning the current model by raising optuna.TrialPruned(), and
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.qei_candidates_func`	Quasi MC-based batch Expected Improvement (qEI).
`optuna.integration.botorch.qehvi_candidates_func`	Quasi MC-based batch Expected Hypervolume Improvement (qEHVI).
`optuna.integration.botorch.qnehvi_candidates_func`	Quasi MC-based batch Expected Noisy 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.

fast.ai

`optuna.integration.FastAIV1PruningCallback`	FastAI callback to prune unpromising trials for fastai.
`optuna.integration.FastAIV2PruningCallback`	FastAI callback to prune unpromising trials for fastai.
`optuna.integration.FastAIPruningCallback`	alias of `FastAIV2PruningCallback`

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.

MXNet

optuna.integration.MXNetPruningCallback

MXNet callback to prune unpromising trials.

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.

SHAP

optuna.integration.ShapleyImportanceEvaluator

Shapley (SHAP) parameter importance evaluator.

TensorFlow

optuna.integration.TensorBoardCallback

Callback to track Optuna trials with TensorBoard.

XGBoost

optuna.integration.XGBoostPruningCallback

Callback for XGBoost to prune unpromising trials.