如何限制 tensorflow 内存使用？

Question

I am new to TensorFlow. I found it took up too much memory when I run a simple script. I do not mean GPU memory, I mean CPU memory.

Here is my script:

# -*- coding: utf-8 -*-

import time
import tensorflow as tf
tf_config = tf.ConfigProto()
tf_config.gpu_options.allow_growth = False
with tf.Session(config=tf_config) as sess:
    print('Listening.....')
    time.sleep(100) 

Memory usage of the python program above

According to my observation, 'import tensorflow as tf' takes about 100MB, and tf.Session takes others.

Well, I wonder if there is any way to optimize it?

Answer 1

In recent TensorFlow 2.0 we could specify the required amount of memory explicitly.

import tensorflow as tf
assert tf.version.VERSION.startswith('2.')

gpus = tf.config.experimental.list_physical_devices('GPU')

tf.config.experimental.set_virtual_device_configuration(gpus[0], 
   [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)])

(Source: https://github.com/tensorflow/tensorflow/issues/25138#issuecomment-533936197 )