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

Question

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

理想情况下，Keras应该能够确定哪个GPU正在繁忙地训练模型，然后使用另一个GPU来训练其他模型。然而，事实并非如此。似乎默认情况下Keras只使用第一个GPU（因为第二个GPU的Volatile GPU-Util始终为0%）。

- NickZeng

Answer 1

也许与我之前的问题重复了

这有点复杂。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，并将适当的值设置为变量。

- Daniel GL

---

太好了！谢谢你！这正是我想要的！ - NickZeng

这对于一个做很多事情的程序来说并不实用...还有其他方法吗？ - rjurney

如果有人能找到更简单的方法，我会很高兴！ - Daniel GL

Answer 2

如果您正在使用培训脚本，您可以在调用脚本之前在命令行中设置它。

CUDA_VISIBLE_DEVICES=1 python train.py 

- Omri Bahat Treidel

Answer 3

如果你想在云GPU上训练模型（例如来自AWS的GPU实例），可以尝试这个库：

!pip install aibro==0.0.45 --extra-index-url https://test.pypi.org/simple

from aibro.train import fit
machine_id = 'g4dn.4xlarge' #instance name on AWS
job_id, trained_model, history = fit(
    model=model,
    train_X=train_X,
    train_Y=train_Y,
    validation_data=(validation_X, validation_Y),
    machine_id=machine_id
)

教程：https://colab.research.google.com/drive/19sXZ4kbic681zqEsrl_CZfB5cegUwuIB#scrollTo=ERqoHEaamR1Y

- Yuqi Lee