如何限制TensorFlow GPU内存使用？

Question

我在Ubuntu 18.04上使用CUDA 10.0和Nvidia GeForce RTX 2070 (驱动版本：415.27)时，使用了tensorflow-gpu 1.13.1。像下面的代码一样用于管理tensorflow内存使用量。我的GPU内存约为8Gb，所以tensorflow不能分配超过1Gb的GPU内存。但是当我使用nvidia-smi命令查看内存使用情况时，发现它实际使用了约1.5Gb，尽管我已经通过GPUOptions限制了内存数量。

memory_config = tf.ConfigProto(gpu_options=tf.GPUOptions(per_process_gpu_memory_fraction=0.12))

memory_config.gpu_options.allow_growth = False

with tf.Session(graph=graph, config=memory_config) as sess:
    output_dict = sess.run(tensor_dict,
                           feed_dict={image_tensor: np.expand_dims(image, 0)})

为什么会这样？我应该如何避免或者至少计算每个会话所需的内存？我需要对每个进程进行严格限制，因为我有几个并行实例，它们具有不同的会话，所以我需要确保没有资源竞争。

顺便说一下，我已经尝试将memory_config.gpu_options.allow_growth设置为False，但是它没有产生任何影响。Tensorflow仍然独立于此标志值分配内存。这也似乎很奇怪。

- Silicon

Answer 1

解决方案尝试使用gpu_options.allow_growth = True，查看在tf.Session创建时消耗了多少默认内存。该内存将始终被分配，而不考虑值。

根据您的结果，它应该小于500MB。因此，如果您希望每个进程真正地拥有1GB的内存，请计算：

(1GB minus default memory)/total_memory

原因

当您创建一个 tf.Session 时，无论您的配置如何，Tensorflow 设备都会在 GPU 上创建。而这个 设备 需要一些最小内存。

import tensorflow as tf

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

给定allow_growth=True，应该不会有gpu分配。但实际情况是：

2019-04-05 18:44:43.460479：I tensorflow/core/common_runtime/gpu/gpu_device.cc:1053] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 15127 MB memory) -> physical GPU (device: 0, name: Tesla P100-PCIE-16GB, pci bus id: 0000:03:00.0, compute capability: 6.0)

其占用的内存只占了一小部分（根据我的过去经验，这个量会因gpu型号而异）。 注意：设置allow_growth占用的内存几乎与设置per_process_gpu_memory=0.00001相同，但后者无法正确创建会话。

在这种情况下，它是345MB:

这就是你所遇到的偏移量。让我们看一下per_process_gpu_memory的情况：

conf = tf.ConfigProto()
conf.gpu_options.per_process_gpu_memory_fraction=0.1
session = tf.Session(config=conf)

由于这个GPU有16,276MB的内存，设置per_process_gpu_memory_fraction = 0.1 可能 让你认为只有大约1,627MB被分配。但事实上：

1,971MB被分配，但是这恰好与默认内存（345MB）和预期内存（1,627MB）之和相同。

- Chris

---

有没有驱动级别的解决方案，不依赖于框架？基本上说，防止所有CUDA应用程序默认占用所有GPU内存？ - Florin Andrei