如何在张量流中设置特定的 gpu？

Question

我想为运行我的进程指定 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

Answer 1

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

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

Answer 2

有 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.

Answer 3

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

Bash 解决方案。 在开始 python 或 jupyter notebook 之前，在 terminal/console 中设置 CUDA_VISIBLE_DEVICES=0,1：

CUDA_VISIBLE_DEVICES=0,1 python script.py

Python 解决方案。 运行构造会话之前的下两行代码

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

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

Answer 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 设备的方法。第二个命令只选择您指定的第一个可用设备。

Answer 6

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

在进程中唯一对我有用的东西是将特定的 GPU 分配给池中的每个进程。

Answer 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

Answer 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")

如何在张量流中设置特定的 gpu？

How to set specific gpu in tensorflow?

tensorflow

tensorflow-gpu