如何限制张量流内存使用？答案

【问题标题】：How to limit tensorflow memory usage?如何限制张量流内存使用？
【发布时间】：2018-04-24 12:06:36
【问题描述】：

我是 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)

Memory usage of the python program above

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

嗯，不知道有没有办法优化呢？

【问题讨论】：

您是在尝试优化主系统 RAM 使用率还是 GPU RAM 使用率？
tf_config.gpu_options.allow_growth = False 允许 Tensorflow 分配所有 GPU 的 RAM。如果您想要更精细的粒度分配，请将其放在True（但您可能会在性能上有所损失）
哦，我不是说 GPU，它只是内存。
这不是一个好问题，）：

标签： python tensorflow

【解决方案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）

【讨论】：

像 OP 一样，我需要限制 CPU 内存，而不是 GPU 内存。当我在 CPU 设备而不是 GPU 设备上尝试这种方法时，我收到以下错误：ValueError: Setting memory limit on CPU virtual devices is currently not supported