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

Question

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

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

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

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

Answer 1

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

Answer 2

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

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

Answer 3

这是本书的摘录 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内存量时,这才有用.

Answer 4

对于 Tensorflow 2.0 和 2.1 版，请使用以下代码段：

 import tensorflow as tf
 gpu_devices = tf.config.experimental.list_physical_devices('GPU')
 tf.config.experimental.set_memory_growth(gpu_devices[0], True)

对于以前的版本，以下代码段曾经对我有用：

import tensorflow as tf
tf_config=tf.ConfigProto()
tf_config.gpu_options.allow_growth=True
sess = tf.Session(config=tf_config)

Answer 5

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

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

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

Answer 6

上面的所有答案都假定使用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)

Answer 7

您可以使用

TF_FORCE_GPU_ALLOW_GROWTH=true

在您的环境变量中。

在张量流代码中：

bool GPUBFCAllocator::GetAllowGrowthValue(const GPUOptions& gpu_options) {
  const char* force_allow_growth_string =
      std::getenv("TF_FORCE_GPU_ALLOW_GROWTH");
  if (force_allow_growth_string == nullptr) {
    return gpu_options.allow_growth();
}

Answer 8

无耻的插件：如果你安装了 GPU 支持的 Tensorflow，无论你将其设置为仅使用 CPU 还是 GPU，会话都会首先分配所有 GPU。我可能会添加我的提示，即使您将图形设置为仅使用 CPU，您也应该设置相同的配置（如上回答:)）以防止不必要的 GPU 占用。

而在像 IPython 和 Jupyter 这样的交互界面中，你也应该设置配置，否则，它会分配所有内存，几乎没有给其他人留下。这有时很难注意到。

Answer 9

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，请参阅： 此答案

Answer 10

如果您使用的是 Tensorflow 2，请尝试以下操作：

config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.compat.v1.Session(config=config)

Answer 11

上面的所有答案都是指在TensorFlow 1.X版本中将内存设置到一定程度，或者允许TensorFlow 2.X.

该方法 tf.config.experimental.set_memory_growth确实适用于在分配/预处理期间允许动态增长。然而，人们可能希望从一开始就分配一个特定的 GPU 内存。

分配特定 GPU 内存背后的逻辑也是为了防止训练期间出现 OOM 内存。例如，如果在打开消耗视频内存的 Chrome 选项卡/任何其他视频消耗过程时进行训练，tf.config.experimental.set_memory_growth(gpu, True)可能会导致抛出 OOM 错误，因此在某些情况下需要从一开始就分配更多内存。

在 TensorFlow 2.X 中为每个 GPU 分配内存的推荐且正确的方式是通过以下方式完成的：

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  # Restrict TensorFlow to only allocate 1GB of memory on the first GPU
  try:
    tf.config.experimental.set_virtual_device_configuration(
        gpus[0],
        [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)]

Answer 12

对于Tensorflow 2.0，这个解决方案对我有用。（TF-GPU 2.0、Windows 10、GeForce RTX 2070）

physical_devices = tf.config.experimental.list_physical_devices('GPU')
assert len(physical_devices) > 0, "Not enough GPU hardware devices available"
tf.config.experimental.set_memory_growth(physical_devices[0], True)