tensorflow - 以编程方式获取 Tensorflow 操作的设备布局

我想编写一个断言，以确保我的图中的某些操作将在特定设备上运行。如何以编程方式确定操作的设备布局，以便我可以编写这样的断言？

最佳答案

您可以使用

确保操作在特定设备上运行

with tf.device('/gpu:0'):

在操作定义之前(有关详细信息，请参阅 here)。

编辑:

每个可用的 GPU 都有自己的索引:“/gpu:0”、“/gpu:1”、“/gpu:2”等。这样您就可以将特定操作绑定(bind)到特定 GPU。当您导入tensorflow时，它会打印出哪些GPU可用以及它们分配的索引。

(例如，它打印:创建 TensorFlow 设备 (/gpu:0) -> (设备: 0, 名称: GeForce GTX 1070, pci 总线 id: ...))

关于tensorflow - 以编程方式获取 Tensorflow 操作的设备布局，我们在Stack Overflow上找到一个类似的问题： https://stackoverflow.com/questions/43171456/