考虑到形状为{{1}的张量a和形状为(n, f)的{{1}}，我创建了一个函数来计算这两个张量之间的欧式距离
b到目前为止，代码运行得还算快（当使用(m, f)时，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张量是2级的张量的平化版本。我这样做是为了能够使用具有相同等级（更简单）的两个张量来评估欧几里得距离。但是在评估了距离之后，我需要知道最接近的每个元素在原始张量上的位置。为此，我创建了自定义lambda函数n= 1000, m=1600, f=4，以将其转换回3级张量坐标。
b可惜的是，这fn花费了巨大的时间，大约 300ms 。
为了进行比较，如果我在完全相同的数据（但是一个numpy数组）的数据集中执行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)
，则占用空间几乎无法察觉，大约为50微秒，而张量流等效值为300ms。
对此tf.map_fn而言，张量流中有更好的方法吗？
TF版本2.1.0和CUDA已启用并正常工作。
只需添加一些时间
np.apply_along_axis前两个时间使用此处显示的自定义功能，第一个时间使用mapping，第二个时间使用%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（开销在{{1}上） }，经过测试。
最后两次是基于scipy的vectorized_map的实现。第一个使用stack的方式与向量化地图中使用的方式完全相同。我们可以看到numpy数组中的开销要小得多。
答案 0 :(得分：1)
您可以尝试tf.vectorized_map。 https://www.tensorflow.org/api_docs/python/tf/vectorized_map
如果需要更改数据类型，则可以尝试更改map_fn参数中的parallel_iterations值，该值在急切模式下默认设置为1。