limit gpu memory usage of keras

keras_limit.py

#put the these lines before importing any module from keras.
import tensorflow as tf

from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.gpu_options.visible_device_list = "0" #only the gpu 0 is allowed
set_session(tf.Session(config=config))

Thank You!

For TensorFlow v2 and Keras this call can be updated to

import tf
from keras.backend import set_session

config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True
config.gpu_options.visible_device_list = "0"  # only the gpu 0 is allowed
set_session(tf.compat.v1.Session(config=config))

or

import tf

physical_devices = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(physical_devices[0], True)