tf.nn.embedding_lookup_sparse

Выполняет поиск эмбеддингов для заданных идентификаторов и весов из списка тензоров.

tf.nn.embedding_lookup_sparse(
    params,
    sp_ids,
    sp_weights,
    combiner=None,
    max_norm=None,
    name=None,
    allow_fast_lookup=False
)

params — это плотный тензор или список плотных тензоров, а sp_ids — это двумерный tf.SparseTensor или tf.RaggedTensor, указывающий индексы params для выборки.

Эту операцию лучше всего объяснить на примере. Предположим, что params — это таблица эмбеддингов размера (4, 2), а sp_ids содержит 3 строки. Поскольку sp_ids является разреженным или ragged-тензором, не в каждой строке одинаковое количество элементов. Форма результата — (3, 2). Каждая строка sp_ids представляет собой список индексов, каждый из которых выбирает строку params. Для заданной строки sp_ids строки params выбираются по индексам в sp_ids, а затем объединяются посредством суммирования или вычисления среднего.

params = tf.constant([[1, 2], [3, 4], [5, 6], [7, 8]], dtype=tf.float32)
sp_ids = tf.SparseTensor(indices=[[0, 0], [0, 1], [1, 0], [2, 0]],
                         values=[0, 1, 3, 2], dense_shape=(3, 2))
tf.nn.embedding_lookup_sparse(params, sp_ids, sp_weights=None,
                              combiner='sum').numpy()
array([[4., 6.], [7., 8.], [5., 6.]], dtype=float32)

В этом примере sp_ids содержит 3 строки, поэтому результат также содержит 3 строки. Строка 0 sp_ids содержит значения 0 и 1, поэтому из params выбираются строки 0 и 1 — [1, 2] и [3, 4]. Поскольку combiner='sum', строки суммируются, в результате чего получается строка результата [4, 6].

Поскольку в строках 1 и 2 sp_ids содержится только по одному значению, в качестве соответствующей строки результата просто выбирается строка params. Строка 1 содержит значение 3, поэтому выбираются элементы params [7, 8], а строка 2 содержит значение 2, поэтому выбираются элементы params [5, 6].

Если задан sparse_weights, он должен иметь ту же форму, что и sp_ids. sparse_weights используется для назначения веса каждому срезу params. Например:

params = tf.constant([[1, 2], [3, 4], [5, 6], [7, 8]], dtype=tf.float32)
sp_ids = tf.SparseTensor(indices=[[0, 0], [0, 1], [1, 0], [2, 0]],
                         values=[0, 1, 3, 2], dense_shape=(3, 2))
sparse_weights = tf.SparseTensor(indices=[[0, 0], [0, 1], [1, 0], [2, 0]],
                                 values=[0.1, 1.0, 0.5, 2.0],
                                 dense_shape=(3, 2))
tf.nn.embedding_lookup_sparse(params, sp_ids, sp_weights=sparse_weights,
                              combiner='sum').numpy()
array([[3.1, 4.2], [3.5, 4.], [10., 12.]], dtype=float32)

В общем случае params может иметь форму (p0, ..., pn), а sp_ids может содержать M строк, каждая из которых может иметь любое количество элементов. Форма результата — (M, p1, ..., pn). Каждый срез результата output[i, ...] вычисляется следующим образом: аргумент combiner используется для объединения значений params[sp_ids[i, j], ...] * sparse_weights[i, j] для каждого j в range(0, len(sp_ids[i])), например посредством суммирования или вычисления среднего этих значений.

Эта операция предполагает, что для каждой строки плотного тензора, представленного sp_ids, имеется хотя бы один идентификатор (то есть нет строк с пустыми признаками), а все индексы sp_ids расположены в каноническом порядке по строкам.

sp_ids и sp_weights (если не None) являются SparseTensor или RaggedTensor ранга 2. Для SpareTensor с выровненными по левому краю ненулевыми элементами, которые можно представить как RaggedTensor, использование RaggedTensor может повысить производительность.

Эта операция предполагает, что все значения id лежат в диапазоне [0, p0), где p0 — это params.shape[0]. Если вам нужна версия этой операции, отбрасывающая значения id меньше 0, см. tf.nn.safe_embedding_lookup_sparse

Если len(params) > 1, каждый элемент sp_ids распределяется между элементами params согласно стратегии разбиения "div", то есть идентификаторы назначаются разделам последовательно. Например, 13 идентификаторов распределяются по 5 разделам так: [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10], [11, 12]].

Если пространство идентификаторов не делится на число разделов без остатка, каждому из первых (max_id + 1) % len(params) разделов назначается на один идентификатор больше.

[0, 0]: id 1, weight 2.0
[0, 1]: id 3, weight 0.5
[1, 0]: id 0, weight 1.0
[2, 3]: id 1, weight 3.0

output[0, :] = (params[1, :] * 2.0 + params[3, :] * 0.5) / (2.0 + 0.5)
output[1, :] = (params[0, :] * 1.0) / 1.0
output[2, :] = (params[1, :] * 3.0) / 3.0