tf.math.approx_min_k
Приближённо возвращает минимальные значения k и их индексы во входном operand.
tf.math.approx_min_k(
operand,
k,
reduction_dimension=-1,
recall_target=0.95,
reduction_input_size_override=-1,
aggregate_to_topk=True,
name=None
)
Подробную информацию об алгоритме см. на странице https://arxiv.org/abs/2206.14286. В настоящее время эта операция оптимизирована только для TPU.
| Аргументы | |
|---|---|
operand | Массив, в котором выполняется поиск минимальных k значений. Тип данных должен быть вещественным числом. |
k | Задаёт количество минимальных k значений. |
reduction_dimension | Целочисленная размерность, вдоль которой выполняется поиск. По умолчанию: -1. |
recall_target | Целевое значение полноты для аппроксимации. |
reduction_input_size_override | Если задано положительное значение, оно переопределяет размер, определяемый с помощью operand[reduction_dim] при оценке полноты. Этот параметр полезен, когда заданное operand представляет собой лишь подмножество общего вычисления в SPMD- или распределённых конвейерах, где истинный размер входных данных нельзя определить по форме operand. |
aggregate_to_topk | Если значение равно true, объединяет приближённые результаты до top-k. Если значение равно false, возвращает приближённые результаты. Число приближённых результатов определяется реализацией и не меньше заданного k. |
name | Необязательное имя операции. |
| Возвращаемое значение | |
|---|---|
Кортеж из двух массивов. Массивы содержат наименьшие значения k и соответствующие индексы вдоль reduction_dimension входного operand. Размерности массивов совпадают с размерностями входного operand, за исключением reduction_dimension: если aggregate_to_topk равно true, размерность свёртки равна k; в противном случае она не меньше k, размер которой определяется реализацией. |
Рекомендуем оборачивать approx_min_k в jit. Ниже приведён пример поиска ближайших соседей по квадрату расстояния l2:
import tensorflow as tf
@tf.function(jit_compile=True)
def l2_ann(qy, db, half_db_norms, k=10, recall_target=0.95):
dists = half_db_norms - tf.einsum('ik,jk->ij', qy, db)
return tf.nn.approx_min_k(dists, k=k, recall_target=recall_target)
qy = tf.random.uniform((256,128))
db = tf.random.uniform((2048,128))
half_db_norms = tf.norm(db, axis=1) / 2
dists, neighbors = l2_ann(qy, db, half_db_norms)В приведённом выше примере мы вычисляем db_norms/2 - dot(qy, db^T) вместо qy^2 - 2 dot(qy, db^T) + db^2 по соображениям производительности. Первый вариант требует меньше арифметических операций и возвращает тот же набор соседей.
© 2022 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 4.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/api_docs/python/tf/math/approx_min_k
