如何防止TensorFlow分配GPU内存的全部？

Question

如何防止TensorFlow分配GPU内存的全部？我在一个共享计算资源的环境中工作，也就是说，我们有几台服务器机器，每个服务器都配备了几个Nvidia Titan X GPU。对于中小型型号，12 GB的土卫六X通常足以让2-3人同时在同一个GPU上进行训练。如果模型足够小，单个模型不能充分利用土卫六X的所有计算单元，那么与一个接一个的训练过程相比，这实际上会导致加速。即使在对GPU的并发访问确实减慢了单个培训时间的情况下，仍然可以灵活地让多个用户同时在GPU上运行。TensorFlow的问题是，默认情况下，当GPU启动时，它会在GPU上分配全部可用内存。即使对于一个小的两层神经网络，我看到12 GB的土卫六X已经用完了。是否有办法使TensorFlow只分配，例如，4GB的GPU内存，如果一个人知道这个数量足够一个给定的模型？

Answer 1

可以将GPU内存的部分设置为在构造tf.Session通过一个tf.GPUOptions作为任择条款的一部分config论点：

# Assume that you have 12GB of GPU memory and want to allocate ~4GB:gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

这个per_process_gpu_memory_fraction作为在同一台机器上的每个GPU上的进程将使用的GPU内存量的硬上限。目前，这个分数被统一地应用于同一台机器上的所有GPU；没有办法在每个GPU的基础上设置它。

Answer 2

以下是这本书的摘录Deep Learning with TensorFlow

在某些情况下，进程只需要分配可用内存的子集，或者只在进程需要时增加内存使用量。TensorFlow提供两种配置会话中控制此操作的选项。第一个是allow_growth选项，它尝试根据运行时分配的GPU内存，它开始分配很少的内存，并且随着会话的运行和需要更多的GPU内存，我们扩展了TensorFlow进程所需的GPU内存区域。

1)允许增长：(更灵活)

config = tf.ConfigProto()config.gpu_options.allow_growth = Truesession = tf.Session(config=config, ...)

第二种方法是per_process_gpu_memory_fraction选项，该选项确定each应该分配可见的GPU。注：不需要释放内存，这样做甚至会使内存碎片恶化。

2)分配固定内存:

只分配40%每个GPU的总内存：

config = tf.ConfigProto()config.gpu_options.per_process_gpu_memory_fraction = 0.4session = tf.Session(config=config, ...)

注：不过，只有当您真正想绑定TensorFlow进程上可用的GPU内存时，这才是有用的。