启用急切执行时如何运行并行map_fn

时间:2018-10-12 07:21:23

标签: python tensorflow parallel-processing

考虑以下张量流代码段:

import time
import numpy as np
import tensorflow as tf

def fn(i):
    # do some junk work
    for _ in range(100):
        i ** 2
    return i

n = 1000
n_jobs = 8
stuff = np.arange(1, n + 1)
eager = False
t0 = time.time()
if eager:
    tf.enable_eager_execution()
res = tf.map_fn(fn, stuff, parallel_iterations=n_jobs)
if not eager:
    with tf.Session() as sess:
        res = sess.run(res)
        print(sum(res))
else:
    print(sum(res))
dt = time.time() - t0
print("(eager=%s) Took %ims" % (eager, dt * 1000))

如果使用eager = True运行,则比使用eager = False运行时要慢10倍。我做了一些打印,发现在eager = True模式下,map_fn调用是顺序运行的,而不是产生8个并行线程。

问题

所以我的问题是如何在急切的执行模式下使用map_fn(parallel_iterations> 1)?

4 个答案:

答案 0 :(得分:1)

粗略地说,tf.map_fn(fn, data)的简写为:

for e in data:
  fn(e)

启用急切执行后,操作将在Python解释器遇到操作时执行,因此没有机会进行“整个程序优化”。

在执行TensorFlow图时,TensorFlow运行时会看到要执行的完整计算,因此可以应用优化,例如“来自并行循环中多次迭代的fn中的执行操作”。这是将计算表示为图形的好处之一。

在TensorFlow中启用急切执行后,您仍然可以使用tf.contrib.eager.defun有选择地将图优化应用于程序的某些部分。

例如(其中大多数代码与上面的代码相同,然后更改一行以使用tf.contrib.eager.defun来获得图形优化的好处)

import time
import numpy as np
import tensorflow as tf

tf.enable_eager_execution()

def fn(i):
    # do some junk work
    for _ in range(100):
        i ** 2
    return i

n = 1000
n_jobs = 8
stuff = np.arange(1, n + 1)

def my_computation(x):
  return tf.map_fn(fn, x, parallel_iterations=n_jobs)

t0 = time.time()
my_computation(stuff)
dt = time.time() - t0
print("my_computation took %ims" % (dt * 1000))

my_computation = tf.contrib.eager.defun(my_computation)
# On the very first call, a graph is constructed, so let's discount
# graph construction time
_ = my_computation(stuff)

# And then time it
t0 = time.time()
my_computation(stuff)
dt = time.time() - t0
print("my_computation took %ims" % (dt * 1000))

一些其他注意事项:

  • 在上面提供的特定示例中,TensorFlow运行时可能还会检测到fn(i)减少为return i并可以优化range(100)的不必要循环,因为不会影响输出。因此,性能上的差异相当大(例如,当急切地执行fn(i)时,Python解释器无法知道其中的for循环是无用的,因此它将执行该循环)。

  • 如果将fn()中的计算更改为更有意义的内容,请说:

    def fn(i):   对于_范围(2):     我=我** 2   返回我

    然后您会看到明显的不同。

  • 请注意,并非所有可以在Python中表达的内容都可以被“破坏”。有关更多详细信息,请参见tf.contrib.eager.defun的文档,并为TensorFlow 2.0提出了更详细的规范和实现(请参见RFC

希望有帮助。

答案 1 :(得分:1)

这不仅仅是对OP问题的答案,它是对OP的扩展,显示了为什么其他答案不能解决真正的问题,因为tf.function不足以强制并行化。


首先,使用tf.function不会强制并行化。它会强制进行跟踪,并且仅构建一次图形,因此,其他答案中使用的time.sleep()仅在需要跟踪时才运行,这就是为什么看到{{1} }。但是在更改tf.function时,您仍然看不到任何区别。

让我们用parallel_iterations看一下区别:

py_fuction

不使用装饰器(或直接调用)def op(x): time.sleep(1) return 2 * x.numpy() def op_tf(x): print('Tracing') return tf.py_function(op, [x], Tout=tf.int32) ,对tf.function的任何调用都将始终显示“正在跟踪”(尽管在这种情况下无法跟踪)

op_tf

In [57]: op_tf(1) Tracing Out[57]: <tf.Tensor: shape=(), dtype=int32, numpy=2> In [58]: op_tf(1) Tracing Out[58]: <tf.Tensor: shape=(), dtype=int32, numpy=2> 中,我们仅看到一次跟踪(如果使用相同的参数):

tf.function

之所以会发生这种情况,是因为该函数必须为每个新参数构建一个新图形,如果我们直接传递一个签名,就可以避免这种情况发生:

In [67]: @tf.function
    ...: def op_tf(x):
    ...:     print("Tracing")
    ...:     return tf.py_function(op, [x], Tout=tf.int32)
    ...: 

In [68]: op_tf(1)
Tracing
Out[68]: <tf.Tensor: shape=(), dtype=int32, numpy=2>

In [69]: op_tf(2)
Tracing
Out[69]: <tf.Tensor: shape=(), dtype=int32, numpy=4>

In [70]: op_tf(3)
Tracing
Out[70]: <tf.Tensor: shape=(), dtype=int32, numpy=6>

In [71]: op_tf(3)
Out[71]: <tf.Tensor: shape=(), dtype=int32, numpy=6>

如果我们首先调用方法In [73]: @tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.int32)]) ...: def op_tf(x): ...: print("Tracing") ...: return tf.py_function(op, [x], Tout=tf.int32) ...: ...: In [74]: op_tf(1) Tracing Out[74]: <tf.Tensor: shape=(), dtype=int32, numpy=2> In [75]: op_tf(2) Out[75]: <tf.Tensor: shape=(), dtype=int32, numpy=4> In [76]: op_tf(3) Out[76]: <tf.Tensor: shape=(), dtype=int32, numpy=6> ,也会发生同样的情况:

