Question

使用TensorFlow的Keras顺序API时，是否有任何方法可以强制在某种硬件上训练我的模型？我的理解是，如果有要使用的GPU（并且我安装了tensorflow-gpu），则默认情况下，我将在GPU上进行训练。

我是否必须切换到其他API才能更好地控制模型的部署位置？

Answer 1

我是keras用户，我在ubuntu上工作。我指定一个特定的GPU，如下所示：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

其中0是GPU的数量。默认情况下，如果您的计算机上有多个GPU，则tensorflow使用第一个GPU（其编号为0）。您可以通过在终端上键入以下命令来获取GPU的信息：

nvidia-smi

或

watch -n 1 -d nvidia-smi

如果您想每秒刷新一次终端。下图显示了我的GPU的信息，并用红色框圈出了它的数量。