Skip to content

Instantly share code, notes, and snippets.

@rohit-gupta

rohit-gupta/tf_limit_gpu_memory.py

Created February 8, 2019 09:34
Show Gist options
  • Save rohit-gupta/62a2e8b560cb5cd20e2d90607d15bb31 to your computer and use it in GitHub Desktop.
Save rohit-gupta/62a2e8b560cb5cd20e2d90607d15bb31 to your computer and use it in GitHub Desktop.
Download ZIP
Limit GPU memory usage in Keras
Raw
tf_limit_gpu_memory.py
import tensorflow as tf
import keras.backend as K
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.25)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
K.set_session(sess)
## Use normally after this
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment