Вычислять попарное расстояние в партии без репликации тензора в 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 'потребляет много памяти.
Я ищу другой способ реализовать это без увеличения использования памяти и просто сохранить только конечный тензор расстояния. Вид двойной петли исходного тензора.
Ответы
Ответ 1
Вы можете использовать некоторую линейную алгебру, чтобы превратить ее в матричные операционные системы. Обратите внимание, что вам нужна матрица D, где a[i] - это строка i th вашей исходной матрицы и
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 A A' + 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)
Ответ 2
Использование 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) для очень больших векторов, в то время как подход @Ярослав Булатов не делает. Итак, я думаю, что разложение операции дает меньший объем памяти (который, как я думал, squared_difference будет лучше работать под капотом).
Ответ 3
Если вы хотите вычислить другой метод, измените порядок модулей 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
Ответ 4
Вот более общее решение для двух тензоров координат A и 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
Обратите внимание, что это квадратное расстояние. Если вы хотите изменить это на евклидово расстояние, выполните tf.sqrt в результате. Если вы хотите это сделать, не забудьте добавить небольшую константу, чтобы компенсировать неустойчивость с плавающей запятой: dist = tf.sqrt(squared_dist(A, B) + 1e-6).
