python - tensorflow-gpu 库是否自动在 GPU 上运行 tensorflow 代码(非 GPU)？

Question

我是否必须自定义为非 GPU Tensorflow 库编写的代码以适应tensorflow-gpu 库？

我有一个 GPU，想运行仅为非 GPU tensorflow 库编写的 Python 代码。我可以简单地安装 tensorflow GPU模块并运行代码吗？或者为了在 GPU 上运行代码，我必须进行任何代码更改吗？

Answer 1

如果您已经安装了 tensorflow-gpu 并且一切运行正常，默认情况下，它应该在您的 GPU 上运行，除非:

• 您在代码中明确指出该图应在 CPU 上运行，例如 with tf.device('/cpu:0')，
• 您可以使用 os.environ['CUDA_VISIBLE_DEVICES'] = '' 从 tensorflow 中隐藏 GPU。

您可以使用 ( credit to this answer ) 检查 TensorFlow 是否找到您的 GPU

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

此外，如果您有 NVIDIA GPU，您可以使用 nvidia-smi 检查 GPU 的使用情况，并查看 TensorFlow 确实正在使用您的 GPU。

现在，如果您想分配特定的 GPU 和/或限制图形可以在 GPU 内存上分配的内存，您可能需要查看 here了解更多详情。