Limit Keras GPU Memory with Tensorflow v2

limit_keras_gpu_mem.py

"""
For TensorFlow v2 and Keras you can limit gpu memory usage a couple of different ways, just add one of the following early on in your scripts
"""
import tensorflow as tf

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

"""
or
"""
import tensorflow as 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))

"""
For Tensorflow v1 
"""
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))