Share GPU memory for different users in keras and tensorflow

2.4k Views Asked by At

We had only one GPU installed with CUDA drivers and whenever one user runs the code, the whole memory is assigned to that user. And the other users are unable to use the GPU. Is there a way to get rid of this behavior?

1

There are 1 best solutions below

3
vvvvv On

If you are using keras, add this at the beginning of your script:

from keras import backend as K

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

This will prevent tensorflow to take all the memory as can be seen here.

If you are using tensorflow without keras, add this:

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.Session(config=config, ...)

As shown here.