tf.compat.v1.gather_nd
Извлечь фрагменты из params в тензор с формой, заданной indices.
tf.compat.v1.gather_nd(
params, indices, name=None, batch_dims=0
)
indices — это Tensor индексов в params. Векторы индексов расположены вдоль последней оси indices.
Это аналогично tf.gather, в котором indices определяет фрагменты в первом измерении params. В tf.gather_nd, indices определяет фрагменты в первых N измерениях params, где N = indices.shape[-1].
Извлечение скаляров
В простейшем случае векторы в indices индексируют весь ранг params:
tf.gather_nd(
indices=[[0, 0],
[1, 1]],
params = [['a', 'b'],
['c', 'd']]).numpy()
array([b'a', b'd'], dtype=object)В этом случае результат имеет на 1 ось меньше, чем indices, и каждый вектор индекса заменяется скаляром, индексированным из params.
В этом случае соотношение форм:
index_depth = indices.shape[-1] assert index_depth == params.shape.rank result_shape = indices.shape[:-1]
Если у indices ранг K, полезно рассматривать indices как (K-1)-мерный тензор индексов в params.
Извлечение фрагментов
Если векторы индексов не индексируют весь ранг params, то каждое местоположение в результате содержит фрагмент params. Этот пример собирает строки из матрицы:
tf.gather_nd(
indices = [[1],
[0]],
params = [['a', 'b', 'c'],
['d', 'e', 'f']]).numpy()
array([[b'd', b'e', b'f'],
[b'a', b'b', b'c']], dtype=object)Здесь indices содержит [2] векторов индексов, каждый длиной 1. Векторы индексов ссылаются на строки матрицы params. Каждая строка имеет форму [3], поэтому форма выходных данных — [2, 3].
В этом случае соотношение форм:
index_depth = indices.shape[-1] outer_shape = indices.shape[:-1] assert index_depth <= params.shape.rank inner_shape = params.shape[index_depth:] output_shape = outer_shape + inner_shape
В этом случае полезно рассматривать результаты как тензоры-тензоров. Форма внешнего тензора задаётся ведущими измерениями indices. В то время как форма внутренних тензоров — форма одного фрагмента.
Пачки
Кроме того, как у params, так и у indices могут быть M ведущих измерений пачки, которые точно совпадают. В этом случае batch_dims должно быть установлено в M.
Например, чтобы собрать одну строку из каждой из пачки матриц, вы можете установить ведущие элементы векторов индексов в их расположение в пачке:
tf.gather_nd(
indices = [[0, 1],
[1, 0],
[2, 4],
[3, 2],
[4, 1]],
params=tf.zeros([5, 7, 3])).shape.as_list()
[5, 3]Аргумент batch_dims позволяет опустить эти ведущие размеры расположения из индекса:
tf.gather_nd(
batch_dims=1,
indices = [[1],
[0],
[4],
[2],
[1]],
params=tf.zeros([5, 7, 3])).shape.as_list()
[5, 3]Это эквивалентно вызову отдельного gather_nd для каждого расположения в измерениях пачки.
params=tf.zeros([5, 7, 3]) indices=tf.zeros([5, 1]) batch_dims = 1 index_depth = indices.shape[-1] batch_shape = indices.shape[:batch_dims] assert params.shape[:batch_dims] == batch_shape outer_shape = indices.shape[batch_dims:-1] assert index_depth <= params.shape.rank inner_shape = params.shape[batch_dims + index_depth:] output_shape = batch_shape + outer_shape + inner_shape output_shape.as_list() [5, 3]
Дополнительные примеры
Индексирование в 3-мерный тензор:
tf.gather_nd(
indices = [[1]],
params = [[['a0', 'b0'], ['c0', 'd0']],
[['a1', 'b1'], ['c1', 'd1']]]).numpy()
array([[[b'a1', b'b1'],
[b'c1', b'd1']]], dtype=object)tf.gather_nd(
indices = [[0, 1], [1, 0]],
params = [[['a0', 'b0'], ['c0', 'd0']],
[['a1', 'b1'], ['c1', 'd1']]]).numpy()
array([[b'c0', b'd0'],
[b'a1', b'b1']], dtype=object)tf.gather_nd(
indices = [[0, 0, 1], [1, 0, 1]],
params = [[['a0', 'b0'], ['c0', 'd0']],
[['a1', 'b1'], ['c1', 'd1']]]).numpy()
array([b'b0', b'b1'], dtype=object)Примеры ниже относятся к случаю, когда у индексов есть дополнительные ведущие измерения. Если у 'params' и 'indices' есть ведущие измерения пачки, используйте параметр 'batch_dims' для выполнения gather_nd в режиме пачки.
Индексирование пачкой в матрице:
tf.gather_nd(
indices = [[[0, 0]], [[0, 1]]],
params = [['a', 'b'], ['c', 'd']]).numpy()
array([[b'a'],
[b'b']], dtype=object)Индексирование пачкой срезов в матрице:
tf.gather_nd(
indices = [[[1]], [[0]]],
params = [['a', 'b'], ['c', 'd']]).numpy()
array([[[b'c', b'd']],
[[b'a', b'b']]], dtype=object)Индексирование пачкой в 3-мерный тензор:
tf.gather_nd(
indices = [[[1]], [[0]]],
params = [[['a0', 'b0'], ['c0', 'd0']],
[['a1', 'b1'], ['c1', 'd1']]]).numpy()
array([[[[b'a1', b'b1'],
[b'c1', b'd1']]],
[[[b'a0', b'b0'],
[b'c0', b'd0']]]], dtype=object)tf.gather_nd(
indices = [[[0, 1], [1, 0]], [[0, 0], [1, 1]]],
params = [[['a0', 'b0'], ['c0', 'd0']],
[['a1', 'b1'], ['c1', 'd1']]]).numpy()
array([[[b'c0', b'd0'],
[b'a1', b'b1']],
[[b'a0', b'b0'],
[b'c1', b'd1']]], dtype=object)tf.gather_nd(
indices = [[[0, 0, 1], [1, 0, 1]], [[0, 1, 1], [1, 1, 0]]],
params = [[['a0', 'b0'], ['c0', 'd0']],
[['a1', 'b1'], ['c1', 'd1']]]).numpy()
array([[b'b0', b'b1'],
[b'd0', b'c1']], dtype=object)Примеры с пакетными 'params' и 'indices':
tf.gather_nd(
batch_dims = 1,
indices = [[1],
[0]],
params = [[['a0', 'b0'],
['c0', 'd0']],
[['a1', 'b1'],
['c1', 'd1']]]).numpy()
array([[b'c0', b'd0'],
[b'a1', b'b1']], dtype=object)tf.gather_nd(
batch_dims = 1,
indices = [[[1]], [[0]]],
params = [[['a0', 'b0'], ['c0', 'd0']],
[['a1', 'b1'], ['c1', 'd1']]]).numpy()
array([[[b'c0', b'd0']],
[[b'a1', b'b1']]], dtype=object)tf.gather_nd(
batch_dims = 1,
indices = [[[1, 0]], [[0, 1]]],
params = [[['a0', 'b0'], ['c0', 'd0']],
[['a1', 'b1'], ['c1', 'd1']]]).numpy()
array([[b'c0'],
[b'b1']], dtype=object)См. также tf.gather.
| Аргументы | |
|---|---|
params | Tensor. Тензор, из которого извлекаются значения. |
indices | Tensor. Должен быть одного из следующих типов: int32, int64. Тензор индексов. |
name | Имя операции (необязательно). |
batch_dims | Целое число или скалярный 'Tensor'. Количество измерений пачки. |
| Возвращаемые значения | |
|---|---|
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/compat/v1/gather_nd
