tf.gather

Извлечение срезов из params по оси axis в соответствии с индексами. (аргументы устарели)

tf.gather(
    params, indices, validate_indices=None, axis=None, batch_dims=0, name=None
)

Используется в блокнотах

Используется в руководстве	Используется в учебниках
Введение в графы и tf.function Введение в срезы тензоров Типы расширений Обработка BERT с помощью TF Text Токенизаторы подслов	Простое распознавание аудио: распознавание ключевых слов Сегментация изображений Руководство по многоуровневому моделированию в TensorFlow Probability MoViNet для распознавания действий в потоке Эффективное в клиентской части обучение больших моделей с помощью `federated_select` и разреженной агрегации

Используется в руководстве

Используется в учебниках

Извлечение срезов из params по оси axis в соответствии с indices. indices должен быть целочисленным тензором любой размерности (часто 1-D).

Tensor.getitem работает для скаляров, tf.newaxis и python-срезами

tf.gather расширяет индексацию для обработки тензоров индексов.

В самом простом случае он идентичен скалярной индексации:

params = tf.constant(['p0', 'p1', 'p2', 'p3', 'p4', 'p5'])
params[3].numpy()
b'p3'
tf.gather(params, 3).numpy()
b'p3'

Наиболее распространённый случай – это передача тензора индексов по одной оси (это нельзя выразить как python-срез, поскольку индексы не являются последовательными):

indices = [2, 0, 2, 5]
tf.gather(params, indices).numpy()
array([b'p2', b'p0', b'p2', b'p5'], dtype=object)

Индексы могут иметь любую форму. Когда у params одна ось, форма выходного тензора равна форме входного тензора:

tf.gather(params, [[2, 0], [2, 5]]).numpy()
array([[b'p2', b'p0'],
       [b'p2', b'p5']], dtype=object)

У params также может быть любая форма. gather может выбирать срезы по любой оси в зависимости от аргумента axis (который по умолчанию равен 0). Ниже он используется для извлечения первых строк, а затем столбцов из матрицы:

params = tf.constant([[0, 1.0, 2.0],
                      [10.0, 11.0, 12.0],
                      [20.0, 21.0, 22.0],
                      [30.0, 31.0, 32.0]])
tf.gather(params, indices=[3,1]).numpy()
array([[30., 31., 32.],
       [10., 11., 12.]], dtype=float32)
tf.gather(params, indices=[2,1], axis=1).numpy()
array([[ 2.,  1.],
       [12., 11.],
       [22., 21.],
       [32., 31.]], dtype=float32)

В общем случае: форма выходного тензора имеет такую же форму, как и входной тензор, с осью индексации, заменённой формой индексов.

def result_shape(p_shape, i_shape, axis=0):
  return p_shape[:axis] + i_shape + p_shape[axis+1:]

result_shape([1, 2, 3], [], axis=1)
[1, 3]
result_shape([1, 2, 3], [7], axis=1)
[1, 7, 3]
result_shape([1, 2, 3], [7, 5], axis=1)
[1, 7, 5, 3]

Вот несколько примеров:

params.shape.as_list()
[4, 3]
indices = tf.constant([[0, 2]])
tf.gather(params, indices=indices, axis=0).shape.as_list()
[1, 2, 3]
tf.gather(params, indices=indices, axis=1).shape.as_list()
[4, 1, 2]

params = tf.random.normal(shape=(5, 6, 7, 8))
indices = tf.random.uniform(shape=(10, 11), maxval=7, dtype=tf.int32)
result = tf.gather(params, indices, axis=2)
result.shape.as_list()
[5, 6, 10, 11, 8]

Потому что каждый индекс берёт срез из params и помещает его в соответствующее место в выходном тензоре. Для примера выше

# For any location in indices
a, b = 0, 1
tf.reduce_all(
    # the corresponding slice of the result
    result[:, :, a, b, :] ==
    # is equal to the slice of `params` along `axis` at the index.
    params[:, :, indices[a, b], :]
).numpy()
True

Обработка пакетов:

Аргумент batch_dims позволяет извлекать разные элементы из каждого элемента пакета.

Использование batch_dims=1 эквивалентно внешнему циклу по первой оси params и indices:

params = tf.constant([
    [0, 0, 1, 0, 2],
    [3, 0, 0, 0, 4],
    [0, 5, 0, 6, 0]])
indices = tf.constant([
    [2, 4],
    [0, 4],
    [1, 3]])

tf.gather(params, indices, axis=1, batch_dims=1).numpy()
array([[1, 2],
       [3, 4],
       [5, 6]], dtype=int32)

Это эквивалентно:

def manually_batched_gather(params, indices, axis):
  batch_dims=1
  result = []
  for p,i in zip(params, indices):
    r = tf.gather(p, i, axis=axis-batch_dims)
    result.append(r)
  return tf.stack(result)
manually_batched_gather(params, indices, axis=1).numpy()
array([[1, 2],
       [3, 4],
       [5, 6]], dtype=int32)

Более высокие значения batch_dims эквивалентны множеству вложенных циклов по внешним осям params и indices. Поэтому общая функция формы

def batched_result_shape(p_shape, i_shape, axis=0, batch_dims=0):
  return p_shape[:axis] + i_shape[batch_dims:] + p_shape[axis+1:]

batched_result_shape(
    p_shape=params.shape.as_list(),
    i_shape=indices.shape.as_list(),
    axis=1,
    batch_dims=1)
[3, 2]

tf.gather(params, indices, axis=1, batch_dims=1).shape.as_list()
[3, 2]

Это естественным образом возникает, если вам нужно использовать индексы операции, такой как tf.argsort или tf.math.top_k, где последняя размерность индексов индексирует последнюю размерность входного тензора в соответствующем месте. В этом случае вы можете использовать tf.gather(values, indices, batch_dims=-1).

См. также:

tf.Tensor.getitem: Прямая операция индексации тензора (t[]), обрабатывающая скаляры и python-срезами tensor[..., 7, 1:-1]
tf.scatter: Коллекция операций, аналогичных __setitem__ (t[i] = x)
tf.gather_nd: Операция, аналогичная tf.gather, но собирающая по нескольким осям одновременно (она может собирать элементы матрицы вместо строк или столбцов)
tf.boolean_mask, tf.where: Бинарная индексация.
tf.slice и tf.strided_slice: Для более низкого уровня доступа к реализации обработки python-срезов в __getitem__ (t[1:-1:2])

Аргументы
`params`	Тензор, из которого необходимо извлечь значения. Должен быть не менее ранга `axis + 1`.
`indices`	Индекс `Tensor`. Должен быть одного из следующих типов: `int32`, `int64`. Значения должны быть в диапазоне `[0, params.shape[axis])`.
`validate_indices`	Устарело, ничего не делает. Индексы всегда проверяются на CPU, но никогда не проверяются на GPU. Предупреждение: На CPU, если обнаружен индекс за пределами границ, генерируется ошибка. На GPU, если обнаружен индекс за пределами границ, в соответствующем значении выходного тензора хранится 0.
`axis`	Ось `Tensor`. Должен быть одного из следующих типов: `int32`, `int64`. Ось `axis` в `params`, для извлечения `indices`. Должно быть больше или равно `batch_dims`. По умолчанию – первая ось, не являющаяся осью пакета. Поддерживает отрицательные индексы.
`batch_dims`	Количество осей пакета. Должно быть меньше или равно `rank(indices)`.
`name`	Имя операции (необязательно).

Возвращаемое значение
Тензор `Tensor`. Имеет тот же тип, что и `params`.

© 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/gather