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

我相信你需要设置CUDA_VISIBLE_DEVICES=1。或者您想要使用哪种GPU。如果只显示一个GPU，则无论您将环境变量设置为什么，都会将其称为张量流中的/gpu:0。

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

Answer 2

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

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

Answer 3

如果没有另外说明，TF将在每个可见GPU上分配所有可用内存。这里有四种方法可以坚持只有一个（或几个）GPU。

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

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)

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