tensorflow - 如何在 Python 界面中指定 GPU 的数量？

Question

当我启动tensorflow Session时，它会立即在我的 dekstop（设备 0 和 2）中的两个 GPU 上分配 100% 的内存。

如何使用Python API将其限制为仅 1 个 GPU？

Answer 1

例如，将其限制为仅使用 GPU #0 的一种粗略方法是定义：

export CUDA_VISIBLE_DEVICES=0

Answer 2

如果您希望仅对特定设备执行某些操作，您可以执行以下操作：

with tf.device('/gpu:0'):
   ...
   # all your operations

其中设备以下列方式指定："/cpu:0", "/gpu:0", "/gpu:1". 我建议您阅读有关使用 GPU 设备的主题中的更多详细信息。

如果你想看一些教程，看看cifar10 教程。