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

Question

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

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

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

2. 我理解通过创建sessionwith log_device_placement=True,tensorflow输出所有变量/ Ops的设备位置.但是,有没有一种方法可以让我只检查一个Op的设备分配？

提前致谢!

Answer 1

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

放置是由误导性命名的simple_placer.cc完成的,虽然注释指定了机制,但仍然有一些错误被挖出(即,这里),所以最好的方法是在实践中检查它.

当你说在GPU上创建变量时,实际上有两种放置 - 显式,当你在with tf.device块内部创建相关的op时,以及隐藏的,在这种块之外.在外部创建ops with tf.device相当于在with 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