如何防止 tensorflow 分配整个 GPU 内存？

Question

我在一个共享计算资源的环境中工作，即我们有几台服务器机器，每台机器都配备了几个 Nvidia Titan X GPU。

对于中小型模型，Titan X 的 12 GB 通常足以让 2–3 人在同一 GPU 上同时进行运行训练。如果模型足够小以至于单个模型无法充分利用 GPU 的所有计算单元，与运行一个接一个地训练过程相比，这实际上可以带来加速。即使在并发访问 GPU 确实会减慢个人训练时间的情况下，让多个用户同时在 GPU 上训练的灵活性仍然很好。

TensorFlow 的问题在于，默认情况下，它会在启动时分配全部可用 GPU 内存。即使对于一个小型的两层神经网络，我也看到所有 12 GB 的 GPU 内存都用完了。

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

Answer 1

您可以在构建 tf.Session by passing a tf.GPUOptions 作为可选 config 参数的一部分时设置要分配的 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

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

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

Answer 4

这是本书的节选 Deep Learning with TensorFlow

In some cases it is desirable for the process to only allocate a subset of the available memory, or to only grow the memory usage as it is needed by the process. TensorFlow provides two configuration options on the session to control this. The first is the allow_growth option, which attempts to allocate only as much GPU memory based on runtime allocations, it starts out allocating very little memory, and as sessions get run and more GPU memory is needed, we extend the GPU memory region needed by the TensorFlow process.

1) 允许增长：（更灵活）

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

第二种方法是per_process_gpu_memory_fraction选项，它决定了each可见GPU应该分配的内存总量的比例。 注意：不需要释放内存，完成后甚至会加剧内存碎片。

2)分配固定内存:

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

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

注：这只有在您真的想绑定 TensorFlow 进程上可用的 GPU 内存量时才有用。

Answer 5

以上所有答案都假定执行 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)

编辑：2018 年 4 月 11 日 例如，如果你要使用 tf.contrib.gan.train，那么你可以使用类似于下面的东西：

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

Answer 6

我尝试在 voc 数据集上训练 unet，但由于图像太大，内存不足。我尝试了以上所有技巧，甚至尝试了批量大小==1，但没有任何改进。有时 TensorFlow 版本也会导致内存问题。尝试使用

pip install tensorflow-gpu==1.8.0

Answer 7

好吧，我是 tensorflow 的新手，我有 Geforce 740m 或带有 2GB 内存的 GPU，我是运行 mnist 手写的母语示例，训练数据包含 38700 张图像和 4300 张测试图像并试图使用以下代码获得精度，召回，F1，因为 sklearn 没有给我精确的结果。一旦我将它添加到我现有的代码中，我就开始出现 GPU 错误。

TP = tf.count_nonzero(predicted * actual)
TN = tf.count_nonzero((predicted - 1) * (actual - 1))
FP = tf.count_nonzero(predicted * (actual - 1))
FN = tf.count_nonzero((predicted - 1) * actual)

prec = TP / (TP + FP)
recall = TP / (TP + FN)
f1 = 2 * prec * recall / (prec + recall)

加上我猜我的模型很重，我在 147、148 个时期后出现内存错误，然后我想为什么不为任务创建函数所以我不知道它在 tensrorflow 中是否以这种方式工作，但我想如果使用局部变量并且超出范围时它可能会释放内存并且我在模块中定义了上述用于训练和测试的元素，我能够毫无问题地实现 10000 个纪元，我希望这会有所帮助..

Answer 8

对于 TensorFlow 2.0 和 2.1 (docs):

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

对于 TensorFlow 2.2+ (docs):

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
  tf.config.experimental.set_memory_growth(gpu, True)

文档还列出了更多方法：

将环境变量 TF_FORCE_GPU_ALLOW_GROWTH 设置为 true。
使用 tf.config.experimental.set_virtual_device_configuration 对虚拟 GPU 设备设置硬限制。

Answer 9

您可以使用

TF_FORCE_GPU_ALLOW_GROWTH=true

在你的环境变量中。

在tensorflow代码中：

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 10

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

参考文献：

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

对于 Tensorflow 2.0 Alpha，请参阅：

Answer 11

# allocate 60% of GPU memory 
from keras.backend.tensorflow_backend import set_session
import tensorflow as tf 
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.6
set_session(tf.Session(config=config))

Answer 12

对于 Tensorflow 2.0，这个 this solution 对我有用。 (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)

Answer 13

对于 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 14

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

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

Answer 15

这段代码对我有用：

import tensorflow as tf
config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.compat.v1.InteractiveSession(config=config)

Answer 16

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

tf.config.experimental.set_memory_growth 方法确实适用于在 allocation/preprocessing 期间实现动态增长。尽管如此，人们可能希望从一开始就分配 specific-upper 限制 GPU 内存。

分配特定 GPU 内存背后的逻辑也是为了防止训练期间出现 OOM 内存。例如，如果一个训练同时打开 video-memory 消费 Chrome tabs/any 其他视频消费过程， 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)]

如何防止 tensorflow 分配整个 GPU 内存？

How to prevent tensorflow from allocating the totality of a GPU memory?

python

tensorflow

nvidia-titan

tensorflow2.0

tensorflow2.x

Tensorflow 2.0 Beta 和（可能）超越