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

给定形状为 indices 的 [batch_size, sequence_len] ，形状为 updates 的 [batch_size, sequence_len, sampled_size] ，形状为 to_shape 的 [batch_size, sequence_len, vocab_size] ，其中 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 吗？提前致谢!

最佳答案

所以假设你有:

形状为 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 维索引。

关于python - 在tensorflow-r1.2中正确使用 `tf.scatter_nd`，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/45162998/