¿Calcular la distancia de pares en un lote sin replicar el tensor en Tensorflow?

Quiero calcular la distancia cuadrada por pares de un lote de características en Tensorflow. Tengo una implementación simple utilizando las operaciones + y * al hacer un mosaico del tensor original:

def pairwise_l2_norm2(x, y, scope=None): with tf.op_scope([x, y], scope, 'pairwise_l2_norm2'): size_x = tf.shape(x)[0] size_y = tf.shape(y)[0] xx = tf.expand_dims(x, -1) xx = tf.tile(xx, tf.pack([1, 1, size_y])) yy = tf.expand_dims(y, -1) yy = tf.tile(yy, tf.pack([1, 1, size_x])) yy = tf.transpose(yy, perm=[2, 1, 0]) diff = tf.sub(xx, yy) square_diff = tf.square(diff) square_dist = tf.reduce_sum(square_diff, 1) return square_dist 

Esta función toma como entrada dos matrices de tamaño (m, d) y (n, d) y calcula la distancia al cuadrado entre cada vector de fila. La salida es una matriz de tamaño (m, n) con el elemento ‘d_ij = dist (x_i, y_j)’.

El problema es que tengo un lote grande y características de alta intensidad ‘m, n, d’ replicando que el tensor consume mucha memoria. Estoy buscando otra forma de implementar esto sin boost el uso de la memoria y solo almacenar el tensor de distancia final. Tipo de doble bucle del tensor original.

Puedes usar un poco de álgebra lineal para convertirlo en operaciones de matriz. Tenga en cuenta que lo que necesita es la matriz D donde a[i] es la fila i th de su matriz original y

 D[i,j] = (a[i]-a[j])(a[i]-a[j])' 

Puedes reescribir eso en

 D[i,j] = r[i] - 2 a[i]a[j]' + r[j] 

Donde r[i] es la norma cuadrada de la fila i de la matriz original.

En un sistema que admite reglas de transmisión estándar, puede tratar r como un vector de columna y escribir D como

 D = r - 2 AA' + r' 

En TensorFlow puedes escribir esto como

 A = tf.constant([[1, 1], [2, 2], [3, 3]]) r = tf.reduce_sum(A*A, 1) # turn r into column vector r = tf.reshape(r, [-1, 1]) D = r - 2*tf.matmul(A, tf.transpose(A)) + tf.transpose(r) sess = tf.Session() sess.run(D) 

resultado

 array([[0, 2, 8], [2, 0, 2], [8, 2, 0]], dtype=int32) 

Usando squared_difference :

 def squared_dist(A): expanded_a = tf.expand_dims(A, 1) expanded_b = tf.expand_dims(A, 0) distances = tf.reduce_sum(tf.squared_difference(expanded_a, expanded_b), 2) return distances 

Una cosa que noté es que esta solución que usa tf.squared_difference me deja sin memoria (OOM) para vectores muy grandes, mientras que el enfoque de @YaroslavBulatov no lo hace. Por lo tanto, creo que al descomponer la operación se obtiene una huella de memoria más pequeña (lo que pensé que squared_difference manejaría mejor bajo el capó).

Aquí hay una solución más general para dos tensores de coordenadas A y B :

 def squared_dist(A, B): assert A.shape.as_list() == B.shape.as_list() row_norms_A = tf.reduce_sum(tf.square(A), axis=1) row_norms_A = tf.reshape(row_norms_A, [-1, 1]) # Column vector. row_norms_B = tf.reduce_sum(tf.square(B), axis=1) row_norms_B = tf.reshape(row_norms_B, [1, -1]) # Row vector. return row_norms_A - 2 * tf.matmul(A, tf.transpose(B)) + row_norms_B 

Tenga en cuenta que esta es la distancia cuadrada. Si desea cambiar esto a la distancia euclidiana, realice un tf.sqrt en el resultado. Si desea hacerlo, no olvide agregar una pequeña constante para compensar las inestabilidades de punto flotante: dist = tf.sqrt(squared_dist(A, B) + 1e-6) .

Si desea calcular otro método, cambie el orden de los módulos tf.

 def compute_euclidean_distance(x, y): size_x = x.shape.dims[0] size_y = y.shape.dims[0] for i in range(size_x): tile_one = tf.reshape(tf.tile(x[i], [size_y]), [size_y, -1]) eu_one = tf.expand_dims(tf.sqrt(tf.reduce_sum(tf.pow(tf.subtract(tile_one, y), 2), axis=1)), axis=0) if i == 0: d = eu_one else: d = tf.concat([d, eu_one], axis=0) return d