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
如何防止这种情况发生？

注意

我不想使用手动设置设备，也不想设置CUDA_VISIBLE_DEVICES，因为我想让tensorflow自动找到最佳的(空闲)GPU可用

当我尝试启动另一个run时，它使用的是另一个tensorflow进程已经使用的GPU，即使还有其他几个可用的GPU(除了其上的内存分配)也是如此。

我正在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分配，因为脚本取决于内存分配，并且内存分配需要几秒钟的时间。