在 Tensorflow 中通过索引张量对张量进行切片

Question

我有两个以下张量（请注意，它们是 both Tensorflow 张量，这意味着在我启动以下切片操作之前，它们实际上仍然是象征性的tf.Session()):

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

我想构造一个返回以下张量的操作：

• output: 形状为 (64,784)，其中

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

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

ps：我尝试使用tf.gather，但无法利用它来执行我上面描述的操作。

非常感谢。
-最好的

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 组合了三个独立的张量，最后一个是索引。前两个张量指定前两个维度的顺序（参数/索引的轴 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的使用可以举例说明：假设params和index的前两个轴的形状为(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 的最后部分将三个张量组合在一个新的第三轴上，因此结果与参数具有相同的三个轴。

---

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

Answer 2

你想要的就像一个自定义的归约函数。如果您在indices 保留类似最大值的索引，那么我建议使用tf.reduce_max：

max_params = tf.reduce_max(params_tensor, reduction_indices=[2])

否则，这是获得所需内容的一种方法（张量对象不可分配，因此我们创建张量的二维列表并使用 tf.pack 打包它）：

import tensorflow as tf
import numpy as np
with tf.Graph().as_default():

    params_tensor = tf.pack(np.random.randint(1,256, [5,5,10]).astype(np.int32))

    indices = tf.pack(np.random.randint(1,10,[5,5]).astype(np.int32))

    output = [ [None for j in range(params_tensor.get_shape()[1])] for i in range(params_tensor.get_shape()[0])]  
    for i in range(params_tensor.get_shape()[0]):
        for j in range(params_tensor.get_shape()[1]):
            output[i][j] = params_tensor[i,j,indices[i,j]]
    output = tf.pack(output)

    with tf.Session() as sess:
        params_tensor,indices,output = sess.run([params_tensor,indices,output])

        print params_tensor
        print indices
        print output