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

Question

考虑以下两行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。

这是在分配的内存量受限制的环境中的问题，例如，在网格引擎设置中。

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

更新1

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

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

更新2

我已经尝试重新编译并禁用jemalloc并看到相同的行为。

Answer 1

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

来自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)

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