简体   繁体   中英

How to measure GPU memory usage of TensorFlow model

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?

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?

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM