如何在tensorflow中设置特定的gpu？答案

【问题标题】：How to set specific gpu in tensorflow?如何在tensorflow中设置特定的gpu？
【发布时间】：2016-10-16 11:47:58
【问题描述】：

我想指定 gpu 来运行我的进程。我设置如下：

import tensorflow as tf
with tf.device('/gpu:0'):
    a = tf.constant(3.0)
with tf.Session() as sess:
    while True:
        print sess.run(a)

但是它仍然在我的两个 GPU 中分配内存。

|    0      7479    C   python                         5437MiB 
|    1      7479    C   python                         5437MiB

【问题讨论】：

TensorFlow 初始化它看到的所有 GPU，你需要设置 CUDA_VISIBLE_DEVICES 来限制可见的 GPU
好的，我明白了。谢谢。

标签： tensorflow tensorflow-gpu

【解决方案1】：

有3种方法可以实现：

使用CUDA_VISIBLE_DEVICES 环境变量。通过设置环境变量CUDA_VISIBLE_DEVICES="1" 使只有设备1 可见，通过设置CUDA_VISIBLE_DEVICES="0,1" 使设备0 和1 可见。您可以在 python 中执行此操作，方法是在导入 os 包后添加一行 os.environ["CUDA_VISIBLE_DEVICES"]="0,1"。
使用with tf.device('/gpu:2') 并创建图表。然后它将使用 GPU 设备 2 运行。
使用config = tf.ConfigProto(device_count = {'GPU': 1})，然后使用sess = tf.Session(config=config)。这将使用 GPU 设备 1。

【讨论】：

