tf.tensor_scatter_nd_update

Разброс updates в существующий тензор в соответствии с indices.

tf.tensor_scatter_nd_update(
    tensor, indices, updates, name=None
)

Эта операция создает новый тензор, применяя разброс updates к входному тензору tensor. Это похоже на присвоение по индексу.

# Not implemented: tensors cannot be updated inplace.
tensor[indices] = updates

Если на ЦП найден индекс за пределами границ, возвращается ошибка.

• Если найден индекс за пределами границ, этот индекс игнорируется.
• Порядок применения обновлений не определён, поэтому результат будет не детерминированным, если indices содержит дубликаты.

Эта операция очень похожа на tf.scatter_nd, за исключением того, что обновления разбрасываются по существующему тензору (в отличие от нулевого тензора). Если память для существующего тензора не может быть повторно использована, создаётся копия и обновляется.

В общем:

• indices — это целочисленный тензор — индексы для обновления в tensor.
• indices имеет **по крайней мере две** оси, последняя ось — глубина векторов индексов.
• Для каждого вектора индексов в indices существует соответствующая запись в updates.
• Если длина векторов индексов соответствует рангу tensor, то векторы индексов указывают на скаляры в tensor, и каждое обновление — скаляр.
• Если длина векторов индексов меньше ранга tensor, то векторы индексов указывают на срезы tensor, а форма обновлений должна совпадать с этим срезом.

В целом, это приводит к следующим ограничениям на форму:

assert tf.rank(indices) >= 2
index_depth = indices.shape[-1]
batch_shape = indices.shape[:-1]
assert index_depth <= tf.rank(tensor)
outer_shape = tensor.shape[:index_depth]
inner_shape = tensor.shape[index_depth:]
assert updates.shape == batch_shape + inner_shape

Типичное использование часто намного проще, чем эта общая форма, и его лучше понять, начав с простых примеров:

Скалярные обновления

Самое простое использование — вставка скалярных элементов в тензор по индексу. В этом случае, index_depth должна быть равна рангу входного тензора tensor, каждый столбец indices — это индекс в оси входного тензора tensor.

В этом простом случае ограничения на форму следующие:

