在Tensorflow中获得欧几里得距离的所有组合？

Question

我想在 Tensorflow 中形成一个损失函数，它依赖于一个矩阵，该矩阵包含一组嵌入的（平方）欧几里德距离的所有组合。在 numpy 中，像这样：

# E is (batch_size,N,32)
N=100
D = np.zeros((batch_size,N,N))
for x in range(N):
  for y in range(N):
    D[:,x,y] = np.sum(np.square(E[:,x,:]-E[:,y,:]),axis=1)

如何在 Tensorflow/Keras 中编写代码而不使用嵌套的 for 循环，或者根本不使用 for 循环？

Answer 1

应该这样做：

D = tf.reduce_sum((E[:, None, :] - E[:, :, None])**2, axis=-1)

D 将是(batch_size, N, N)。这也适用于 numpy（显然使用 np.sum），因此您可以使用它来检查循环版本的等价性以确保。

此解决方案通过广播工作：None 用于插入轴，以使大小为 N 的轴与大小为 1 的轴匹配，而后者被广播（重复）以匹配前者。这导致将所有元素与所有其他元素（每个批次元素）进行比较。由于我们在这里处理的是四维张量，所以用文字描述有点困难，也很难形象化......