如何限制tensorflow GPU内存使用？

Question

我在 Ubuntu 18.04 中使用了tensorflow-gpu 1.13.1，并在 Nvidia GeForce RTX 2070 上使用了 CUDA 10.0（驱动程序版本：415.27）。

下面的代码用于管理张量流内存使用情况。我有大约 8Gb GPU 内存，因此 TensorFlow 不得分配超过 1Gb 的 GPU 内存。但是，当我使用命令查看内存使用情况时nvidia-smi，我发现它使用了 ~1.5 Gb，尽管我使用 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 仍然以相同的方式分配内存，独立于此标志值。而且看起来也很奇怪

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：我tensorflow/core/common_runtime/gpu/gpu_device.cc:1053]创建了TensorFlow设备（/job:localhost/replica:0/task:0/device:GPU:0 with 15127 MB内存）->物理GPU（设备：0，名称：Tesla P100-PCIE-16GB，pci总线id：0000：03：00.0，计算能力：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) 的总和一致。