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

Question

给定indicesshape [batch_size, sequence_len]，updatesshape [batch_size, sequence_len, sampled_size]，to_shapeshape [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_shape。请注意，sequence_len和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正确使用吗？提前致谢！

Answer 1

所以假设你有：

• 一个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需要一个indices张量，一个updates张量和一些形状。updates是原始张量，而形状正是所需的输出形状，所以[batch_size, sequence_len, vocab_size]. 现在，indices情况更加复杂。由于您的输出有 3 个维度（等级 3），对于您中的每个元素，updates您需要 3 个索引来确定每个元素将放置在输出中的位置。因此indices参数的形状应该updates与大小为 3 的附加维度相同。在这种情况下，我们希望第一个维度相同，但我们仍然必须指定 3 个索引。所以我们tf.meshgrid用来生成我们需要的索引，并沿着第三维（最后一维中每个元素向量的第一个和第二个索引）将它们平铺updates是一样的）。最后，我们将这些索引与之前创建的映射索引叠加在一起，我们就有了完整的 3 维索引。