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:

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.

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.