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

Если на CPU обнаружен выход за пределы границ, возвращается ошибка.

• Если обнаружен выход за пределы границ, индекс игнорируется.
• Порядок применения обновлений не определён, поэтому результат будет непредсказуемым, если 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()
inner_shape = tensor.shape[:index_depth]
outer_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 должен иметь соответствующую форму (оси до 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)