张量流中的广播和降维

Question

我有以下

  tensor A with A.shape = (N,2) 
  tensor B with B.shape = (3,2)

物理上我将 A 可视化为二维中的 N 个数据点。

B 是同一个二维空间中的 3 个中心。

我的objective是计算A到3个中心的距离的平方，然后相加（即系统到3个中心的惯性总和）。我要计算

$$ D = \Sum_{i,j} (A(i,j) - B(1,j))^2 + (A(i,j) - B(2,j))^2 + (A(i,j) - B(3,j))^2 $$

谁能帮我弄清楚如何在 tensorflow + python 中实现这一点。提前致谢

Answer 1

人们可能想到的第一个解决方案可能是

tf.reduce_sum(tf.square(A-B[0])+tf.square(A-B[1])+tf.square(A-B[2]))

，这是将您的公式直接翻译成代码。不过，使用 tensorflow 提供的隐式广播稍微更有效。

tf.reduce_sum(tf.square(A[:,None,:]-B[None,:,:]))

微基准测试代码 1（大型数据集）：

A=tf.random.normal((2**25,2))
B=tf.random.normal((3,2))
%timeit tf.reduce_sum(tf.square(A[:,None,:]-B[None,:,:]))
%timeit tf.reduce_sum(tf.square(A-B[0])+tf.square(A-B[1])+tf.square(A-B[2]))

输出：

13.1 ms ± 38.9 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
15.7 ms ± 7.83 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

微基准测试代码 2（小数据集）：

A=tf.random.normal((3,2))
B=tf.random.normal((3,2))
%timeit tf.reduce_sum(tf.square(A[:,None,:]-B[None,:,:]))
%timeit tf.reduce_sum(tf.square(A-B[0])+tf.square(A-B[1])+tf.square(A-B[2]))

输出：

175 µs ± 731 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)
391 µs ± 18 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

张量流中的广播和降维

broadcasting and reducing dimension in tensorflow

python

numpy

broadcast

tensorflow

tensor