Tensorflow仅使用GPU0

问题描述

测试代码：

import tensorflow as tf

tf.debugging.set_log_device_placement(True)
tf.config.set_soft_device_placement(False)

device = '/GPU:1'
with tf.device(device):
    print("Input device : " + device)
    a = tf.constant([[1.0,2.0,3.0],[4.0,5.0,6.0]])
    b = tf.constant([[1.0,2.0],[3.0,4.0],[5.0,6.0]])

c = tf.matmul(a,b)
print(c)

当前行为我正在尝试为GPU1定义代码范围。但是，该代码在GPU 0上执行。例如，上面给出的测试代码的输出是

Input device : /GPU:1
2020-08-29 00:15:42.093921: I tensorflow/core/common_runtime/eager/execute.cc:574] Executing op MatMul in device /job:localhost/replica:0/task:0/device:GPU:0
2020-08-29 00:15:42.094117: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libcublas.so.10.0
tf.Tensor(
[[22. 28.]
[49. 64.]],shape=(2,2),dtype=float32)

GPUS 0和GPU 1都可用

预期行为该代码应在GPU 1上执行

系统信息

OS Platform and distribution (e.g.,Linux Ubuntu 16.04): Linux
Mobile device (e.g. iPhone 8,Pixel 2,Samsung galaxy) if the issue happens on mobile device: na
TensorFlow installed from (source or binary): conda
TensorFlow version (use command below): 2.0.0
Python version: 3.6.5
CUDA/cuDNN version: 7.6.5

GPU模型和内存：

name: "/device:GPU:0"
device_type: "GPU"
memory_limit: 32040203060
physical_device_desc: "device: 0,name: Tesla V100-SXM2-32GB,pci bus id: 0000:61:00.0,compute capability: 7.0"
name: "/device:GPU:1"
device_type: "GPU"
memory_limit: 32040203060
physical_device_desc: "device: 1,pci bus id: 0000:62:00.0,compute capability: 7.0"

解决方法

通过设置环境变量，一些简单的行可以处理您甚至不希望TF知道其他GPU的情况。例如

import os
os.environ["CUDA_VISIBLE_DEVICES"]="3" # sets only GPU3 to be visible
os.environ["TF_FORCE_GPU_ALLOW_GROWTH"]="true"

还请注意，TF教程的这一行而不是您使用的行：

with tf.device('/device:GPU:2'): # It includes /device:... in the device name

也许首先尝试第二个选项，但是最上面的一个将锁定您的会话，使其仅使用一个GPU。您应该在使用TF之前先调用它。

您可以在运行python文件时从终端中选择要使用的gpu：

CUDA_VISIBLE_DEVICES=1 python app.py

用于计算c的语句是缩进的，因此它将计算分配给默认计算机GPU0。

cudnn gpu python tensorflow