在循环中评估 Tensorflow 操作非常慢

Evaluating Tensorflow operation is very slow in a loop

我正在尝试通过编写一些简单的问题来学习 tensorflow:我正在尝试使用直接采样 Monte Carlo 方法找到 pi 的值。

运行 时间比我想象的要长得多,当使用 for loop 来做这件事时。我看过其他关于类似事情的帖子,并且我尝试按照解决方案进行操作,但我认为我仍然做错了什么。

下面附上我的代码:

import tensorflow as tf
import numpy as np
import time

n_trials = 50000

tf.reset_default_graph()


x = tf.random_uniform(shape=(), name='x')
y = tf.random_uniform(shape=(), name='y')
r = tf.sqrt(x**2 + y**2)

hit = tf.Variable(0, name='hit')

# perform the monte carlo step
is_inside = tf.cast(tf.less(r, 1), tf.int32)
hit_op = hit.assign_add(is_inside) 

with tf.Session() as sess:
    init_op = tf.global_variables_initializer()
    sess.run(init_op)

    # Make sure no new nodes are added to the graph
    sess.graph.finalize()

    start = time.time()   

    # Run monte carlo trials  -- This is very slow
    for _ in range(n_trials):
        sess.run(hit_op)

    hits = hit.eval()
    print("Pi is {}".format(4*hits/n_trials))
    print("Tensorflow operation took {:.2f} s".format((time.time()-start)))

>>> Pi is 3.15208
>>> Tensorflow operation took 8.98 s

相比之下,在 numpy 中执行 for loop 类型的解决方案要快一个数量级

start = time.time()   
hits = [ 1 if np.sqrt(np.sum(np.square(np.random.uniform(size=2)))) < 1 else 0 for _ in range(n_trials) ]
a = 0
for hit in hits:
    a+=hit
print("numpy operation took {:.2f} s".format((time.time()-start)))
print("Pi is {}".format(4*a/n_trials))

>>> Pi is 3.14032
>>> numpy operation took 0.75 s

下面附上了不同试验次数的总体执行时间差异图。

请注意:我的问题不是关于 "how to perform this task the fastest",我知道有更有效的计算 Pi 的方法。我只将它用作基准测试工具来检查 tensorflow 与我熟悉的东西 (numpy) 的性能。

简单,session.run开销很大,它不是为那样使用而设计的。通常,例如一个神经网络,你会调用一个 session.run 来进行十几个大矩阵的乘法运算,那么这 0.2 毫秒就根本不重要了。 至于你的情况,你可能想要这样的东西。它在我的机器上运行速度比 numpy 版本快 5 倍。

顺便说一句,你在 numpy 中做完全相同的事情。如果你使用循环而不是 np.sum 来减少它会慢得多。

    import tensorflow as tf
    import numpy as np
    import time

    n_trials = 50000

    tf.reset_default_graph()

    x = tf.random_uniform(shape=(n_trials,), name='x')
    y = tf.random_uniform(shape=(), name='y')
    r = tf.sqrt(x**2 + y**2)

    hit = tf.Variable(0, name='hit')

    # perform the monte carlo step
    is_inside = tf.cast(tf.less(r, 1), tf.int32)
    hit2= tf.reduce_sum(is_inside)
        #hit_op = hit.assign_add(is_inside) 

    with tf.Session() as sess:
    #    init_op = tf.global_variables_initializer()
        sess.run(tf.initialize_all_variables())

        # Make sure no new nodes are added to the graph
        sess.graph.finalize()

        start = time.time()   

        # Run monte carlo trials  -- This is very slow
        #for _ in range(n_trials):
        sess.run(hit2)

        hits = hit2.eval()
        print("Pi is {}".format(4*hits/n_trials))
        print("Tensorflow operation took {:.2f} s".format((time.time()-start)))

速度慢与 Python 和 sess.run 中的 Tensorflow 之间的一些通信开销有关,它在你的循环中执行了多次。我建议使用 tf.while_loop 在 Tensorflow 中执行计算。这将是比 numpy.

更好的比较
import tensorflow as tf
import numpy as np
import time

n_trials = 50000

tf.reset_default_graph()

hit = tf.Variable(0, name='hit')

def body(ctr):
    x = tf.random_uniform(shape=[2], name='x')
    r = tf.sqrt(tf.reduce_sum(tf.square(x))
    is_inside = tf.cond(tf.less(r,1), lambda: tf.constant(1), lambda: tf.constant(0))
    hit_op = hit.assign_add(is_inside)
    with tf.control_dependencies([hit_op]):
        return ctr + 1

def condition(ctr):
    return ctr < n_trials

with tf.Session() as sess:
    tf.global_variables_initializer().run()
    result = tf.while_loop(condition, body, [tf.constant(0)])

    start = time.time()
    sess.run(result)

    hits = hit.eval()
    print("Pi is {}".format(4.*hits/n_trials))
    print("Tensorflow operation took {:.2f} s".format((time.time()-start)))