如何批量使用 tf.gather？

Question

我有一个 A = 10x1000 张量和一个 B = 10x1000 索引张量。张量 B 的值介于 0-999 之间，它用于从 A 收集值（B[0,:] 从 A[0,:] 收集，B[1,:] 从 A[1,:] 收集，等等...）。

但是，如果我使用 tf.gather(A, B)，当我期待一个 10x1000 张量返回时，我会得到一个形状为 (10, 1000, 1000) 的数组。有什么想法可以解决这个问题吗？

编辑

假设A= [[1, 2, 3],[4,5,6]] 和B = [[0, 1, 1],[2,1,0]] 我想要的是能够使用相应的B 对A 进行采样。这应该导致C = [[1, 2, 2],[6,5,4]]。

Answer 1

1. 张量的维度是预先知道的。

首先，我们沿第一个维度“取消堆叠”参数和索引（分别为A 和B）。然后我们应用tf.gather() 使得A 的行对应B 的行。最后，我们将结果叠加在一起。

import tensorflow as tf
import numpy as np

def custom_gather(a, b):
    unstacked_a = tf.unstack(a, axis=0)
    unstacked_b = tf.unstack(b, axis=0)
    gathered = [tf.gather(x, y) for x, y in zip(unstacked_a, unstacked_b)]
    return tf.stack(gathered, axis=0)

a = tf.convert_to_tensor(np.array([[1, 2, 3], [4, 5, 6]]), tf.float32)
b = tf.convert_to_tensor(np.array([[0, 1, 1], [2, 1, 0]]), dtype=tf.int32)

gathered = custom_gather(a, b)
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print(sess.run(gathered))
# [[1. 2. 2.]
#  [6. 5. 4.]]

对于形状 1000x10 的初始情况，我们得到：

a = tf.convert_to_tensor(np.random.normal(size=(10, 1000)), tf.float32)
b = tf.convert_to_tensor(np.random.randint(low=0, high=999, size=(10, 1000)), dtype=tf.int32)
gathered = custom_gather(a, b)
print(gathered.get_shape().as_list()) # [10, 1000]

更新

1. 第一个维度未知（即None）

只有在预先知道第一个维度的情况下，前面的解决方案才有效。如果维度未知，我们按如下方式解决：

• 我们将两个张量堆叠在一起，这样两个张量的行就堆叠在一起：

# A = [[1, 2, 3], [4, 5, 6]]        [[[1 2 3]
#                            --->     [0 1 1]]
#                                    [[4 5 6]
# B = [[0, 1, 1], [2, 1, 0]]          [2 1 0]]]

• 我们迭代这个堆叠张量的元素（由堆叠在一起的A 和B 行组成）并使用tf.map_fn() 函数我们应用tf.gather()。

• 我们将使用tf.stack() 获得的元素堆叠起来

import tensorflow as tf
import numpy as np

def custom_gather_v2(a, b):
    def apply_gather(x):
        return tf.gather(x[0], tf.cast(x[1], tf.int32))
    a = tf.cast(a, dtype=tf.float32)
    b = tf.cast(b, dtype=tf.float32)
    stacked = tf.stack([a, b], axis=1)
    gathered = tf.map_fn(apply_gather, stacked)
    return tf.stack(gathered, axis=0)

a = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32)
b = np.array([[0, 1, 1], [2, 1, 0]], dtype=np.int32)

x = tf.placeholder(tf.float32, shape=(None, 3))
y = tf.placeholder(tf.int32, shape=(None, 3))

gathered = custom_gather_v2(x, y)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print(sess.run(gathered, feed_dict={x:a, y:b}))
# [[1. 2. 2.]
#  [6. 5. 4.]]

Answer 2

将tf.gather 与batch_dims=-1 一起使用：

import numpy as np
import tensorflow as tf

rois = np.array([[1, 2, 3],[3, 2, 1]])
ind = np.array([[0, 2, 1, 1, 2, 0, 0, 1, 1, 2],
        [0, 1, 2, 0, 2, 0, 1, 2, 2, 2]])
tf.gather(rois, ind, batch_dims=-1)
# output:
# <tf.Tensor: shape=(2, 10), dtype=int64, numpy=
# array([[1, 3, 2, 2, 3, 1, 1, 2, 2, 3],
#       [3, 2, 1, 3, 1, 3, 2, 1, 1, 1]])>