Managing Tensorflow GPU Memory Usage

GPU Mem Growth

Described in https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/protobuf/config.proto

Tensorflow will start small ... and increase the use of memory over time.

config = tf.ConfigProto()
config.gpu_options.allow_growth = True

tf.Session(config=config)

If you want a hard limit

config.gpu_options.per_process_gpu_memory_fraction = 0.05

In this case it is 5% of total memory


In [3]:
import tensorflow as tf

tf.reset_default_graph()

a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
c = tf.matmul(a, b)

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
# memory hard limit
config.gpu_options.per_process_gpu_memory_fraction = 0.05

with tf.Session(config=config) as sess:
    print(sess.run(c))


[[ 22.  28.]
 [ 49.  64.]]


Konpat Preechakul

Read more posts by this author.