给定形状{
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数组中的开销要小得多
你可以试试tf矢量化的地图https://www.tensorflow.org/api_docs/python/tf/vectorized_map
如果需要更改de数据类型,可以尝试更改map_fn参数中的parallel_iterations值,该值在急切模式下默认设置为1
相关问题 更多 >
编程相关推荐