Optuna - Memory Issues

Question

I am trying to free memory in between Optuna optimization runs. I am using python 3.8 and the latest version of Optuna. What happens is I run the commands: optuna.create_study(), then I call optuna.optimize(...) in a loop, with a new objective function each time. When I monitor my memory usage, each time the command optuna.create_study() is called, memory usage keeps on increasing to the point that my processor just kills the program eventually. Just for a more clear picture, the first run takes over 3% memory and it eventually builds up to >80%. Any thoughts on how I can remove a study from memory in between successive calls of create_study()?

Answer 1

I had a similar problem when running my trial (although I did not loop over multiple optimizations). Creating a callback that frees garbage at the end of each epoch solved my problem and will probably already help you free quite some space. Try the following:

import tensorflow as tf
import gc

class MyCustomCallback(tf.keras.callbacks.Callback):
  def on_epoch_end(self, epoch, logs=None):
    gc.collect()
...
_ = model.fit(x_train, y_train, ...
          callbacks=[MyCustomCallback()])```

Solution based on this issue.