Учитывая тензоры a формы (n, f) и b формы (m, f), я создал функцию для вычисления евклидовых расстояний между этими двумя тензорами
import tensorflow as tf
nr = tf.reduce_sum(tf.square(a), 1)
nw = tf.reduce_sum(tf.square(b), 1)
nr = tf.reshape(nr, [-1, 1])
nw = tf.reshape(nw, [1, -1])
res = nr - 2*tf.matmul(a, b, False, True) + nw
res = tf.argmin(res, axis=1)
Пока все хорошо, код работает немного быстро (я получил лучшую производительность с cKDTree, когда n= 1000, m=1600, f=4, но сейчас это не проблема). Я проверю производительность по сравнению с различными размерами входных данных позже.
В этом примере тензор b - это сплющенная версия тензора ранга 3. Я делаю это, чтобы иметь возможность оценить евклидовы расстояния, используя два тензора с одинаковым рангом (что проще). Но после оценки расстояний мне нужно знать, где в исходном тензоре находится каждый из ближайших элементов. Для этого я создал собственную лямбда-функцию fn для преобразования обратно в тензорные координаты ранга 3.
fn = lambda x: (x//N, x%N)
# This map takes a enormous amount of time
out = tf.map_fn(fn, res, dtype=(tf.int64, tf.int64))
return tf.stack(out, axis=1)
Но, к сожалению, для этого tf.map_fn требуется ОГРОМНОЕ время для запуска, около 300 мс .
Просто для сравнения, если я выполню np.apply_along_axis в наборе данных, который содержит точно те же данные (но массив numpy), объем отпечатка едва заметен, около 50 микросекунды против 300 мс эквивалентного тензорного потока.
Для этого mapping?
TF версии 2.1.0 есть лучшие подходы, и CUDA включен и работает.
Просто добавим немного времени
%timeit eucl_dist_tf_vecmap(R_tf, W_tf)
28.1 ms ± 128 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit eucl_dist_tf_nomap(R_tf, W_tf)
2.07 ms ± 122 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit eucl_dist_ckdtree_applyaxis(R, W)
878 µs ± 2.34 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit eucl_dist_ckdtree_noapplyaxis(R, W)
817 µs ± 51 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
Первые два времени используют показанную здесь пользовательскую функцию: первое с vectorized_map, второе без vectorized_map и stack (накладные расходы на vectorized_map, протестировано.
И последние два раза являются реализациями, основанными на cKDTree Сципи. Первый использует np.apply_along_axis точно так же, как в векторизованной карте. Мы видим, что издержки намного меньше в массиве numpy.