如何使用tf.gather类似于numpy的切片

Question

我有一个索引列表，代表我想要访问的行和列。

r= np.array([0, 1, 0, 1, 0, 1, 2, 3, 2, 3, 2, 3])
c = np.array([0, 0, 1, 1, 4, 4, 1, 1, 3, 3, 5, 5])

当我将它们用作 [r,c] 时，我能够在 numpy 数组中获取相应的元素。

arr = np.array([[517.0, 1876.0, 4716.0, 2725.0, 2138.0, 2213.0],
                [517.0, 1876.0, 4716.0, 2725.0, 2138.0, 2213.0],
                [3745.0, 3103.0, 885.0, 3482.0, 4196.0, 1802.0], 
                [3745.0, 3103.0, 885.0, 3482.0, 4196.0, 1802.0],
                [1548.0, 610.0, 3936.0, 905.0, 3791.0, 3657.0], 
                [1548.0, 610.0, 3936.0, 905.0, 3791.0, 3657.0], 
                [971.0, 573.0, 4756.0, 2137.0, 1407.0, 4388.0], 
                [971.0, 573.0, 4756.0, 2137.0, 1407.0, 4388.0], 
                [2786.0, 4769.0, 3391.0, 940.0, 2188.0, 1823.0], 
                [2786.0, 4769.0, 3391.0, 940.0, 2188.0, 1823.0], 
                [3225.0, 3262.0, 3444.0, 783.0, 3931.0, 1546.0], 
                [3225.0, 3262.0, 3444.0, 783.0, 3931.0, 1546.0]])
print(mst_array[r, c])
[ 517.  517. 1876. 1876. 2138. 2138. 3103. 3103. 3482. 3482. 1802. 1802.]

我想在 tf 中执行相同的操作，我正在使用 tf.gather。

https://www.tensorflow.org/api_docs/python/tf/gather

t_r = tf.convert_to_tensor(r, dtype= tf.int32)
t_c =  tf.convert_to_tensor(c, dtype= tf.int32)
t_arr = tf.convert_to_tensor(arr, dtype= tf.int32)
print(t_r, "\n", t_c, "\n")
tf.Tensor([0 1 0 1 0 1 2 3 2 3 2 3], shape=(12,), dtype=int32) 
tf.Tensor([0 0 1 1 4 4 1 1 3 3 5 5], shape=(12,), dtype=int32)

我尝试了不同的命令，但无法获得与我在 numpy 中相同的结果。

tf.gather(t_arr, [t_r, t_c], axis=-1)
tf.gather(t_arr, [t_r, t_c], axis=-1).shape_as_list()

Answer 1

尝试 tf.stack 和 tf.gather_nd:

z = tf.gather_nd(t_arr, tf.stack([t_r, t_c], axis=1))

tf.Tensor([ 517  517 1876 1876 2138 2138 3103 3103 3482 3482 1802 1802], shape=(12,), dtype=int32)

如何使用tf.gather类似于numpy的切片

How to use tf.gather similar to the numpy slicing

python

numpy

tensorflow

numpy-slicing

tensorflow2.0