Tensorflow：将 allow_growth 设置为 true 仍然会分配我所有 GPU 的内存

Question

我有多个 GPU，但我只想使用一个 GPU 进行训练。我正在使用以下选项：

config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=True)
config.gpu_options.allow_growth = True

with tf.Session(config=config) as sess:

尽管设置/使用了所有这些选项，我所有的 GPU 都会分配内存和

#processes = #GPUs

如何防止这种情况发生？

注意

1. 我不想手动设置设备，也不想设置 CUDA_VISIBLE_DEVICES，因为我希望 tensorflow 自动找到可用的最佳（空闲）GPU
2. 当我尝试启动另一个 run 时，它使用的 GPU 已被另一个 tensorflow 进程使用，即使还有其他几个空闲 GPU（除了它们上的内存分配）
3. 我在 docker 容器中运行 tensorflow：tensorflow/tensorflow:latest-devel-gpu-py

Answer 1

我可以为您提供一个方法mask_busy_gpus 在这里定义：https://github.com/yselivonchyk/TensorFlow_DCIGN/blob/master/utils.py

函数的简化版：

import subprocess as sp
import os

def mask_unused_gpus(leave_unmasked=1):
  ACCEPTABLE_AVAILABLE_MEMORY = 1024
  COMMAND = "nvidia-smi --query-gpu=memory.free --format=csv"

  try:
    _output_to_list = lambda x: x.decode('ascii').split('\n')[:-1]
    memory_free_info = _output_to_list(sp.check_output(COMMAND.split()))[1:]
    memory_free_values = [int(x.split()[0]) for i, x in enumerate(memory_free_info)]
    available_gpus = [i for i, x in enumerate(memory_free_values) if x > ACCEPTABLE_AVAILABLE_MEMORY]

    if len(available_gpus) < leave_unmasked: ValueError('Found only %d usable GPUs in the system' % len(available_gpus))
    os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(map(str, available_gpus[:leave_unmasked]))
  except Exception as e:
    print('"nvidia-smi" is probably not installed. GPUs are not masked', e)

用法：

mask_unused_gpus()
with tf.Session()...

先决条件：nvidia-smi

使用这个脚本，我正在解决下一个问题：在多 GPU 集群上，只使用单个（或任意）数量的 GPU，从而允许它们被自动分配。

脚本的缺点：如果您同时启动多个脚本，随机分配可能会导致相同的 GPU 分配，因为脚本取决于内存分配，而内存分配需要几秒钟才能启动。

Answer 2

我自己也遇到过这个问题。设置 config.gpu_options.allow_growth = True 没有成功，所有的 GPU 内存仍然被 Tensorflow 消耗。 绕过它的方法是未记录的环境变量TF_FORCE_GPU_ALLOW_GROWTH（我在 https://github.com/tensorflow/tensorflow/blob/3e21fe5faedab3a8258d344c8ad1cec2612a8aa8/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc#L25)

设置TF_FORCE_GPU_ALLOW_GROWTH=true 效果很好。

在Python代码中，可以设置

os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'