在Tensorflow中通过索引张量切片张量

Question

我有两个以下的张量(注意,它们都 Tensorflow张量,这意味着他们仍然当时我构建了以下切片运算实际上象征之前,我推出tf.Session()):

• params:有形状(64,784,256)
• indices:有形状(64,784)

我想构造一个返回以下张量的op:

• output:有形状(64,784)在哪里

output[i,j] = params_tensor[i,j, indices[i,j] ]

Tensorflow中最有效的方法是什么？

ps:我试过tf.gather但是无法利用它来执行我上面描述的操作.

非常感谢.
-Bests

Answer 1

你可以准确地得到你想要的东西tf.gather_nd.最后的表达是:

tf.gather_nd(params, tf.stack([tf.tile(tf.expand_dims(tf.range(tf.shape(indices)[0]), 1), [1, tf.shape(indices)[1]]), tf.transpose(tf.tile(tf.expand_dims(tf.range(tf.shape(indices)[1]), 1), [1, tf.shape(indices)[0]])), indices], 2))

该表达式具有以下说明:

• tf.gather_nd 做你所期望的并使用指数来收集参数的输出

• tf.stack结合了三个独立的张量,最后一个是指数.前两个张量指定前两个维度的顺序(params/indices的轴0和轴1)

  • 对于所提供的示例,对于轴0,该排序仅为0,1,2,...,63,对于轴1,该排序仅为0,1,2,... 783.这些序列分别用tf.range(tf.shape(indices)[0])和获得tf.range(tf.shape(indices)[1]).

  • 对于提供的示例,索引具有形状(64,784).从上一点开始的另外两个张量需要具有相同的形状才能与之结合tf.stack

    • 首先,使用两个序列中的每一个添加额外的维度/轴tf.expand_dims.

    • 使用tf.tile和tf.transpose可以通过示例显示:假设参数和索引的前两个轴具有形状(5,3).我们希望第一个张量为:

      [[0, 0, 0], [1, 1, 1], [2, 2, 2], [3, 3, 3], [4, 4, 4]]
      

      我们希望第二个张量为:

      [[0, 1, 2], [0, 1, 2], [0, 1, 2], [0, 1, 2], [0, 1, 2]]
      

      这两个张量几乎就像为相关索引指定网格中的坐标一样.

• 最后一部分tf.stack将三个张量组合在一个新的第三轴上,因此结果与params具有相同的3个轴.

请记住,如果您的轴数多于或少于问题,则需要相应地修改坐标指定张量的数量tf.stack.