Tensorflow 在使用 tf.device('/cpu:0') 时分配 GPU 内存

Question

系统信息：1.1.0、GPU、Windows、Python 3.5，代码在 ipython 控制台中运行。

我正在尝试运行两个不同的 Tensorflow 会话，一个在 GPU 上（执行一些批处理工作），一个在我用于快速测试的 CPU 上，而另一个在运行。

问题在于，当我生成指定 with tf.device('/cpu:0') 的第二个会话时，该会话尝试分配 GPU 内存并导致我的另一个会话崩溃。

我的代码：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = ""
import time

import tensorflow as tf

with tf.device('/cpu:0'):
  with tf.Session() as sess:
    # Here 6 GBs of GPU RAM are allocated.
    time.sleep(5)

如何强制 Tensorflow 忽略 GPU？

更新：

根据@Nicolas 的评论建议，我看了at this answer 并跑了

import os
os.environ["CUDA_VISIBLE_DEVICES"] = ""
import tensorflow as tf

from tensorflow.python.client import device_lib
print(device_lib.list_local_devices())

哪个打印：

[name: "/cpu:0"
device_type: "CPU"
memory_limit: 268435456
locality {
}
incarnation: 2215045474989189346
, name: "/gpu:0"
device_type: "GPU"
memory_limit: 6787871540
locality {
  bus_id: 1
}
incarnation: 13663872143510826785
physical_device_desc: "device: 0, name: GeForce GTX 1080, pci bus id: 0000:02:00.0"
]

在我看来，即使我明确告诉脚本忽略任何 CUDA 设备，它仍然会找到并使用它们。这可能是 TF 1.1 的错误吗？

Answer 1

事实证明，将CUDA_VISIBLE_DEVICES 设置为空字符串不会屏蔽脚本可见的 CUDA 设备。

来自documentation of CUDA_VISIBLE_DEVICES（重点由我添加）：

只有序列中存在索引的设备对 CUDA 应用程序，它们按顺序列举 顺序。 如果其中一个索引无效，则只有其设备 无效索引之前的索引对 CUDA 应用程序可见。为了 例如，将 CUDA_VISIBLE_DEVICES 设置为 2,1 会导致设备 0 不可见，设备 2 在设备 1 之前枚举。设置 CUDA_VISIBLE_DEVICES 到 0,2,-1,1 导致设备 0 和 2 可见 并且设备 1 不可见。

空字符串似乎曾经被处理为“不存在有效设备”，但含义有所改变，因为文档中没有提及。

将代码更改为os.environ["CUDA_VISIBLE_DEVICES"] = "-1" 可以解决问题。跑步

import os
os.environ["CUDA_VISIBLE_DEVICES"]="-1"    
import tensorflow as tf

from tensorflow.python.client import device_lib
print(device_lib.list_local_devices())

现在打印

[name: "/cpu:0"
device_type: "CPU"
memory_limit: 268435456
locality {
}
incarnation: 14097726166554667970
]

并且实例化 tf.Session 不再占用 GPU 内存。

Answer 2

您介意尝试这些配置选项之一吗？

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
# or config.gpu_options.per_process_gpu_memory_fraction = 0.0
with tf.Session(config=config) as sess:
    ...

根据documentation，它应该可以帮助您管理此特定会话的 GPU 内存，因此您的第二个会话应该能够在 GPU 上运行。

编辑：根据这个answer 你也应该试试这个：

import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"   # see issue #152
os.environ["CUDA_VISIBLE_DEVICES"]="-1"