API Reference for Optuna-Integration

The Optuna-Integration package contains classes used to integrate Optuna with external machine learning frameworks.

All of these classes can be imported in two ways. One is “from optuna.integration import xxx” like a module in Optuna, and the other is “from optuna_integration import xxx” as an Optuna-Integration specific module. The former is provided for backward compatibility.

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.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.qkg_candidates_func`	Quasi MC-based batch Knowledge Gradient (qKG).
`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.
`optuna_integration.botorch.qhvkg_candidates_func`	Quasi MC-based batch Hypervolume Knowledge Gradient (qHVKG).

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.

Comet

optuna_integration.CometCallback

A callback for logging Optuna study trials to a Comet ML Experiment.

Dask

optuna_integration.DaskStorage

Dask-compatible storage class.

fast.ai

`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.

MXNet

optuna_integration.MXNetPruningCallback

MXNet callback to prune unpromising trials.

pycma

optuna_integration.PyCmaSampler

A Sampler using cma library as the backend.

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.