Hydra optuna sweeper plugin which suports pruning

[DirkKuhn]

Hi,

I have been using hydra and optuna for a couple of months now. Since I wanted to use a pruner with hydra, I have developed an optuna sweeper which supports pruning among other features.

Here are its advantages over the original Optuna Sweeper plugin:

• Pruning with

    from hydra_plugins.hydra_optuna_pruning_sweeper import trial_provider
    ...
    trial = trial_provider.trial
    ...
    trial.report(score, step)
    if trial.should_prune():
        ...
    ...

  in the objective function.
• Manual specification of hyperparameters which are used at the beginning of the hyperparameter search (see https://optuna.readthedocs.io/en/stable/tutorial/20_recipes/008_specify_params.html).
• Simple parallelization across processes with the optuna.integration.DaskStorage
  simply by specifying n_workers>1.
• Parallelization across nodes with the optuna.integration.DaskStorage by specifying a dask Client.
• More powerful custom search spaces. For example, a variable length list of values can be suggested.
• It internally uses optuna.study.Study.optimize
  which can be freely configured.
• It uses optuna>=3.1.0. This is only an advantage until keisuke-umezawa's Pull Request is merged.

Before anyone gets too exited there are also some disadvantages/limitations:

• This plugin is supposed to be used with the BasicLauncher. Different launchers like the Submitit Launcher plugin are not supported.
• The deprecated parameter search_space from the original has been removed.
• This plugin internally uses hydra_zen for configuration. It requires importing optuna and dask.distributed. Both might be imported even if this plugin is not used and thus, once installed, it can slow down starting hydra.
• This plugin does not seek to improve the architecture of hydra nor does it allow the usage of hyperparameter search libraries other than optuna.
• Fewer people use this plugin. Therefore, it might not be as robust as the original.

If you are interested please check out the repository. In the README file I have outlined the usage which is mostly similar to the original optuna sweeper and my design decisions. In the examples package you can find the adapted examples from the original optuna sweeper. I have also implemented a deep-learning example with pytorch lightning which showcases the new sweeper with a more complicated configuration using hydra and hydra-zen.

Since I have not published any code before, I would be happy about any feedback.

Thank you for your interest!

[devanshamin]

I implemented Optuna pruning with Hydra without any external dependencies.

• litgnn/optimize.py
• litgnn/trainer/optuna.py
• litgnn/trainer/hydra_trainer.py

The objective function uses the Hydra internals to suggest hparams and PyTorchLightningPruningCallback is passed to the Lightning Trainer.