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

Question

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

import tensorflow as tf
s=tf.Session()

如果这些命令的Ubuntu Linux操作系统14.04机，采用水蟒Python的2.7.13和TensorFlow R1.3（从来源汇编）中，用32G物理存储器和2分的GPU（一个GTX泰坦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

在这些试验中，通过查看中的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)

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