TensorFlow 在会话启动时分配大量主内存答案

【问题标题】：TensorFlow allocating large amounts of main memory at session startup timeTensorFlow 在会话启动时分配大量主内存
【发布时间】：2017-09-14 05:19:37
【问题描述】：

考虑以下两行 Python/TensorFlow 交互会话：

import tensorflow as tf
s=tf.Session()

如果这些命令在 Ubuntu Linux 14.04 机器上执行，使用 Anaconda Python 2.7.13 和 TensorFlow r1.3（从源代码编译），具有 32G 物理内存和 2 个 GPU（一个 GTX Titan X 和一个 GTX 970），而CUDA_VISIBLE_DEVICES 未设置（即两个 GPU 都可见）生成的 python 进程分配了 59.7G 内存！请注意，它实际上只使用了 754M。

如果CUDA_VISIBLE_DEVICES=0（即只有 Titan X 可见），则分配 55.2G 并使用 137M。

如果CUDA_VISIBLE_DEVICES=1（即只有 970 可见）则分配 47.0G 并使用 325M。

如果CUDA_VISIBLE_DEVICES=（即两个 GPU 都不可见），则仅分配 2.5G 且仅使用 131M。

这在分配的内存量受到限制的环境中是一个问题，例如在 Grid Engine 设置中。

有没有办法限制 TensorFlow 在使用 CUDA 时分配的主内存量？

更新 1

在这些试验中，分配的内存量是通过查看htop 中的VIRT 列来确定的。

TensorFlow r1.3 编译时大多使用默认的configure 答案。唯一的变化是通往 CUDA 和 cuDNN 的路径。因此，jemalloc 正在被使用。

更新 2

我尝试在禁用 jemalloc 的情况下重新编译并看到相同的行为。

【问题讨论】：

可能是由此处讨论的 CUDA 驱动程序问题引起的：stackoverflow.com/questions/11631191/…
如何检查分配了多少内存？也可以尝试使用不同的分配器运行——sudo apt-get install google-perftools; export LD_PRELOAD="/usr/lib/libtcmalloc.so.4"
谢谢@YaroslavBulatov。我尝试使用tcmalloc，但它似乎对行为没有影响。当直接使用 CUDA 而不是通过 TensorFlow 时，我听说过类似的行为，所以我认为最好的 TF 可能是在加载 CUDA 驱动程序时对其进行不同的配置。
啊，你在看虚拟内存。我认为这在 Unix 上相当典型，我经常看到进程分配的巨大“virt”分配没有任何实际缺点
不幸的是，TF+CUDA 进程无法在监视虚拟内存分配的执行容器中运行，以确定进程是否表现自己并保持在指定的约束范围内，例如通过ulimit。 Open Grid Scheduler/Grid Engine 就是一个例子。

标签： python tensorflow

【解决方案1】：

TensorFlow 在 GPU 上的默认行为是使用所有可用内存。但是，如果您想避免这种行为，您可以指定给会话动态分配内存。

来自ConfigProto 声明：

// allow_growth
// If true, the allocator does not pre-allocate the entire specified
// GPU memory region, instead starting small and growing as needed.

为此，请在创建会话时将 ConfigProto 对象传递给您的会话：

session_config = tf.ConfigProto()
session_config.gpu_options.allow_growth=True
sess = tf.Session(config=session_config)

如果您想限制使用的内存量，这取决于您的批量大小和模型中的参数数量。

【讨论】：

这会影响 TF 使用的 GPU 内存量，但我的问题是关于启用 CUDA 时 TF 使用主系统内存的情况。
@DanielRenshaw 我也遇到了同样的问题，你找到解决这个问题的方法了吗？
@OliverHu 不。除非在内存控制的环境中运行，例如 SGE，否则不会造成太大问题，因为实际使用的内存量对我来说仍然足够低。我不得不避免将 SGE 用于 TF 工作。
@DanielRenshaw 对，对于单节点训练，它工作正常。但是在多租户环境下，我们在每个节点上运行了大约 8 个作业，每个作业占用了大约 300 GB 的 VM，机器被吹走了。 :(
我们找到解决方案了吗？ tf Session 在主系统中使用了大量的内存。这样的事情可以在 tf issues 中归档吗？ @OliverHu