首页 > 解决方案 > 为什么 tensorflow.map_fn 慢,下面的代码有什么问题?

问题描述

我正在尝试使用 tensorflow map_fn 进行并行计算。但是在我看来,性能提升并不显着。

以下是在 Ubuntu 14.04 LTS 上运行 Python 3.6.5、Tensorflow 版本 1.12.0、28 双核(Intel(R) Xeon(R) CPU E5-2697 v3 @ 2.60GHz)= 56 个处理器的示例代码

在 Amazon AWS SagerMaker ml-p3-xlarge 上运行的这些相同代码甚至需要更长的时间,227 秒。

import tensorflow as tf
import time
# version 1
tic = time.time()
elems = np.array(range(1,1000000), dtype=np.float64)
output = tf.map_fn(lambda x: x**6 , elems, dtype=tf.float64,  parallel_iterations=56)
sess = tf.Session()

res = sess.run(output)
toc = time.time() - tic
print("elapsed=", toc)  # 29.47 (seconds)

# version 2
tic = time.time()
elems = np.array(range(1,1000000), dtype=np.float64)
output = tf.map_fn(lambda x: x**6 , elems, dtype=tf.float64,   parallel_iterations=56)
n_cpus=28


with  tf.Session(
config=tf.ConfigProto(log_device_placement=True, 
device_count={ "CPU": n_cpus },
inter_op_parallelism_threads=n_cpus,
intra_op_parallelism_threads=1,

)) 作为 sess:
res = sess.run(output)

toc = time.time() - tic
print("elapsed=", toc)  # 29.26 (seconds)

# version 3
tic = time.time()
elems = np.array(range(1,1000000), dtype=np.float64)
x6 = [ x**6 for x in elems]
toc = time.time() - tic
print("elapsed time=", toc) # 0.5 seconds

上面的代码有什么问题?没有 map_fn,顺序执行版本 3 只需 0.5(秒)。

标签: pythontensorflow

解决方案


推荐阅读