Tensorflow - 根据批次位置进行索引

Question

我正在处理屏蔽 r-cnn，但在根据标签索引屏蔽时遇到问题。

这是我想要实现的目标：我有一个张量(?,28,28,c)，其中? 是未知的batch_size，“28x28”是二维坐标，c 代表不同的标签，然后我有一个索引列表（基本上我的标签预测）（？，）int32。现在我想根据批次索引提取给定标签的掩码 -> 使其成为(?,28,28,1) 张量。

我试过self.masks_sigmoids = tf.gather(self.final_conv, self.label_predictions, axis=3)，但形状还是一样。

我也看了tf.gather_nd这里http://www.riptutorial.com/tensorflow/example/29069/how-to-use-tf-gather-nd，我想这是正确的路径，但我不知道如何根据批处理索引合并我想要的索引（在numpy中(b_i,:,:,c_i)）

我也觉得我的问题有点类似于Batched 4D tensor Tensorflow indexing，尽管我的问题似乎不那么复杂。然而，就 tensorflow 的快速发展而言，这个问题已经过时了，所以我要求一个可能更好、更清晰的解决方案。编辑：即使是肮脏的解决方案也可能有益，因为我没有在链接的 SO 中得到问题（已经写了评论要求澄清问题），因此我从唯一的答案中没有得到太多。这也可能对社区有益，因为这个问题更简单，这意味着它将更清楚地展示解决方案。

Answer 1

解决方案 1：更通用

你可以看看答案here，和你的问题基本一样，只是维度不同。

那里描述的解决方案是创建一个[?, 28, 28, 4]形张量indices where indices[i, x, y, :] = [i, x, y, self.label_predictions[i]]，然后使用tf.gather_nd：

self.masks_sigmoids = tf.gather_nd(self.final_conv, indices=indices)

构建indices 不是很优雅，如this answer 所示（为您增​​加了一个维度），但本身很容易。

解决方案 2：更优雅并适应您的问题

此解决方案与第一个解决方案非常相似，但避免创建indices 的[x, y] 部分。这个想法是使用gather_nd 的切片功能，通过在收集数据之前转置数据来避免在indices 中为每个(i, x, y) 写入[x, y]。我把整个代码放在这里，包括如何创建indices以及如何测试：

import numpy as np
import tensorflow as tf

N_CHANNELS = 5
pl=tf.placeholder(dtype=tf.int32, shape=(None, 28, 28, N_CHANNELS))

# Indices we'll use. batch_size = 4 here.
label_predictions = tf.constant([0, 2, 0, 3])

# Indices of shape [?, 2], with indices[i] = [i, self.label_predictions[i]],
# which is easy to do with tf.range() and tf.stack()
indices = tf.stack([tf.range(tf.size(label_predictions)), label_predictions], axis=-1)
# [[0, 0], [1, 2], [2, 0], [3, 3]]

transposed = tf.transpose(pl, perm=[0, 3, 1, 2])
gathered = tf.gather_nd(transposed, indices)  # Should be of shape (4, 2, 3)
result = tf.expand_dims(gathered, -1)

initial_value = np.arange(4*28*28*N_CHANNELS).reshape((4, 28, 28, N_CHANNELS))
sess = tf.InteractiveSession()
res = sess.run(result, feed_dict={pl: initial_value})
# print(res)

print("checking validity")
for i in range(4):
    for x in range(28):
        print(x)
        for y in range(28):
            assert res[i, x, y, 0] == initial_value[i, x, y, indices[i, 1].eval()]
print("All assertions passed")