get_concrete_function

然后,仅添加In [79]: @tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.int32)]) ...: def op_tf(x): ...: print("Tracing") ...: return tf.py_function(op, [x], Tout=tf.int32) ...: ...: In [80]: op_tf = op_tf.get_concrete_function() Tracing In [81]: op_tf(1) Out[81]: <tf.Tensor: shape=(), dtype=int32, numpy=2> In [82]: op_tf(2) Out[82]: <tf.Tensor: shape=(), dtype=int32, numpy=4> 就足以并行执行的答案并不完全正确:

tf.function

相比之下,如果用于睡眠和打印的python指令位于py_function内,则它们将始终被调用:

In [84]: def op(x):
    ...:     print("sleep")
    ...:     time.sleep(0.1)
    ...:     return 1.
    ...: 

In [85]: x = tf.ones(shape=(10,))

In [86]: _ = tf.map_fn(op, x, parallel_iterations=10)
sleep
sleep
sleep
sleep
sleep
sleep
sleep
sleep
sleep
sleep

In [87]: @tf.function
    ...: def my_map(*args, **kwargs):
    ...:     return tf.map_fn(*args, **kwargs)
    ...: 

In [88]: my_map(op, x, parallel_iterations=10)
sleep
Out[88]: <tf.Tensor: shape=(10,), dtype=float32, numpy=array([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.], dtype=float32)>

现在,我们已经很清楚地知道函数的跟踪给我们带来了一些困惑,让我们删除打印内容以计时调用:

In [96]: x = tf.ones(shape=(10,), dtype=tf.int32)

In [97]: def op(x):
    ...:     print("sleep")
    ...:     time.sleep(0.1)
    ...:     return 2 * x.numpy()
    ...: 

In [98]: @tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.int32)])
    ...: def op_tf(x):
    ...:     print("Tracing")
    ...:     return tf.py_function(op, [x], Tout=tf.int32)
    ...: 

In [99]: _ = my_map(op_tf, x, parallel_iterations=1)
Tracing
sleep
sleep
sleep
sleep
sleep
sleep
sleep
sleep
sleep
sleep

运行以下脚本并使用tensorboard,我们可以确切地看到正在发生的事情:

In [106]: def op(x):
     ...:     time.sleep(0.1)
     ...:     return 2 * x.numpy()
     ...: 

In [107]: @tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.int32)])
     ...: def op_tf(x):
     ...:     return tf.py_function(op, [x], Tout=tf.int32)
     ...: 

In [108]: %timeit tf.map_fn(op_tf, x, parallel_iterations=1)
1.02 s ± 554 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [109]: %timeit tf.map_fn(op_tf, x, parallel_iterations=10)
1.03 s ± 509 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

我们在Tensorboard中获得以下内容: enter image description here

