How to add report_tensor_allocations_upon_oom to RunOptions in Keras

TF1 solution:

Its not as hard as it seems, what you need to know is that according to the documentation, the **kwargs parameter passed to model.compile will be passed to session.run

So you can do something like:

import tensorflow as tf
run_opts = tf.RunOptions(report_tensor_allocations_upon_oom = True)

model.compile(loss = "...", optimizer = "...", metrics = "..", options = run_opts)

And it should be passed directly each time session.run is called.

TF2:

The solution above works only for tf1. For tf2, unfortunately, it appears there is no easy solution yet.


OOM means out of memory. May be it is using more memory at that time. Decrease batch_size significantly. I set to 16, then it worked fine


Currently, it is not possible to add the options to model.compile. See: https://github.com/tensorflow/tensorflow/issues/19911