Keras: real amount of GPU memory used

Question

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 l

Answer 1

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