Skip to content

Instantly share code, notes, and snippets.

@vikramsoni2
Last active July 23, 2020 09:23
Show Gist options
  • Save vikramsoni2/9e9632054551c6d44bb6920249fb9705 to your computer and use it in GitHub Desktop.
Save vikramsoni2/9e9632054551c6d44bb6920249fb9705 to your computer and use it in GitHub Desktop.
Tensorflow dynamic gpu resource allocation
# actual code you can copy-paste into your Keras code to have Tensorflow dynamically allocate the GPU memory:
import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True # dynamically grow the memory used on the GPU
config.log_device_placement = True # to log device placement (on which device the operation ran)
# (nothing gets printed in Jupyter, only if you run it standalone)
sess = tf.Session(config=config)
set_session(sess) # set this TensorFlow session as the default session for Keras
#####################################################################################
from keras import backend as K
cfg = K.tf.ConfigProto()
cfg.gpu_options.allow_growth = True
K.set_session(K.tf.Session(config=cfg))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment