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

我自己遇到了这个问题。设置config.gpu_options.allow_growth = True并没有解决问题，Tensorflow仍然消耗了所有GPU内存。解决方法是未记录的环境变量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。

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()...

Prerequesities： nvidia-smi

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

脚本的缺点：如果您一次启动多个脚本，随机分配可能会导致相同的GPU分配，因为脚本依赖于内存分配，内存分配需要几秒钟才能完成。

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

2 个答案: