张量流中的广播和降维

Question

我在物理上有以下内容

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

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

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

我的目标是计算 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]))

，它将公式直接转换为代码。不过，使用张量流提供的隐式广播效率稍高一些。

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)