Tensorflow:动态调用具有足够可用内存的GPU

Question

我的桌面有两个gpus，可以运行/gpu:0或/gpu:1规范的Tensorflow。但是，如果我不指定运行代码的图形处理器，Tensorflow将默认调用/gpu:0，这一点我们都知道。

现在我想设置系统，以便它可以根据每个gpu的空闲内存动态分配gpu。例如，如果一个脚本没有指定运行代码的gpu，系统将首先为它分配/gpu:0；然后，如果另一个脚本现在运行，它将检查/gpu:0是否有足够的空闲内存。如果是，则继续为其分配/gpu:0，否则将为其分配/gpu:1。我怎样才能做到这一点？

后续:我相信上面的问题可能与GPU的虚拟化问题有关。也就是说，如果我可以将桌面中的多个GPU虚拟化为一个GPU，我就可以得到我想要的东西。因此，除了Tensorflow的任何设置方法之外，也欢迎任何关于虚拟化的想法。

Answer 1

TensorFlow通常假设它不会与任何人共享图形处理器，所以我在TensorFlow内部看不到这样做的方法。但是，您可以从外部执行以下操作--调用nvidia-smi的shell脚本，解析出具有更多内存的GPU，然后设置"CUDA_VISIBLE_DEVICES=k“并调用TensorFlow脚本

Answer 2

灵感来自：How to set specific gpu in tensorflow?

def leave_gpu_with_most_free_ram():
  try:
    command = "nvidia-smi --query-gpu=memory.free --format=csv"
    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)]
    least_busy_idx = memory_free_values.index(max(memory_free_values))

    # update CUDA variable
    gpus =[least_busy_idx]
    setting = ','.join(map(str, gpus))
    os.environ["CUDA_VISIBLE_DEVICES"] = setting
    print('Left next %d GPU(s) unmasked: [%s] (from %s available)'
          % (leave_unmasked, setting, str(available_gpus)))
  except FileNotFoundError as e:
    print('"nvidia-smi" is probably not installed. GPUs are not masked')
    print(e)
  except sp.CalledProcessError as e:
    print("Error on GPU masking:\n", e.output)

在导入tensorflow之前添加对此函数的调用