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

Posted 2023-02-23

【中文标题】Tensorflow：将 allow_growth 设置为 true 仍然会分配我所有 GPU 的内存

【英文标题】：Tensorflow: Setting allow_growth to true does still allocate memory of all my GPUs

【发布时间】：2018-06-03 06:54:41

【问题描述】：

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

【问题讨论】：

看起来很奇怪。能否请您尝试发布完整代码和您正在使用的 TF 版本？

您是否尝试指定初始内存分数？ gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333) sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))***.com/questions/34199233/…

完整的代码超过5个脚本，所以我不能给你看完整的代码，很遗憾，但我想我的意思说清楚了？或者你有什么特别想看的吗？我已经添加了我正在开发的 tensorflow 版本。 @MatanHugi

不，我没有，但我相信这不会帮助我解决我的问题。 @yauheni_selivonchyk

Tensorflow 没有找到可用的野兽（空闲）GPU 的逻辑。

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

【参考方案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'

【讨论】：

太棒了！但请注意，如果您使用 SLURM，则需要在 sbatch 脚本中添加“export TF_FORCE_GPU_ALLOW_GROWTH=true”，因为您将无法在 python 中设置环境变量！