这三种方法我都试过了。第二个还是不行。我的演示代码如上所示。
理论上这应该可行，但有时您必须物理切换显卡。为什么？不知道！但是当这种方法失败时它对我有用。
我不认为第三部分是完全正确的。顾名思义，device_count 只设置正在使用的设备数量，而不是哪个。来自tf源码：message ConfigProto { // Map from device type name (e.g., "CPU" or "GPU" ) to maximum // number of devices of that type to use. If a particular device // type is not found in the map, the system picks an appropriate // number. map<string, int32> device_count = 1;，见github.com/tensorflow/tensorflow/blob/r1.11/tensorflow/core/…
请注意，如果设置了tf.ConfigProto(device_count = {'GPU': 0})，所有将在CPU上运行，所以它不是gpu_id，而是gpu的数量。使用gpu_options = tf.GPUOptions(allow_growth=True, visible_device_list=str(gpu_id))

【解决方案2】：

如果没有另行通知，TF 将在每个可见 GPU 上分配所有可用内存。以下是仅使用一个（或几个）GPU 的 5 种方法。

Bash 解决方案。 在启动 python 或 jupyter notebook 之前，在终端/控制台中设置 CUDA_VISIBLE_DEVICES=0,1：

CUDA_VISIBLE_DEVICES=0,1 python script.py

Python 解决方案。在构建会话之前运行接下来的 2 行代码

import os
os.environ["CUDA_VISIBLE_DEVICES"]="0,1"

自动化解决方案。以下方法会自动检测其他脚本未使用的 GPU 设备，并为您设置 CUDA_VISIBLE_DEVICES。在构建会话之前，您必须调用 mask_unused_gpus。它将根据当前内存使用情况过滤掉 GPU。这样您就可以一次运行多个脚本实例，而无需更改代码或设置控制台参数。

功能：

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: raise 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(2)

限制：如果同时启动多个脚本，可能会导致冲突，因为在构建会话时不会立即分配内存。如果对您有问题，您可以使用原始source code: mask_busy_gpus() 中的随机版本

Tensorflow 2.0 提出了另一种方法：

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  # Restrict TensorFlow to only use the first GPU
  try:
    tf.config.experimental.set_visible_devices(gpus[0], 'GPU')
  except RuntimeError as e:
    # Visible devices must be set at program startup
    print(e)

Tensorflow/Keras 还允许指定要与会话配置一起使用的 gpu。仅当设置环境变量不是选项（即 MPI 运行）时，我才能推荐它。因为它往往是所有方法中最不可靠的，尤其是对于 keras。

config = tf.ConfigProto()
config.gpu_options.visible_device_list = "0,1"
with tf.Session(config) as sess:
#or K.set_session(tf.Session(config))

【讨论】：

所选答案在技术上是正确的，但这是处理这种情况的最佳方式。
对于 Bash 解决方案，不应该是 CUDA_VISIBLE_DEVICES=0,1 python script.py，在 env 变量 CUDA_VISIBLE_DEVICES 之前没有 $ 吗？

【解决方案3】：

我相信你需要设置CUDA_VISIBLE_DEVICES=1。或者您想使用哪个 GPU。如果您只让一个 GPU 可见，无论您将环境变量设置为什么，您都将在 tensorflow 中将其称为 /gpu:0。

有关该环境变量的更多信息：https://devblogs.nvidia.com/cuda-pro-tip-control-gpu-visibility-cuda_visible_devices/

【讨论】：

@ScottYang，试试tensorflow.org/guide/…
在其他答案中使用 os.environ["CUDA_VISIBLE_DEVICES"]="0" 允许选择在哪个 GPU 上使用用户输入。

【解决方案4】：

您可以通过在 python 脚本的开头添加以下内容来修改 GPU 选项设置：

gpu_options = tf.GPUOptions(visible_device_list="0")
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

“0”在这里是您要使用的 GPU 的名称。您可以通过在终端提示符下键入命令 nvidia-smi 来获取可用的 GPU 列表。

对于 Keras，这两个函数允许选择 CPU 或 GPU，如果是 GPU，则允许选择将使用的内存部分。

import os
from keras.backend.tensorflow_backend import set_session
import tensorflow as tf



def set_cpu_option():
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"  # see issue #152
    os.environ["CUDA_VISIBLE_DEVICES"] = ""
    os.environ["CUDA_VISIBLE_DEVICES"] = ""


def set_gpu_option(which_gpu, fraction_memory):
    config = tf.ConfigProto()
    config.gpu_options.per_process_gpu_memory_fraction = fraction_memory
    config.gpu_options.visible_device_list = which_gpu
    set_session(tf.Session(config=config))
    return

set_gpu_option("0", 0.9)
# or 
set_cpu_option()

【讨论】：

它给了我undefined name set_session, undefined name logger

【解决方案5】：

在我的多核 gpu 设置中，我见过的最优雅、最简洁的方式是：

import os
os.environ["CUDA_VISIBLE_DEVICES"]="1"
tf_device='/gpu:0'

这会将任务分配给 gpu 设备 1。

同样，做一些事情就行了：

import os 
os.environ["CUDA_VISIBLE_DEVICES"]="2"
tf_device='/gpu:0'

os.environ 命令可以看作是一种只让你打算运行代码的 GPU 设备暴露的方法。第二个命令只选择您指定的第一个可用设备。

【讨论】：

不行，设置同一个进程的环境变量不优雅。您应该在运行进程之前设置它们，例如：CUDA_VISIBLE_DEVICES=XYZ python script.py

【解决方案6】：

import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"   # see issue #152
os.environ["CUDA_VISIBLE_DEVICES"]="3"

唯一对我有用的方法是在进程中为池中的每个进程分配特定的 GPU。

【讨论】：

【解决方案7】：

def set_specific_gpu(ID):
    gpus_all_physical_list = tf.config.list_physical_devices(device_type='GPU')    
    tf.config.set_visible_devices(gpus_all_physical_list[ID], 'GPU')

参考https://www.tensorflow.org/api_docs/python/tf/config/list_physical_devices

【讨论】：

【解决方案8】：

import tensorflow as tf


gpu_number = 2 #### GPU number 
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    tf.config.experimental.set_visible_devices(gpus[gpu_number], 'GPU') 
    logical_gpus = tf.config.experimental.list_logical_devices('GPU')
    print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPU")

【讨论】：