Question

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

这是我的剧本：

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

根据我的观察，'导入张量流为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）

如何限制tensorflow内存使用？

1 个答案: