Calcule la distancia de pares entre cada par de las dos colecciones de entradas en TensorFlow

Tengo dos colecciones. Uno consta de m 1 puntos en k dimensiones y otro de m 2 puntos en k dimensiones. Necesito calcular la distancia por pares entre cada par de las dos colecciones.

Básicamente tiene dos matrices A m 1 , k y B m 2 , k Necesito obtener una matriz C m 1 , m 2 .

Puedo hacer esto fácilmente en scipy usando distance.sdist y seleccionar una de las muchas métricas de distancia, y también puedo hacer esto en TF en un bucle, pero no puedo descubrir cómo hacer esto con manipulaciones matriciales incluso para la distancia Euclediana .

Después de unas horas finalmente encontré cómo hacer esto en Tensorflow. Mi solución funciona solo para la distancia euclediana y es bastante detallada. Tampoco tengo una prueba matemática (solo un montón de handwaving, que espero hacer más riguroso):

 import tensorflow as tf import numpy as np from scipy.spatial.distance import cdist M1, M2, K = 3, 4, 2 # Scipy calculation a = np.random.rand(M1, K).astype(np.float32) b = np.random.rand(M2, K).astype(np.float32) print cdist(a, b, 'euclidean'), '\n' # TF calculation A = tf.Variable(a) B = tf.Variable(b) p1 = tf.matmul( tf.expand_dims(tf.reduce_sum(tf.square(A), 1), 1), tf.ones(shape=(1, M2)) ) p2 = tf.transpose(tf.matmul( tf.reshape(tf.reduce_sum(tf.square(B), 1), shape=[-1, 1]), tf.ones(shape=(M1, 1)), transpose_b=True )) res = tf.sqrt(tf.add(p1, p2) - 2 * tf.matmul(A, B, transpose_b=True)) with tf.Session() as sess: sess.run(tf.global_variables_initializer()) print sess.run(res) 

Esto lo hará para tensores de dimensionalidad arbitraria (es decir, que contienen (…, N, d) vectores). Tenga en cuenta que no es entre colecciones (es decir, no como scipy.spatial.distance.cdist ) sino que está dentro de un único lote de vectores (es decir, como scipy.spatial.distance.pdist )

 import tensorflow as tf import string def pdist(arr): """Pairwise Euclidean distances between vectors contained at the back of tensors. Uses expansion: (x - y)^T (x - y) = x^Tx - 2x^Ty + y^Ty :param arr: (..., N, d) tensor :returns: (..., N, N) tensor of pairwise distances between vectors in the second-to-last dim. :rtype: tf.Tensor """ shape = tuple(arr.get_shape().as_list()) rank_ = len(shape) N, d = shape[-2:] # Build a prefix from the array without the indices we'll use later. pref = string.ascii_lowercase[:rank_ - 2] # Outer product of points (..., N, N) xxT = tf.einsum('{0}ni,{0}mi->{0}nm'.format(pref), arr, arr) # Inner product of points. (..., N) xTx = tf.einsum('{0}ni,{0}ni->{0}n'.format(pref), arr, arr) # (..., N, N) inner products tiled. xTx_tile = tf.tile(xTx[..., None], (1,) * (rank_ - 1) + (N,)) # Build the permuter. (sigh, no tf.swapaxes yet) permute = list(range(rank_)) permute[-2], permute[-1] = permute[-1], permute[-2] # dists = (x^Tx - 2x^Ty + y^Tx)^(1/2). Note the axis swapping is necessary to 'pair' x^Tx and y^Ty return tf.sqrt(xTx_tile - 2 * xxT + tf.transpose(xTx_tile, permute))