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

Question

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

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

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

我的代码:

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(强调由我添加):

Only the devices whose index is present in the sequence are visible to CUDA applications and they are enumerated in the order of the sequence. If one of the indices is invalid, only the devices whose index precedes the invalid index are visible to CUDA applications. For example, setting CUDA_VISIBLE_DEVICES to 2,1 causes device 0 to be invisible and device 2 to be enumerated before device 1. Setting CUDA_VISIBLE_DEVICES to 0,2,-1,1 causes devices 0 and 2 to be visible and device 1 to be invisible.

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

将代码更改为 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 内存。