# tf.gather tf.gather_nd 和 tf.batch_gather 使用方法

在计算机视觉的项目中，会遇到根据索引获取数组特定下标元素的任务。常用的函数有tf.gather , tf.gather_nd 和 tf.batch_gather。

1.tf.gather(params,indices,validate_indices=None,name=None,axis=0)

indices：一维索引张量

name：返回张量名称

返回值：通过indices获取params下标的张量。

``````import tensorflow as tf
tensor_a = tf.Variable([[1,2,3],[4,5,6],[7,8,9]])
tensor_b = tf.Variable([1,2,0],dtype=tf.int32)
tensor_c = tf.Variable([0,0],dtype=tf.int32)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
print(sess.run(tf.gather(tensor_a,tensor_b)))
print(sess.run(tf.gather(tensor_a,tensor_c)))``````

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

2.tf.gather_nd(params,indices,name=None)

``````import tensorflow as tf
tensor_a = tf.Variable([[1,2,3],[4,5,6],[7,8,9]])
tensor_b = tf.Variable([[1,0],[1,1],[1,2]],dtype=tf.int32)
tensor_c = tf.Variable([[0,2],[2,0]],dtype=tf.int32)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
print(sess.run(tf.gather_nd(tensor_a,tensor_b)))
print(sess.run(tf.gather_nd(tensor_a,tensor_c)))``````

tf.gather_nd(tensor_a,tensor_b)值为[4,5,6],tf.gather_nd(tensor_a,tensor_c)的值为[3,7].

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

3.tf.batch_gather(params,indices,name=None)

``````import tensorflow as tf
tensor_a = tf.Variable([[1,2,3],[4,5,6],[7,8,9]])
tensor_b = tf.Variable([[0],[1],[2]],dtype=tf.int32)
tensor_c = tf.Variable([[0],[0],[0]],dtype=tf.int32)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
print(sess.run(tf.batch_gather(tensor_a,tensor_b)))
print(sess.run(tf.batch_gather(tensor_a,tensor_c)))``````

tf.gather_nd(tensor_a,tensor_b)值为[1,5,9],tf.gather_nd(tensor_a,tensor_c)的值为[1,4,7].

tensor_a的三个元素[1,2,3],[4,5,6],[7,8,9]分别对应索引元素的第一，第二和第三个值。[1,2,3]的第0个元素为1,[4,5,6]的第1个元素为5,[7,8,9]的第2个元素为9,所以索引[[0],[1],[2]]的返回值为[1,5,9],同样地，索引[[0],[0],[0]]的返回值为[1,4,7].

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

在深度学习的模型训练中，有时候需要对一个batch的数据进行类似于tf.gather_nd的操作，但tensorflow中并没有tf.batch_gather_nd之类的操作，此时需要tf.map_fn和tf.gather_nd结合来实现上述操作。