如何限制 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 ）