运行tensorflow时如何指定使用哪个GPU？答案

【问题标题】：How to specify which GPU to use when running tensorflow?运行tensorflow时如何指定使用哪个GPU？
【发布时间】：2020-01-14 00:32:46
【问题描述】：

我们在实验室中有一台 DGX-1。

我看到许多任务在不同的 GPU 上运行。

对于 MLperf docker 应用程序，我可以使用 NV_GPU=x 来分配要使用的 GPU。

但是，我有一个 python Keras/TensorFlow 代码，我用同样的方式，加载不去指定的 GPU。

【问题讨论】：

【解决方案1】：

您可以使用CUDA_VISIBLE_DEVICES 指定您的模型要使用的 GPU：

import os 
os.environ["CUDA_VISIBLE_DEVICES"] = 0,1  #Will assign GPUs 0 and 1 to the model

【讨论】：