Skip to content

Preventing Tensorflow Keras from allocating the totality of a GPU memories.

fialhocoelho edited this page Dec 20, 2018 · 3 revisions

To prevent Tensorflow from allocating the totality of a GPU memories, we need enable the allow_growth setting in Tensorflow or Keras. The following code for setting allow_growth memory option in:

Tensorflow:

import tensorflow as tf

config = tf.ConfigProto()
config.gpu_options.allow_growth = True # dynamically grow the memory used on the GPU
session = tf.Session(config=config)

Keras:

from keras.callbacks import ModelCheckpoint
from keras.models import Model, load_model, save_model, Sequential
from keras.layers import Dense, Activation, Dropout, Input, Masking, TimeDistributed, LSTM, Conv1D
from keras.layers import GRU, Bidirectional, BatchNormalization, Reshape
from keras.optimizers import Adam
from keras.backend.tensorflow_backend import set_session

import tensorflow as tf

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)

sess = tf.Session(config=config)
tf.keras.backend.set_session(sess)  # set this TensorFlow session as the default session for Keras

Inserting this code, your application will only allocate the necessary memory in the GPUs.

for more information:

https://www.tensorflow.org/guide/using_gpu