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

0 投票

1 回答

39 浏览

提问于 2025-04-14 16:55

我正在使用一台有多个GPU的服务器，想要选择一个有足够空闲空间的GPU。使用pytorch时，我发现只需要设置一个叫“CUDA_VISIBLE_DEVICES”的环境变量，就能选择想要的GPU，但不知道为什么在tensorflow中用同样的方法却不行。

我找到了解决方案：

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，而不是我选择的那些。

在我代码的后面部分，我创建了一个tf.keras.layers.Embedding层，但它失败了，因为它自动检查GPU 0的空间，而这个GPU已经被占用了。

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

我还应该说明，我是在jupyter notebook上工作，可能这也导致了一些问题？

keras tensorflow gpu cuda_visible_devices jupyter_notebook embedding_layer device_selection resource_management

1 个回答

我用这个，但我不确定它是否能在Jupyter笔记本中正常工作。

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

这里有一个相关的问题，里面提到了一些关于Jupyter的内容：在Jupyter中设置CUDA_VISIBLE_DEVICES的Tensorflow

回答于 2025-04-14 由 Python大师

分享举报

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

1 个回答

撰写回答