I'm using Keras with Tensorflow backend and looking at nvidia-smi is not sufficient to understand how much memory current network architecture need because seems like Tensorflow just allocate all availible memory.
So the question is how to find out real GPU memory usage?
It can be done using Timeline, which can give you a full trace about memory logging. Similar to the code below:
from keras import backend as K from tensorflow.python.client import timeline import tensorflow as tf with K.get_session() as s: run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE) run_metadata = tf.RunMetadata() # your fitting code and s run with run_options to = timeline.Timeline(run_metadata.step_stats) trace = to.generate_chrome_trace_format() with open('full_trace.json', 'w') as out: out.write(trace) If you want to limit the gpu memory usage, it can alse be done from gpu_options. Like the following code:
import tensorflow as tf from keras.backend.tensorflow_backend import set_session config = tf.ConfigProto() config.gpu_options.per_process_gpu_memory_fraction = 0.2 set_session(tf.Session(config=config)) Check the following documentation about the Timeline object
As you use TensorFlow in the backend, you can use tfprof profiling tool
fit, follows by session.run, by the time your instrumented run call starts, bulk of memory allocated by fit will have been deallocated. There's a related issue here: github.com/tensorflow/tensorflow/issues/9868 What's missing is a recipe of getting Keras to use custom run_options, or adding custom ops to keras model (like MaxBytesInUse op like here )You can still use nvidia-smi after telling TensorFlow not to reserve all memory of the GPU, but to grow this reservation on demand:
config = tf.ConfigProto() config.gpu_options.allow_growth = True keras.backend.tensorflow_backend.set_session(tf.Session(config=config)) 