Is there a way of determining how much GPU memory is in use by TensorFlow?

user_349760 0
Viewed 32089

Tensorflow tends to preallocate the entire available memory on it's GPUs. For debugging, is there a way of telling how much of that memory is actually in use?

tensorflow gpu
4 Answers

Here's a practical solution that worked well for me:

Disable GPU memory pre-allocation using TF session configuration: 

config = tf.ConfigProto()  
config.gpu_options.allow_growth=True  
sess = tf.Session(config=config)

run nvidia-smi -l (or some other utility) to monitor GPU memory consumption.

Step through your code with the debugger until you see the unexpected GPU memory consumption.

user_8961110

There's some code in tensorflow.contrib.memory_stats that will help with this:

from tensorflow.contrib.memory_stats.python.ops.memory_stats_ops import BytesInUse
with tf.device('/device:GPU:0'):  # Replace with device you are interested in
  bytes_in_use = BytesInUse()
with tf.Session() as sess:
  print(sess.run(bytes_in_use))
user_99479510
Related