Reserving gpu memory tensorflow before loading data

Question

Is there a way to reserve GPU memory before the graph is built and run in tensorflow?

I have a slow data loading that needs to be done before creating the graph and running it. I am sharing a GPU cluster so sometimes other people would use the GPU that I have designated before I'm done loading the data and the program will encounter OOM error.

Answer 1

Device memory is reserved when you instantiate a tf.Session.

If you need to reserve device memory early on, you could start by creating a session, and then proceed normally to construct your graph. Although it is common in "raw" tensorflow to create the graph first then instantiate a session, it is nonetheless allowed to proceed differently, and it is actually common in the Keras world where one would often start a script with

import tensorflow as tf
sess = tf.Session()

from keras import backend as K
K.set_session(sess)
...