Keras 多 GPU：指定显式 GPU ID

Question

通过查看 tensorflow GitHub 存储库中的文件 keras/utils/multi_gpu_utils.py，我可以看到，如果您指定要使用 x GPU，它会自动分配 GPU来自 range(x) 的 ID，即 0, 1, 2, ..., x - 1.

我需要使用 GPU 4, 5, 6 ,7，因为前 4 个 GPU 已经在处理另一项任务。有没有办法指定它？

Answer 1

在python中你可以使用

import os
os.environ["CUDA_VISIBLE_DEVICES"]="0,1"

或者在bash中设置CUDA_VISIBLE_DEVICES=0,1，然后再开始python脚本

你也可以参考我的回答来自动化这个过程。