Tensorflow，如何确保一切都纯粹在GPU中运行？

Question

我的理解是，tensorflow将采用计算图并在gpu和cpu上运行它以达到最大使用率。如何确保我的代码仅在gpu上运行？

Answer 1

您可以将所有图形构建代码包装在tf.device()块中，如下所示：

with tf.device("/gpu:0"):
    # [Build graph in here.]

使用此注释，放置算法会将所有操作放在GPU设备上，如果任何操作没有GPU实现，则会失败。 （如果没有这个注释，正如您所提到的，它将回退到在CPU设备上没有GPU实现的情况下进行操作。）