Tensorflow`map_fn`需要很长时间才能执行

2024-03-29 02:11:49 发布

您现在位置:Python中文网/ 问答频道 /正文

给定形状{}的{}张量和形状{}的{}张量,我创建了一个函数来计算这两个张量之间的欧几里德距离

import tensorflow as tf

nr = tf.reduce_sum(tf.square(a), 1)
nw = tf.reduce_sum(tf.square(b), 1)

nr = tf.reshape(nr, [-1, 1])
nw = tf.reshape(nw, [1, -1])
res = nr - 2*tf.matmul(a, b, False, True) + nw
res = tf.argmin(res, axis=1)

到目前为止,代码运行速度稍快(我在cKDTree时获得了更好的性能,但现在这不是问题所在)。稍后,我将对照不同的输入大小检查性能

在这个例子中b张量是秩3张量的一个秩2平坦版本。我这样做是为了能够使用两个具有相同秩的张量来计算欧几里德距离(这更简单)。但是在计算距离之后,我需要知道每个最近的元素在原始张量上的位置。为此,我创建了自定义lambda函数fn,以转换回秩3张量坐标

fn = lambda x: (x//N, x%N)
# This map takes a enormous amount of time
out = tf.map_fn(fn, res, dtype=(tf.int64, tf.int64))
return tf.stack(out, axis=1)

但遗憾的是,这个tf.map_fn需要大量的时间运行,大约300ms

只是为了比较,如果我在一个数据集中执行一个np.apply_along_axis,这个数据集与相同的数据(但是一个numpy数组)几乎看不到足迹,大约50微秒,而tensorflow相当于300毫秒

在tensorflow中有更好的方法来解决这个问题mapping

TF版本2.1.0和CUDA已启用并工作

只是为了增加一些时间

%timeit eucl_dist_tf_vecmap(R_tf, W_tf)
28.1 ms ± 128 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)


%timeit eucl_dist_tf_nomap(R_tf, W_tf)
2.07 ms ± 122 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)


%timeit eucl_dist_ckdtree_applyaxis(R, W)  
878 µs ± 2.34 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit eucl_dist_ckdtree_noapplyaxis(R, W)  
817 µs ± 51 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

前两个计时使用此处显示的自定义函数,第一个计时带有vectorized_map,第二个计时不带vectorized_map,第二个计时带有stack(开销在vectorized_map上,已测试)

最后两次是基于scipy的cKDTree的实现。第一次使用np.apply_along_axis与矢量化映射中使用的完全相同。我们可以看到numpy数组中的开销要小得多


Tags: ofloopmapdisttfresmeannr