num_updates, index_depth = indices.shape.as_list()
assert updates.shape == [num_updates]
assert index_depth == tf.rank(tensor)`

Например, чтобы вставить 4 разбросанных элемента в тензор ранга 1 с 8 элементами.

Эта операция разброса будет выглядеть так:

tensor = [0, 0, 0, 0, 0, 0, 0, 0]    # tf.rank(tensor) == 1
indices = [[1], [3], [4], [7]]       # num_updates == 4, index_depth == 1
updates = [9, 10, 11, 12]            # num_updates == 4
print(tf.tensor_scatter_nd_update(tensor, indices, updates))
tf.Tensor([ 0 9  0 10  11  0  0 12], shape=(8,), dtype=int32)

Длина (первая ось) updates должна быть равна длине indices: num_updates. Это количество вставляемых обновлений. Каждое скалярное обновление вставляется в tensor по указанному индексом расположению.

Для входного тензора более высокого ранга tensor скалярные обновления можно вставить, используя index_depth, которая соответствует tf.rank(tensor):

tensor = [[1, 1], [1, 1], [1, 1]]    # tf.rank(tensor) == 2
indices = [[0, 1], [2, 0]]           # num_updates == 2, index_depth == 2
updates = [5, 10]                    # num_updates == 2
print(tf.tensor_scatter_nd_update(tensor, indices, updates))
tf.Tensor(
    [[ 1  5]
     [ 1  1]
     [10  1]], shape=(3, 2), dtype=int32)

Обновления срезов

Когда входной тензор tensor имеет более одной оси, разброс можно использовать для обновления целых срезов.

В этом случае полезно рассматривать входной тензор tensor как двухуровневый массив-из-массивов. Форма этого двухуровневого массива разделяется на outer_shape и inner_shape.

indices индексирует внешний уровень входного тензора (outer_shape) и заменяет подмассив в этом месте соответствующим элементом из списка updates. Форма каждого обновления равна inner_shape.

При обновлении списка срезов ограничения на форму следующие:

num_updates, index_depth = indices.shape.as_list()
outer_shape = tensor.shape[:index_depth]
inner_shape = tensor.shape[index_depth:]
assert updates.shape == [num_updates, inner_shape]

Например, чтобы обновить строки (6, 3) tensor:

tensor = tf.zeros([6, 3], dtype=tf.int32)

Используйте глубину индексарования равную 1.

indices = tf.constant([[2], [4]])     # num_updates == 2, index_depth == 1
num_updates, index_depth = indices.shape.as_list()

outer_shape — это 6, внутренняя форма — 3:

outer_shape = tensor.shape[:index_depth]
inner_shape = tensor.shape[index_depth:]

Индексируется 2 строки, поэтому необходимо предоставить 2 updates. Каждое обновление должно иметь форму, соответствующую форме inner_shape.

# num_updates == 2, inner_shape==3
updates = tf.constant([[1, 2, 3],
                       [4, 5, 6]])

В совокупности это дает:

tf.tensor_scatter_nd_update(tensor, indices, updates).numpy()
array([[0, 0, 0],
       [0, 0, 0],
       [1, 2, 3],
       [0, 0, 0],
       [4, 5, 6],
       [0, 0, 0]], dtype=int32)

Ещё примеры обновлений срезов

Тензор, представляющий пакет видеоклипов одинакового размера, естественным образом имеет 5 осей: [batch_size, time, width, height, channels].

Например:

batch_size, time, width, height, channels = 13,11,7,5,3
video_batch = tf.zeros([batch_size, time, width, height, channels])

Чтобы заменить выборку видеоклипов:

• Используйте глубину индексации 1 (индексируя outer_shape: [batch_size])
• Предоставьте обновления, каждое с формой, соответствующей форме inner_shape: [time, width, height, channels].

Чтобы заменить первые два клипа единицами:

indices = [[0],[1]]
new_clips = tf.ones([2, time, width, height, channels])
tf.tensor_scatter_nd_update(video_batch, indices, new_clips)

Чтобы заменить выборку кадров в видеоклипах:

• indices должна иметь глубину индексации 2 для outer_shape: [batch_size, time].
• updates должна иметь форму, как список изображений. Каждое обновление должно иметь форму, соответствующую форме inner_shape: [width, height, channels].

Чтобы заменить первый кадр первых трёх видеоклипов:

indices = [[0, 0], [1, 0], [2, 0]] # num_updates=3, index_depth=2
new_images = tf.ones([
  # num_updates=3, inner_shape=(width, height, channels)
  3, width, height, channels])
tf.tensor_scatter_nd_update(video_batch, indices, new_images)

Сложенные индексы

В простых случаях удобно представлять indices и updates как списки, но это не является строгим требованием. Вместо плоского списка num_updates, indices и updates могут быть сложены в batch_shape. Это batch_shape - это все оси indices, за исключением самой внутренней index_depth оси.

index_depth = indices.shape[-1]
batch_shape = indices.shape[:-1]

Примечание: Единственное исключение заключается в том, что batch_shape не может быть []. Нельзя обновить одиночный индекс, передавая индексы с формой [index_depth].

updates должна иметь соответствующую форму batch_shape (оси до inner_shape).

assert updates.shape == batch_shape + inner_shape

Примечание: Результат эквивалентен сглаживанию осей batch_shape тензоров indices и updates. Это обобщение просто устраняет необходимость в преобразованиях формы, когда естественнее строить "сложенные" индексы и обновления.

При этом полные ограничения на форму следующие:

assert tf.rank(indices) >= 2
index_depth = indices.shape[-1]
batch_shape = indices.shape[:-1]
assert index_depth <= tf.rank(tensor)
outer_shape = tensor.shape[:index_depth]
inner_shape = tensor.shape[index_depth:]
assert updates.shape == batch_shape + inner_shape

Например, чтобы нарисовать X на матрице (5,5), начните с этих индексов:

tensor = tf.zeros([5,5])
indices = tf.constant([
 [[0,0],
  [1,1],
  [2,2],
  [3,3],
  [4,4]],
 [[0,4],
  [1,3],
  [2,2],
  [3,1],
  [4,0]],
])
indices.shape.as_list()  # batch_shape == [2, 5], index_depth == 2
[2, 5, 2]

Здесь indices не имеют формы [num_updates, index_depth], а имеют форму batch_shape+[index_depth].

Поскольку index_depth равна рангу tensor:

• outer_shape — это (5,5)
• inner_shape — это () — каждое обновление является скаляром
• updates.shape — это batch_shape + inner_shape == (5,2) + ()

updates = [
  [1,1,1,1,1],
  [1,1,1,1,1],
]

Объединив это, получаем:

tf.tensor_scatter_nd_update(tensor, indices, updates).numpy()
array([[1., 0., 0., 0., 1.],
       [0., 1., 0., 1., 0.],
       [0., 0., 1., 0., 0.],
       [0., 1., 0., 1., 0.],
       [1., 0., 0., 0., 1.]], dtype=float32)