import tensorflow as tf import time from datetime import datetime stamp = datetime.now().strftime("%Y%m%d-%H%M%S") logdir = 'logs/func/%s' % stamp # Start tracing. options = tf.profiler.experimental.ProfilerOptions( host_tracer_level=3, python_tracer_level=1, device_tracer_level=1, delay_ms=None ) tf.profiler.experimental.start(logdir, options = options) def op(x): x = x.numpy() start = time.time() while time.time() < start + x / 100: x = (2 * x) % 123 return x @tf.function(input_signature=[tf.TensorSpec([], tf.int32)]) def op_tf(x): return tf.py_function(op, [x], Tout=tf.int32, name='op') @tf.function(input_signature=[tf.TensorSpec([None], tf.int32)]) def my_map(x): return tf.map_fn(op_tf, x, parallel_iterations=16) x = tf.ones(100, tf.int32) print(my_map(x)) tf.profiler.experimental.stop() 有效地使用了多个线程,但不是并行使用。使用py_function,我们可以获得类似的结果 enter image description here

如果我们在脚本的开头添加以下内容

parallel_iterations=1

我们强迫TF使用单个线程进行所有图形计算: enter image description here

因此,目前,只有正确设置内部/内部线程,我们才能获得某种形式的并行执行。

如果我们完全禁用急切执行:

tf.config.threading.set_inter_op_parallelism_threads(1)
tf.config.threading.set_intra_op_parallelism_threads(1)

我们现在可以在Tensorboard中看到并行执行: enter image description here

如果将intra / inter线程和parallel_iterations设置为1,则会返回以前的行为: enter image description here

我希望这有助于阐明import time from datetime import datetime import numpy as np import tensorflow as tf tf.compat.v1.disable_eager_execution() tf.config.threading.set_inter_op_parallelism_threads(128) tf.config.threading.set_intra_op_parallelism_threads(128) stamp = datetime.now().strftime("%Y%m%d-%H%M%S") logdir = f'logs/func/{stamp}' tf.profiler.experimental.start(logdir) def op(x): x = x.numpy() start = time.time() while time.time() < start + x / 100: x = (2 * x) % 123 return x @tf.function(input_signature=[tf.TensorSpec([], tf.int32)]) def op_tf(x): return tf.py_function(op, [x], Tout=tf.int32, name='op') # Create a placeholder. x = tf.compat.v1.placeholder(tf.int32, shape=[None]) with tf.compat.v1.Session() as sess: writer = tf.summary.create_file_writer(logdir) #tf.profiler.experimental.start(logdir, options = options) tf.summary.trace_on(graph=True, profiler=True) print( sess.run( [tf.map_fn(op_tf, x, parallel_iterations=16)], feed_dict={ x: np.ones(4, dtype=np.int) } ) ) tf.profiler.experimental.stop() 在检查完全并行性中的作用。

答案 2 :(得分:0)

在此为TF2.0用户更新。 您可以通过将调用包装到tf.function装饰器中来并行调用tf.map_fn内部运算符:

import tensorflow as tf
import time

x = tf.ones(shape=(10,))

def op(x):
    time.sleep(0.1)
    return 1.

_ = tf.map_fn(op, x, parallel_iterations=10) # will take 1 sec along with the 
                                             # warning message.

# Now wrap tf.map_fn inside tf.function
@tf.function
def my_map(*args, **kwargs):
    return tf.map_fn(*args, **kwargs)

_ = my_map(op, x, parallel_iterations=10) # will take 0.1 sec along with no 
                                          # warning message.

答案 3 :(得分:0)

让我在TF2.1上回答这个问题的答案。

从TF2.x开始,当计算图描述语句可以在紧急模式下执行时,某些tf函数可以自然地在会话模式下运行,因此可以并行执行。

一个简单的解决方案是使用tf.function将那些急切的运行python函数的模式转换为tf_function的运行模式(会话),而无需更改整个编程模式(从急切模式转换为会话模式)。

@RémyDubois解决方案在TF2.1中工作正常。

@tf.function
def my_map(*args, **kwargs):
    return tf.map_fn(*args, **kwargs)

_ = my_map(op, x, parallel_iterations=10) # will take 0.1 sec along with no 
                                          # warning message.

例如,我们还可以通过使用tf.function()转换my_map函数来动态更改它。

def my_map(*args, **kwargs):
    return tf.map_fn(*args, **kwargs)

my_map=tf.function(my_map)
_ = my_map(op, x, parallel_iterations=10) # will take 0.1 sec along with no 
                                          # warning message.

以上两种解决方案现在都可以在TF2.1中正常工作。

TF2.x中关于tf.map_fn()的parallel_iterations的警告消息是错误的,因为:

Setting parallel_iterations > 1 has no effect when executing eagerly. Consider calling map_fn with tf.contrib.eager.defun to execute fn in parallel.

因为tf.contrib.eager.defun已更改为tf.function。