在tensorflow-r1.2中正确使用“tf.scatter_nd”

Given indices有形状[batch_size, sequence_len], updates有形状[batch_size, sequence_len, sampled_size], to_shape有形状[batch_size, sequence_len, vocab_size], where vocab_size>>sampled_size，我想用tf.scatter来映射updates到一个巨大的张量to_shape，使得to_shape[bs, indices[bs, sz]] = updates[bs, sz]。也就是说，我想绘制地图updates to to_shape一行一行。请注意sequence_len and sampled_size是标量张量，而其他则是固定的。我尝试执行以下操作：

new_tensor = tf.scatter_nd(tf.expand_dims(indices, axis=2), updates, to_shape)

但我收到一个错误：

ValueError: The inner 2 dimension of output.shape=[?,?,?] must match the inner 1 dimension of updates.shape=[80,50,?]: Shapes must be equal rank, but are 2 and 1 for .... with input shapes: [80, 50, 1], [80, 50,?], [3]

你能告诉我如何使用吗scatter_nd适当地？提前致谢！

所以假设你有：

• 一个张量updates有形状[batch_size, sequence_len, sampled_size].
• 一个张量indices有形状[batch_size, sequence_len, sampled_size].

然后你就可以：

import tensorflow as tf

# Create updates and indices...

# Create additional indices
i1, i2 = tf.meshgrid(tf.range(batch_size),
                     tf.range(sequence_len), indexing="ij")
i1 = tf.tile(i1[:, :, tf.newaxis], [1, 1, sampled_size])
i2 = tf.tile(i2[:, :, tf.newaxis], [1, 1, sampled_size])
# Create final indices
idx = tf.stack([i1, i2, indices], axis=-1)
# Output shape
to_shape = [batch_size, sequence_len, vocab_size]
# Get scattered tensor
output = tf.scatter_nd(idx, updates, to_shape)

tf.scatter_nd https://www.tensorflow.org/api_docs/python/tf/scatter_nd需要一个indices张量，一个updates张量和一些形状。updates是原始张量，形状就是所需的输出形状，所以[batch_size, sequence_len, vocab_size]. Now, indices比较复杂。由于您的输出有 3 个维度（等级 3），因此对于中的每个元素updates您需要 3 个索引来确定每个元素将放置在输出中的位置。所以形状indices参数应该与updates具有大小为 3 的附加维度。在这种情况下，我们希望第一个维度相同，但我们仍然必须指定 3 个索引。所以我们使用tf.meshgrid https://www.tensorflow.org/api_docs/python/tf/meshgrid生成我们需要的索引，并将它们沿着第三个维度平铺（最后一个维度中每个元素向量的第一个和第二个索引）updates是一样的）。最后，我们将这些索引与之前创建的映射索引堆叠在一起，就得到了完整的 3 维索引。