Skip to content

Instantly share code, notes, and snippets.

@amarion35

amarion35/keras_gpu_options.py

Created October 11, 2018 17:21
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save amarion35/8c4452f0f70b1e7eebe1119c68843c87 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save amarion35/8c4452f0f70b1e7eebe1119c68843c87 to your computer and use it in GitHub Desktop.
Download ZIP
Choose your GPU and limit Keras memory usage
Raw
keras_gpu_options.py
import os
os.environ["CUDA_VISIBLE_DEVICES"]="1" #Before all
import tensorflow as tf
config = tf.ConfigProto()
config.gpu_options.allow_growth = True # Use only necessary memory
#config.gpu_options.per_process_gpu_memory_fraction = 0.5
session = tf.Session(config=config) # Before importing Keras!
from keras.backend.tensorflow_backend import set_session
set_session(session)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment