[Feature Request] Resume study and avoid OOM when optimizing with Optuna plugin

🚀 Feature Request

I would like to be able:

• to persist my study in order to resume my hyperparameter search from where I left it
• set the gc_after_trial parameter of Optuna’s study.optimize()

Motivation

Is your feature request related to a problem? Please describe. I’m always frustrated when my code crashes after 60 trials (out of 100). I suspect an OOM error. Being able to prevent the script from crashing in the first place with gc.collect() would be great. However, at least being able to resume my search from where it stopped would be a game changer.

Pitch

Describe the solution you’d like I would like to set gc_after_trial to True and a path to store my study parameters after each trial in my Optuna sweeper hydra config.

Describe alternatives you’ve considered I read Optuna’s documentation but I’m not sure how to make their examples work with Hydra:

• How do I avoid running out of memory (OOM) when optimizing studies?
• How can I save and resume studies?

Are you willing to open a pull request? (See CONTRIBUTING) I’m not comfortable enough with Optuna’s and Hydra’s library to prepare a pull request.

[omry]

Hydra has callbacks which can probably be used for it. See this.

[omry]

Hi @dianemarquette, I am open to supporting it, although resume study might be harder than in seems (in general resume is not something supported by any Hydra Sweeeper right now).

In any case, we do not have the cycles for it, which means this will happen if someone from the community wants to work toward it.

Supporting gc_after_trial seems like it should be straight forward though.