optuna.integration

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

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.

AllenNLP

optuna.integration.AllenNLPExecutor

AllenNLP extension to use optuna with Jsonnet config file.

optuna.integration.allennlp.dump_best_config

Save JSON config file with environment variables and best performing hyperparameters.

optuna.integration.AllenNLPPruningCallback

AllenNLP callback to prune unpromising trials.

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.

Catalyst

optuna.integration.CatalystPruningCallback

Catalyst callback to prune unpromising trials.

CatBoost

optuna.integration.CatBoostPruningCallback

Callback for catboost to prune unpromising trials.

Chainer

optuna.integration.ChainerPruningExtension

Chainer extension to prune unpromising trials.

optuna.integration.ChainerMNStudy

A wrapper of Study to incorporate Optuna with ChainerMN.

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

Keras

optuna.integration.KerasPruningCallback

Keras callback to prune unpromising trials.

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.

skorch

optuna.integration.SkorchPruningCallback

Skorch callback to prune unpromising trials.

TensorFlow

optuna.integration.TensorBoardCallback

Callback to track Optuna trials with TensorBoard.

optuna.integration.TensorFlowPruningHook

TensorFlow SessionRunHook to prune unpromising trials.

optuna.integration.TFKerasPruningCallback

tf.keras callback to prune unpromising trials.

XGBoost

optuna.integration.XGBoostPruningCallback

Callback for XGBoost to prune unpromising trials.