以编程方式获取Tensorflow op的设备放置

Question

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

Answer 1

您可以使用

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

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

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

修改

每个可用的gpu都有自己的索引：'/ gpu：0'，'/ gpu：1'，'/ gpu：2'等。这样你可以将特定的操作绑定到特定的gpus。导入tensorflow时，会输出可用的gpus以及分配的索引。

（例如，它打印：创建TensorFlow设备（/ gpu：0） - ＆gt;（设备：0，名称：GeForce GTX 1070，pci总线ID：...））