使用 Keras 和 Tensorflow 指定要使用的 3 个 GPU（一台机器上的 4 个）

Question

我正在我的具有 4 个 GPU 的机器上执行深度学习。在训练期间，第三个 GPU 始终丢失（错误出现“GPU 丢失”并且日志表明它是这个特定的 GPU）。我假设它是一个散热问题，并且 GPU 正在被取消。

在解决此硬件问题之前，我想继续使用 3 个 GPU（'/gpu:0'、'/gpu:1'、'/gpu:3'）。有没有办法在 Keras 中指定这些是我想要使用的 GPU（或者，忽略 '/gpu:2'）？

我已经看到很多关于指定 GPU 与 CPU 使用率以及在多 GPU 机器上指定一个 GPU 的内容，但没有看到这个特定问题（隔离了许多特定的 GPU）。非常感谢。

Answer 1

您可以尝试使用 CUDA_VISIBLE_DEVICES 环境

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

可能在导入 keras/tf 之前设置。