Эффективный способ попарного вычисления евклидовых расстояний между всеми векторами в тензоре в 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, но это привело к ошибкам памяти, поскольку сгенерированный тензор был очень большим.

Есть ли более эффективный способ добиться этого?


numpy tensorflow tensorflow2.0 vectorization array-broadcasting
person xro7    schedule 17.11.2019    source источник

Ответы (1)


arrow_upward
0
arrow_downward

Я использовал трюк с евклидовой матрицей расстояний из здесь

Адаптация 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)
person xro7    schedule 20.11.2019