如何使用类似于麻木切片的tf.ather

问题描述

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

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, "
", t_c, "
")
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()

---

解决方案

尝试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)