tensorflow - Keras multi-gpu：指定显式 GPU id

Question

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

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

标签： tensorflowkerasmulti-gpu

Answer 1

在python中你可以使用

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

或者CUDA_VISIBLE_DEVICES=0,1在启动 python 脚本之前在 bash 中设置

您也可以在此处参考我的答案以自动化此过程。