在 TensorFlow 中计算张量中所有向量之间成对欧氏距离的有效方法

Efficient way to compute pairwise euclidean distances between all vectors in a tensor in TensorFlow

feature_map 为张量,其形状为:b,h,w,c = feature_map.shape

要计算通道 (c) 轴中所有向量之间的成对欧氏距离,我使用以下代码:

feature_map = tf.reshape(feature_map,(b,-1,c))
norms = []
for i in range(h*w):
    norm = tf.norm(feature_map-feature_map[:,i:i+1,:],ord=2,axis=-1)
    norms.append(sorted_tensor)
norms = tf.stack(norms,axis=1)

这个程序很慢。为了避免 for-loop 我尝试使用 tf.tile 复制张量,然后应用 norm 函数但导致内存错误,因为生成的张量非常大。

有没有更有效的方法来完成这个?

我使用来自 here

的欧氏距离矩阵技巧

numpy改编如下:

feature_map = tf.reshape(feature_map,(b,-1,c))         
G = np.einsum('bik, bjk->bij', feature_map, feature_map)
D = G.diagonal(axis1=1,axis2=2).reshape(b,-1,1)+ np.transpose(G.diagonal(axis1=1,axis2=2).reshape(b,-1,1),axes=(0,2,1)) - 2*G
norms = np.sqrt(D)

和对应的tensorflow适配:

feature_map = tf.reshape(feature_map,(b,-1,c))     
G = tf.einsum('bik, bjk->bij', feature_map, feature_map)
D = tf.reshape(tf.linalg.diag_part(G),(b,-1,1))+ tf.transpose(tf.reshape(tf.linalg.diag_part(G),(b,-1,1)),perm=(0,2,1)) - 2*G
norms = tf.sqrt(D)