如何防止张量流分配GPU内存的全部？

我在一个共享计算资源的环境中工作,也就是说,我们有一些服务器机器配备了几个Nvidia Titan X GPU.

对于小到中等大小的型号,12GB的Titan X通常足以让2-3人在同一GPU上同时进行训练.如果模型足够小以至于单个模型没有充分利用Titan X的所有计算单元,那么与在另一个训练过程之后运行一个训练过程相比,这实际上可以导致加速.即使在并发访问GPU确实减慢了单个培训时间的情况下,仍然可以灵活地让多个用户同时在GPU上运行.

TensorFlow的问题在于,默认情况下,它在启动时会在GPU上分配全部可用内存.即使对于一个小的2层神经网络,我也看到12 GB的Titan X已用完.

有没有办法让TensorFlow只分配4GB的GPU内存,如果有人知道这个数量对于给定的模型来说足够了？

1> mrry..：

---

您可以tf.Session通过传递a tf.GPUOptions作为可选config参数的一部分来设置构造a时要分配的GPU内存的分数:

# 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的基础上设置它.

---

相关说明:设置CUDA_VISIBLE_DEVICES以将TensorFlow限制为单个GPU对我有用.请参阅http://www.acceleware.com/blog/cudavisibledevices-masking-gpus

---

非常感谢你。该信息在当前文档中相当隐蔽。我永远不会自己找到它：-)如果您可以回答，我想问两个附加信息：1-这是否限制了曾经使用的内存量，或者仅仅是限制了最初分配的内存？（即，如果计算图需要它，它将仍然分配更多的内存）2-是否可以在每个GPU的基础上进行设置？

---

似乎内存分配超出了请求，例如我在24443MiB gpu上请求了per_process_gpu_memory_fraction = 0.0909并获得了占用2627MiB的进程

---

我似乎无法在"MonitoredTrainingSession"中使用它

---

@jeremy_rutman我相信这是由于cudnn和cublas上下文初始化引起的。这仅在您正在执行使用这些库的内核时才有意义。

2> Sergey Demya..：

---

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

https://github.com/tensorflow/tensorflow/issues/1578

---

这个正是我想要的,因为在多用户环境中,指定要在代码本身中保留的GPU内存的确切数量是非常不方便的.

---

另外，如果您将Keras与TF后端一起使用，则可以使用它并从keras import backend运行为K和K.set_session（sess）以避免内存限制

3> user1767754..：

---

这是本书的摘录 Deep Learning with TensorFlow

在某些情况下,希望该过程仅分配可用内存的子集,或者仅增加该过程所需的内存使用量.TensorFlow 在会话中提供了两个配置选项来控制它.第一个是allow_growth选项,它尝试仅基于运行时分配分配尽可能多的GPU内存,它开始分配非常少的内存,并且当会话运行并且需要更多GPU内存时,我们扩展了TensorFlow所需的GPU内存区域处理.

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

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
session = 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.4
session = tf.Session(config=config, ...)

注意: 仅当您真正想要绑定TensorFlow进程上可用的GPU内存量时,这才有用.

4> Theo..：

---

已针对TensorFlow 2.0 Alpha及更高版本进行了更新

从2.0 Alpha文档来看，答案现在只有一行，然后您可以使用TensorFlow进行任何操作：

import tensorflow as tf
tf.config.gpu.set_per_process_memory_growth(True)

5> Urs..：

---

上面的所有答案都假定使用sess.run()调用执行,这将成为例外,而不是最近版本的TensorFlow中的规则.

当使用tf.Estimator框架(TensorFlow 1.4及更高版本)时,将分数传递给隐式创建的方法MonitoredTrainingSession是,

opts = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
conf = tf.ConfigProto(gpu_options=opts)
trainingConfig = tf.estimator.RunConfig(session_config=conf, ...)
tf.estimator.Estimator(model_fn=..., 
                       config=trainingConfig)

同样在Eager模式(TensorFlow 1.5及以上版本)中,

opts = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
conf = tf.ConfigProto(gpu_options=opts)
tfe.enable_eager_execution(config=conf)

编辑:11-04-2018 作为一个例子,如果你要使用tf.contrib.gan.train,那么你可以使用类似于下面的东西:

tf.contrib.gan.gan_train(........, config=conf)

6> mx_muc..：

---

Tensorflow 2.0 Beta和（可能）超越

API再次更改。现在可以在以下位置找到它：

tf.config.experimental.set_memory_growth(
    device,
    enable
)

别名：

tf.compat.v1.config.experimental.set_memory_growth

tf.compat.v2.config.experimental.set_memory_growth

参考文献：

https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/config/experimental/set_memory_growth

https://www.tensorflow.org/guide/gpu#limiting_gpu_memory_growth

另请参阅： Tensorflow-使用GPU：https : //www.tensorflow.org/guide/gpu

对于Tensorflow 2.0 Alpha，请参阅： 此答案