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

Question

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

1. 以下面的伪代码为例。我们可以说:只要 SimpleOp 是在 with tf.device('/gpu:0') 的上下文中创建，它将肯定在 GPU 上运行(假设 SimpleOp 的 GPU 实现可用)，无论其输入变量(in_1 和 in_2)都是在 CPU 或 GPU 上创建？

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

2. 我通过创建一个 session 来理解 log_device_placement=True，tensorflow输出设备 所有变量/操作的放置。但是，有没有一种方法允许 我只检查一个 Op 的设备分配？

提前致谢!

Answer 1

TLDR; 使用 tf.device("/gpu:0") 创建的操作将始终在 GPU 上运行。如果您指定将输入放置在 cpu 上，那么它们将放置在 CPU 上。如果您省略输入的设备规范，它们将被放置在 GPU 上以更接近您的操作。您可以使用 run_metadata 获取包含所有设备分配的 Python 对象，并在那里查找您的操作。

放置是通过误导性命名 simple_placer.cc 完成的，虽然注释指定了机制，但仍然存在一些错误被解决(即 here )，因此最好的方法是在实践中检查它。

当您说变量是在 GPU 上创建时，实际上有两种放置方式:显式放置(当您在 with tf.device block 内创建相关操作时)和隐式放置(在此类外部)堵塞。在 with tf.device 之外创建操作相当于在 with tf.device(None) block 中创建操作。

这是一个简单的实验

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