tf.gather_nd函数

函数原型

tf.gather_nd(
    params, indices, batch_dims=0, name=None
)

函数说明

从指定的张量中根据索引来进行切片。注意tf.gather函数和tf.gather_nd函数的区别。

参数params表示待切片的张量。

参数indices表示进行切片的索引。

参数batch_dims表示批处理的维度，这个参数比较难理解，可以看一下使用示例中的用法。

使用示例

1、一维张量的索引

>>> params = tf.constant([1, 3, 5, 7, 9])
# 对于一位张量，只能进行索引，输入某一个索引对应的值。
>>> indices = tf.constant([0])
>>> tf.gather_nd(params, indices)
<tf.Tensor: shape=(), dtype=int32, numpy=1>

2、二维张量的索引

>>> params = tf.constant([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
>>> params
<tf.Tensor: shape=(3, 3), dtype=int32, numpy=
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])>
>>> indices = tf.constant([[0, 0], [1, 1]])
>>> tf.gather_nd(params, indices)
<tf.Tensor: shape=(2,), dtype=int32, numpy=array([1, 5])>

>>> params = tf.constant([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
>>> params
<tf.Tensor: shape=(3, 3), dtype=int32, numpy=
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])>
>>> indices = tf.constant([[0], [1]])
# 注意使用batch_dims和不使用batch_dims，参数indices的区别
>>> tf.gather_nd(params, indices, batch_dims=1)
<tf.Tensor: shape=(3,), dtype=int32, numpy=array([1, 5,])>

3、二维张量的切片

>>> params = tf.constant([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
>>> params
<tf.Tensor: shape=(3, 3), dtype=int32, numpy=
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])>
>>> indices = tf.constant([[0], [1]])
>>> tf.gather_nd(params, indices)
<tf.Tensor: shape=(2, 3), dtype=int32, numpy=
array([[1, 2, 3],
       [4, 5, 6]])>

4、三维张量的索引

# 没有使用batch_dims
>>> params = tf.constant([[[1, 1], [2, 2], [3, 3]],
		      [[4, 4], [5, 5], [6, 6]],
		      [[7, 7], [8, 8], [9, 9]]])
>>> params
<tf.Tensor: shape=(3, 3, 2), dtype=int32, numpy=
array([[[1, 1],
        [2, 2],
        [3, 3]],

       [[4, 4],
        [5, 5],
        [6, 6]],

       [[7, 7],
        [8, 8],
        [9, 9]]])>
>>> indices = tf.constant([[0, 0], [1, 1], [2, 2]])
>>> tf.gather_nd(params, indices)
<tf.Tensor: shape=(3, 2), dtype=int32, numpy=
array([[1, 1],
       [5, 5],
       [9, 9]])>

# 使用batch_dims
>>> params = tf.constant([[[1, 1], [2, 2], [3, 3]],
		      [[4, 4], [5, 5], [6, 6]],
		      [[7, 7], [8, 8], [9, 9]]])
>>> params
<tf.Tensor: shape=(3, 3, 2), dtype=int32, numpy=
array([[[1, 1],
        [2, 2],
        [3, 3]],

       [[4, 4],
        [5, 5],
        [6, 6]],

       [[7, 7],
        [8, 8],
        [9, 9]]])>
>>> indices = tf.constant([[0], [1], [2]])
>>> tf.gather_nd(params, indices, batch_dims=1)
<tf.Tensor: shape=(3, 2), dtype=int32, numpy=
array([[1, 1],
       [5, 5],
       [9, 9]])>

5、三维张量的切片

>>> params = tf.constant([[[1, 1], [2, 2], [3, 3]],
		      			  [[4, 4], [5, 5], [6, 6]],
		     			  [[7, 7], [8, 8], [9, 9]]])
>>> indices = tf.constant([[0], [1]])
>>> tf.gather_nd(params, indices)
<tf.Tensor: shape=(2, 3, 2), dtype=int32, numpy=
array([[[1, 1],
        [2, 2],
        [3, 3]],

       [[4, 4],
        [5, 5],
        [6, 6]]])>