python - 如何让 Keras 在特定 GPU 上训练模型？

Question

我所在的机构有一台带有 2 个 GPU 的共享服务器。假设有两个团队成员都想同时训练一个模型，那么他们如何让 Keras 在特定的 GPU 上训练他们的模型以避免资源冲突？

理想情况下，Keras 应该找出当前哪个 GPU 正在忙于训练模型，然后使用另一个 GPU 来训练另一个模型。然而，情况似乎并非如此。似乎默认情况下，Keras 只使用第一个 GPU(因为第二个 GPU 的 Volatile GPU-Util 始终为 0%)。

Answer 1

Possibly duplicate with my previous question
它有点复杂。 Keras 将使用两个 GPU 中的内存，尽管默认情况下它只会使用一个 GPU。查询 keras.utils.multi_gpu_model用于使用多个 GPU。
我通过使用环境变量 CUDA_VISIBLE_DEVICES 选择 GPU 找到了解决方案。
您可以在导入 keras 或 tensorflow 之前手动添加它以选择您的 gpu

os.environ["CUDA_VISIBLE_DEVICES"]="0" # first gpu
os.environ["CUDA_VISIBLE_DEVICES"]="1" # second gpu
os.environ["CUDA_VISIBLE_DEVICES"] = "-1" # runs in cpu

为了让它自动生成，我做了一个解析 nvidia-smi 的函数并自动检测哪个 GPU 已被使用，并为变量设置适当的值。