50

我想指定 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)

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

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

7 回答 7

63

有3种方法可以实现这一点：

  1. 使用CUDA_VISIBLE_DEVICES环境变量。通过设置环境变量CUDA_VISIBLE_DEVICES="1"使设备 1 可见，通过设置CUDA_VISIBLE_DEVICES="0,1"使设备 0 和 1 可见。您可以在 python 中通过os.environ["CUDA_VISIBLE_DEVICES"]="0,1"在导入os包后添加一行来执行此操作。

  2. 使用with tf.device('/gpu:2')和创建图表。然后它将使用 GPU 设备 2 运行。

  3. 使用config = tf.ConfigProto(device_count = {'GPU': 1})然后sess = tf.Session(config=config)。这将使用 GPU 设备 1。

于 2017-06-30T13:46:53.087 回答
40

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

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

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)

限制：如果一次启动多个脚本可能会导致冲突，因为在构建会话时不会立即分配内存。如果这对您有问题，您可以使用原始源代码中的随机版本：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))
于 2017-12-27T20:19:54.647 回答
32

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

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

于 2016-12-07T18:44:04.540 回答
5

您可以通过在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()
于 2019-02-13T09:59:20.337 回答
1

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

在此处输入图像描述

于 2020-05-07T11:45:38.273 回答
0 
import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"   # see issue #152
os.environ["CUDA_VISIBLE_DEVICES"]="3"

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

于 2021-01-17T17:34:44.113 回答
0 
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

于 2021-11-10T01:21:23.173 回答