如何衡量 TensorFlow 模型的 GPU 内存使用情况

Question

My goal is to figure out how much GPU memory a TensorFlow model saved as a .pb file uses during inference. By limiting the per_process_gpu_memory_fraction to a value of 0.01 and successfully running the model on an input image, I would expect a memory usage of 120MB (based on a 12,000MB GPU).

conf = tf.ConfigProto()
conf.gpu_options.per_process_gpu_memory_fraction=0.01
session = tf.Session(config=conf)

When I observe the actual GPU memory usage via nvidia-smi , I see a utilization of 550MB. Based on the answer here ( https://stackoverflow.com/a/55532954 ), I measure the default memory needed to create a TF session, which amounts to 150MB. Now the question is where the remaining 550MB (measured) - 150MB (session) - 120MB (model) = 280MB are from. I want to exclude any other factors that use the GPU memory and only quantify how much the model itself uses. Can I simply assume 120MB in this case?

Answer 1

Try setting TF_FORCE_GPU_ALLOW_GROWTH to true before you run you python code:

os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'

Put this before you import tensorflow.

Then measure the GPU Memory utilization. Also it also depends on the Tensorflow Memory allocator on how the memory on GPU is allocated. So it is generally going to be a bit larger than the actual model in memory.

Also try suggestions mentioned here: How to restrict tensorflow GPU memory usage?