python - TensorFlow 中稀疏加法的 scatter_nd add() 示例

Question

我在将 tf.scatter_nd_add() 应用于二维张量时遇到困难。该文档有点不清楚，没有包含稀疏更新的示例，但仅包含完整切片更新的示例。

我的情况如下:

• 更新 - 形状为 [None, 6] 的 2D 张量
• indices - 形状为 [None, 6]
的二维张量
• ref - 形状为 [None, 6]
零点的二维变量

保证 updates、indices 和 ref 的第一个维度始终相等，但该维度的大小可以变化.我要执行的更新看起来像

for i, j:
    k = indices[i][j]
    ref[i][k] += updates[i][j] 

请注意，indices 包含重复项。 tf.scatter_nd_add(ref, indices, updates) 提示形状不匹配，我无法弄清楚我需要如何重组张量才能执行更新。

Answer 1

我想通了。 indices 中的每个 2D 条目实际上必须指定将在 ref 中更新的绝对位置。这意味着 indices 必须是 3D，然后非矢量化更新看起来像:

for i, j:
    r, k = indices[i][j]
    ref[r][k] += updates[i][j]

在上面的问题中，恰好 r 总是等于 i。

这是具有不同形状的完整 Tensorflow 实现。为清楚起见，在以下示例中，col_indices 对应于原始问题中的 indices:

import tensorflow as tf
import numpy as np

updates     = tf.placeholder(dtype=tf.float32,  shape=[None, 6])
col_indices = tf.placeholder(dtype=tf.int32,    shape=[None, 6])
row_indices = tf.cumsum(tf.ones_like(col_indices), axis=0, exclusive=True)
indices     = tf.concat([tf.expand_dims(row_indices, axis=-1), 
                         tf.expand_dims(col_indices, axis=-1)], axis=-1)

tmp_var     = tf.Variable(0, trainable=False, dtype=tf.float32, validate_shape=False)
ref         = tf.assign(tmp_var, tf.zeros_like(updates), validate_shape=False)
# This makes sure that ref is always 0 before scatter_nd_add() runs
with tf.control_dependencies([target_var]):
  result = tf.scatter_nd_add(ref, indices, updates)

sess = tf.Session()
sess.run(tf.global_variables_initializer())

# Create example input data 
np_input = np.arange(0, 6, 1, dtype=np.int32)
np_input = np.tile(np_input[None,:], [10, 1])

res = sess.run(result, feed_dict={updates: np_input, col_indices: np_input})
print(res)