如何限制TensorFlow中的可见设备？

Question

我正在一台具有多个 GPU 的服务器上工作，我想选择一个具有足够可用空间的服务器。 使用 pytorch，我发现只需将操作系统环境变量“CUDA_VISIBLE_DEVICES”设置为所需的 GPU 即可完成此操作，但由于某种原因，相同的方法不适用于张量流。

我找到了这个解决方案：

config = tf.compat.v1.ConfigProto()
config.gpu_options.visible_device_list = "desired_GPU_no"
tf.compat.v1.keras.backend.set_session(tf.compat.v1.Session(config=config))

不幸的是，它对我不起作用，因为运行

tf.config.list_physical_devices('GPU')

显示所有 GPU，而不是我选择的 GPU。

后来在我的代码中，我构建了一个 tf.keras.layers.Embedding 层，但它失败了，因为它会自动检查 GPU 0 上的空间，该空间已被占用。

有没有更可靠的方法来选择我需要的GPU，或者我当前的解决方案是否做错了什么？

我还应该指定我正在使用 jupyter 笔记本，也许这会导致问题？

Answer 1

我使用这个，但我无法确认它是否适用于 Jupyter 笔记本。

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0" # or "1", "2", "3", ...