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.05In 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))