optuna.integration.PyTorchLightningPruningCallback

class optuna.integration.PyTorchLightningPruningCallback(trial, monitor)[source]

PyTorch Lightning callback to prune unpromising trials.

See the example if you want to add a pruning callback which observes accuracy.

Parameters
  • trial – A Trial corresponding to the current evaluation of the objective function.

  • monitor – An evaluation metric for pruning, e.g., val_loss or val_acc. The metrics are obtained from the returned dictionaries from e.g. pytorch_lightning.LightningModule.training_step or pytorch_lightning.LightningModule.validation_epoch_end and the names thus depend on how this dictionary is formatted.

Note

For the distributed data parallel training, the version of PyTorchLightning needs to be higher than or equal to v1.4.0. In addition, Study should be instantiated with RDB storage.

Methods

on_fit_end(trainer, pl_module)

on_init_start(trainer)

on_validation_end(trainer, pl_module)