machine-learning - 如何在 Keras 中使用其他 GPU 和 TensorFlow 后端？

Question

这与 How to enable Keras with Theano to utilize multiple GPUs 相关但我感兴趣的是指定特定模型训练或运行的 GPU，而不是使用多个 GPU。

我的 nvidia-smi 输出如下所示:

+------------------------------------------------------+                       
| NVIDIA-SMI 361.42     Driver Version: 361.42         |                       
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  Tesla K80           Off  | 0000:03:00.0     Off |                    0 |
| N/A   38C    P0    60W / 149W |  11354MiB / 11519MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   1  Tesla K80           Off  | 0000:04:00.0     Off |                    0 |
| N/A   37C    P0    71W / 149W |    224MiB / 11519MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   2  GeForce GTX 750 Ti  Off  | 0000:06:00.0      On |                  N/A |
| 40%   29C    P8     1W /  38W |    120MiB /  2047MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

这个输出当然是在没有任何东西运行的情况下的。问题是我不确定 Keras 中如何指定在哪个 GPU 上运行。当然，使用 TensorFlow，我们可以使用 with tf.device('/cpu:1'): 范例，但我不确定这将如何与 Keras 集成。

谢谢!

Answer 1

除了在 keras.json 文件中指定 tensorflow 作为后端之外，您还可以使用环境变量 CUDA_VISIBLE_DEVICES 限制使用的 GPU 数量和/或使用特定 GPU (http://acceleware.com/blog/cudavisibledevices-masking-gpus)。您可以在此处指定要使用的 GPU。