如何限制 tensorflow 内存使用？

Question

我是 TensorFlow 的新手。 当我运行一个简单的脚本时，我发现它占用了太多内存。 我不是说 GPU 内存，我是说CPU内存。

这是我的脚本：

# -*- 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) 

上面python程序的内存使用情况

根据我的观察，'import tensorflow as tf' 需要大约 100MB，而 tf.Session 需要其他人。

嗯，我想知道有没有什么方法可以优化它？

Answer 1

在最近的 TensorFlow 2.0 中，我们可以明确指定所需的内存量。

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)])

（来源： https : //github.com/tensorflow/tensorflow/issues/25138#issuecomment-533936197 ）