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它使用另一个 tensorflow 进程已经使用的相同 GPU，即使还有其他几个可用的 GPU（除了它们的内存分配）
3. 我在 docker 容器中运行 tensorflow： tensorflow/tensorflow:latest-devel-gpu-py

Answer 1

我自己也有这个问题。设置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'

Answer 2

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 分配，因为脚本取决于内存分配，而内存分配需要几秒钟才能启动。