如何在张量流中设置特定的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)

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

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

Answer 1

有3种方法可以实现这一目标:

1. 使用CUDA_VISIBLE_DEVICES环境变量.通过设置设置环境变量CUDA_VISIBLE_DEVICES="1",仅使设备1可见,并通过设置CUDA_VISIBLE_DEVICES="0,1"使设备0和1可见.

2. 使用os.environ["CUDA_VISIBLE_DEVICES"]="0,1"和创建图表.然后它将使用GPU设备2运行.

3. 使用os然后with tf.device('/gpu:2').这将使用GPU设备1.

Answer 2

我相信你需要设定CUDA_VISIBLE_DEVICES=1.或者您想要使用哪种GPU.如果只显示一个GPU,则/gpu:0无论您将环境变量设置为什么,都将引用它.

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

Answer 3

如果没有另外说明,TF将在每个可见GPU上分配所有可用内存.

你可以设置CUDA_VISIBLE_DEVICES=0,1它是否是一次性的东西.

如果您正在使用群集而您不想跟踪哪个GPU正忙,并且在构建会话之前手动输入您可以调用下一个方法的信息.它将过滤掉已经使用的GPU(没有太多可用内存)并为您设置CUDA_VISIBLE_DEVICES.

功能:

$CUDA_VISIBLE_DEVICES=0,1 python script.py

限制:如果您一次启动多个脚本,它仍然会导致冲突,因为在构造会话时不会立即分配内存.如果您遇到问题,可以使用我原始源代码中的随机版本:mask_busy_gpus()

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