tensorflow如何分配Ops在GPU上运行？

我对tensorflow用于为CPU或GPU分配不同操作的机制感到困惑

以下面的伪代码为例。我们可以说：只要

SimpleOp

是 在使用tf.device（'/gpu:0'）的上下文中创建，它将 一定要在GPU上运行（假设

SimpleOp

可用），无论其输入变量（

在_1

和

在_2

中）是 在CPU或GPU上创建

with tf.device('/gpu:0'):
    out = tf.SimpleOp(in_1, in_2, name='Simple')

我通过创建一个

会话来理解
log\u device\u placement=True
，tensorflow输出设备
所有变量/操作的放置。然而，是否有一种方法允许
让我只检查一个Op的设备分配

提前谢谢

对。事实上，如果指定的设备没有可用的内核，它将失败。但有两件事需要考虑：

• 这可以通过if会话的

  config

  来覆盖

• tf.device

  上下文管理器可以嵌套，因此如果

  SimpleOp

  不是那么简单，它可能已经用tf.device（“/cpu:0”）用

  包装了它的一些部分：

据我所知并非如此（欢迎评论）。如果您使用*nix，

python script.py，则始终可以
grep
您的脚本输出，grep您的_op_名称

。缺点是您需要重新运行脚本两次：第一次运行时使用

log\u device\u placement=True

和

grep

，第二次运行时不使用它们

---

TLDR；使用tf.device（“/gpu:0”）在

中创建的op将始终在gpu上运行。如果指定要放置在
cpu
上的输入，则它们将被放置在cpu上。如果您省略输入的设备规范，它们将被放置在GPU上，以便更接近您的op。您可以使用
run\u metadata
获取包含所有设备分配的Python对象，并在那里查找您的op

放置是通过错误命名来完成的，虽然注释指定了机制，但仍然有一些bug需要解决（即），所以最好的方法是在实践中检查它

当你说变量是在GPU上创建的时，实际上有两种放置方式——显式的，当你用tf.device在
块内创建相关的op时，以及隐式的，在这样的块外。使用tf.device在
之外创建ops相当于使用tf.device（None）
块在
中创建ops

这是一个简单的实验

n = 10**6
def inputs_cpu():
    tf.reset_default_graph()
    with tf.device("/cpu:0"):
        a = tf.ones((n,), name="A")
        b = tf.ones((n,), name="B")
    with tf.device("/gpu:0"):
        c = tf.add(a, b, name="C")
    return c

def inputs_none():
    tf.reset_default_graph()
    a = tf.ones((n,), name="A")
    b = tf.ones((n,), name="B")
    with tf.device("/gpu:0"):
        c = tf.add(a, b, name="C")
    return c

def run_and_summarize(target):
    # turn off graph-rewriting optimizations
    sess = tf.Session(config=tf.ConfigProto(graph_options=tf.GraphOptions(optimizer_options=tf.OptimizerOptions(opt_level=tf.OptimizerOptions.L0))))
    run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
    run_metadata = tf.RunMetadata()
    sess.run(target, options=run_options, run_metadata=run_metadata)

    for device in run_metadata.step_stats.dev_stats:
        device_name = device.device
        if not (device_name.endswith("/cpu:0") or device_name.endswith("/gpu:0")):
            continue
        print(device.device)
        for node in device.node_stats:
            print("   ", node.node_name)

现在你可以这样做了

run_and_summarize(inputs_cpu())

运行时，输入固定在CPU上，您将看到这种放置方式得到了尊重

/job:localhost/replica:0/task:0/gpu:0
    _SOURCE
    C
/job:localhost/replica:0/task:0/cpu:0
    _SOURCE
    A
    B

另一方面，当未指定输入时

run_and_summarize(inputs_none())

您可以看到，现在所有操作都放置在GPU上

/job:localhost/replica:0/task:0/cpu:0
    _SOURCE
/job:localhost/replica:0/task:0/gpu:0
    _SOURCE
    A
    B
    C

棒 极 了我现在拿到了。谢谢！顺便说一句，最近在评论中澄清了这种行为。基本上，如果您显式地设置设备，它将始终位于该设备上。我以前的印象是，只要GPU和这些操作的合适内核可用，TF就会隐式地将操作分配给GPU。从这一描述中，我现在了解到，如果先显式地（例如，
与TF.device（“/gpu:0”）：
）分配了与这些ops具有输入关系的其他ops，则TF将仅隐式地将ops分配给gpu。所以，如果我不显式地向GPU分配任何ops，TF也不会隐式地分配任何ops。对吗？