如何在张量流中设置特定的GPU?
tensorflow
9
0

我想指定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 
参考资料:
Stack Overflow
收藏
评论
共 4 个回答
高赞 时间 活跃

有3种方法可以实现此目的:

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

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

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

收藏
评论

您可以通过在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,这2个功能允许选择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()
收藏
评论

我相信您需要设置CUDA_VISIBLE_DEVICES=1 。或您要使用的GPU。如果仅使一个GPU可见,则无论将环境变量设置为什么,都将在tensorflow中将其称为/gpu:0

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

收藏
评论

除非另有说明,否则TF将在每个可见GPU上分配所有可用内存。坚持使用一个(或几个)GPU的5种方法。

重击解决方案。在启动python或jupyter笔记本之前,在终端/控制台中设置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)

局限性:如果一次启动多个脚本可能会导致冲突,因为构造会话时不会立即分配内存。如果您遇到问题,可以使用原始版本中的随机版本: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))
收藏
评论
新手导航
  • 社区规范
  • 提出问题
  • 进行投票
  • 个人资料
  • 优化问题
  • 回答问题

关于我们

常见问题

内容许可

联系我们

@2020 AskGo
京ICP备20001863号