Вычислять попарное расстояние в партии без репликации тензора в Tensorflow?

Я хочу вычислить парное квадратное расстояние партии функции в Tensorflow. У меня есть простая реализация с использованием операций + и * путем разбиения исходного тензора:

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 

Эта функция принимает в качестве входных данных две матрицы размера (m, d) и (n, d) и вычисляет квадрат расстояния между каждым вектором строки. Вывод представляет собой матрицу размера (m, n) с элементом 'd_ij = dist (x_i, y_j)'.

Проблема в том, что у меня большие партии и сильно тусклые черты, «m, n, d», реплицирующие тензор, потребляют много памяти. Я ищу другой способ реализовать это без увеличения использования памяти и просто сохранить только конечный тензор расстояния. Вид двойной петли первоначального тензора.

Вы можете использовать некоторую линейную алгебру, чтобы превратить ее в матричные операционные системы. Обратите внимание, что вам нужна матрица D где a[i]i я строка вашей исходной матрицы и

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

Вы можете переписать это в

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

Где r[i] – квадратная норма i го ряда исходной матрицы.

В системе, поддерживающей стандартные правила вещания, вы можете рассматривать r как вектор-столбец и записывать D как

 D = r - 2 AA' + r' 

В TensorFlow вы можете записать это как

 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) 

результат

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

Использование 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 

Одна вещь, которую я заметил, это то, что это решение с использованием tf.squared_difference дает мне память (OOM) для очень больших векторов, в то время как подход by @YaroslavBulatov этого не делает. Итак, я думаю, что разложение операции дает меньший объем памяти (который, как я думал, squared_difference будет лучше справляться с ним под капотом).