tf.math.approx_max_k
Приближённо возвращает максимальные k значений и их индексы во входном operand.
tf.math.approx_max_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 | Массив, в котором выполняется поиск max-k. Должен иметь вещественный числовой тип. |
k | Задаёт количество элементов max-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_max_k в jit. Ниже приведён пример поиска максимального скалярного произведения (MIPS):
import tensorflow as tf
@tf.function(jit_compile=True)
def mips(qy, db, k=10, recall_target=0.95):
dists = tf.einsum('ik,jk->ij', qy, db)
# returns (f32[qy_size, k], i32[qy_size, k])
return tf.nn.approx_max_k(dists, k=k, recall_target=recall_target)
qy = tf.random.uniform((256,128))
db = tf.random.uniform((2048,128))
dot_products, neighbors = mips(qy, db, k=20)
© 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_max_k
