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

Question

给定形状 (n, f) 的张量 a 和形状 (m, f) 的 b，我创建了一个函数来计算这两个张量之间的欧氏距离

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 时获得了更好的性能，而在 n= 1000, m=1600, f=4 时，但这不是现在的问题）。稍后我将检查性能与不同输入大小的关系。

在此示例中，b 张量是 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 需要 巨大的 时间才能达到运行，大约 300 毫秒。

只是为了比较，如果我在完全相同的数据（但 numpy 数组）的数据集中执行 np.apply_along_axis，足迹几乎不明显，大约 50 微秒与 300 毫秒的 tensorflow 等效。

在 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 数组中的开销要小得多。

Answer 1

你可以试试 tf.vectorized_map。 https://www.tensorflow.org/api_docs/python/tf/vectorized_map

如果您需要更改数据类型，您可以尝试更改 map_fn 参数中的 parallel_iterations 值，在 eager 模式下默认设置为 1。

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

Tensorflow `map_fn` takes long time to execute

python

mapping

tensorflow

tensorflow2.0