如何将“tf.scatter_nd”与多维张量一起使用

Question

我正在尝试创建一个新的张量（output），其中另一个张量（）的值updates根据idx张量放置。的形状output应该是[batch_size, 1, 4, 4]（如 2x2 像素和一个通道的图像）并且update具有形状[batch_size, 3]。

我已阅读 Tensorflow 文档（我正在使用 GPU 版本 1.13.1）并发现tf.scatter_nd应该可以解决我的问题。问题是我无法让它发挥作用，我认为我在理解如何安排方面遇到了问题idx。

让我们考虑一下batch_size = 2，所以我正在做的是：

updates = tf.constant([[1, 2, 3], [4, 5, 6]])  # shape [2, 3]
output_shape = tf.constant([2, 1, 4, 4])
idx = tf.constant([[[1, 0], [1, 1], [1, 0]], [[0, 0], [0, 1], [0, 2]]])  # shape [2, 3, 2]
idx_expanded = tf.expand_dims(idx, 1)  # so I have shape [2, 1, 3, 2]
output = tf.scatter_nd(idx_expanded, updates, output_shape)

我希望它能工作，但事实并非如此，它给了我这个错误：

ValueError: The outer 3 dimensions of indices.shape=[2,1,3,2] must match the outer 3 dimensions of updates.shape=[2,3]: Shapes must be equal rank, but are 3 and 2 for 'ScatterNd_7' (op: 'ScatterNd') with input shapes: [2,1,3,2], [2,3], [4]

我不明白为什么它期望updates有维度3。我认为必须与（这就是我使用的原因）和（指定三个点的两个索引）idx有意义，但很明显我在这里遗漏了一些东西。output_shapeexpand_dimsupdates

任何帮助，将不胜感激。

Answer 1

我一直在研究这个函数，发现了我的错误。如果有人遇到这个问题，这就是我解决它的方法：

考虑到batch_size=2和3点，idx张量必须具有 shape [2, 3, 4]，其中第一个维度对应于我们从中获取update值的批次，第二个维度必须等于updates（每批次的点数）的第二个维度，第三个维度是4因为我们需要4索引： [批次号、通道、行、列]。按照问题中的示例：

updates = tf.constant([[1., 2., 3.], [4., 5., 6.]])  # [2, 3]
idx = tf.constant([[[0, 0, 0, 1], [0, 0, 0, 0], [0, 0, 1, 0]], [[1, 0, 1, 1], [1, 0, 0, 0], [1, 0, 1, 0]]])  # [2, 3, 4]
output = tf.scatter_nd(idx, updates, [2, 1, 4, 4])

sess = tf.Session()
print(sess.run(output))

[[[[2. 1. 0. 0.]
   [3. 0. 0. 0.]
   [0. 0. 0. 0.]
   [0. 0. 0. 0.]]]


 [[[5. 0. 0. 0.]
   [6. 4. 0. 0.]
   [0. 0. 0. 0.]
   [0. 0. 0. 0.]]]]

这样就可以将特定的数字放入新的张量中。