介绍
根据索引从矩阵取出元素
这个函数有些小复杂
语法
tf.gather_nd(
params,
indices,
batch_dims=0,
name=None
)
参数：
- params：Tensor矩阵
- indices：索引。int32 或 int64。
- batch_dims：整数，批次尺寸的数量。
例子
一维索引
arr=np.arange(18).reshape((3,2,3))
print(arr)
print('-'*50)
res=tf.gather_nd(arr,[1])
print(res)
执行结果：
[[[ 0 1 2]
[ 3 4 5]]
[[ 6 7 8]
[ 9 10 11]]
[[12 13 14]
[15 16 17]]]
--------------------------------------------------
tf.Tensor(
[[ 6 7 8]
[ 9 10 11]], shape=(2, 3), dtype=int64)
返回0轴索引是1的数组
m行1列索引
arr=np.arange(18).reshape((3,2,3))
print(arr)
print('-'*50)
res=tf.gather_nd(arr,[[2],[1]])
print(res)
执行结果：
[[[ 0 1 2]
[ 3 4 5]]
[[ 6 7 8]
[ 9 10 11]]
[[12 13 14]
[15 16 17]]]
--------------------------------------------------
tf.Tensor(
[[[12 13 14]
[15 16 17]]
[[ 6 7 8]
[ 9 10 11]]], shape=(2, 2, 3), dtype=int64)
将0轴索引是2和1的数组，组成新数组并返回
m行n列索引
arr=np.arange(18).reshape((3,2,3))
print(arr)
res=tf.gather_nd(arr,[[2,1],[1,0]])
print(res)
执行结果：
[[[ 0 1 2]
[ 3 4 5]]
[[ 6 7 8]
[ 9 10 11]]
[[12 13 14]
[15 16 17]]]
--------------------------------------------------
tf.Tensor(
[[15 16 17]
[ 6 7 8]], shape=(2, 3), dtype=int64)
执行过程：从矩阵中找出索引是(2,1)、(1,0)的元素，并组成新数组，然后返回
相当于numpy中的：arr[[2,1],[1,0]]
3维数组
arr=np.arange(18).reshape((3,2,3))
print(arr)
print('-'*50)
res=tf.gather_nd(arr,[[2,1,2],[1,0,1]])
print(res)
执行结果：
[[[ 0 1 2]
[ 3 4 5]]
[[ 6 7 8]
[ 9 10 11]]
[[12 13 14]
[15 16 17]]]
--------------------------------------------------
tf.Tensor([17 7], shape=(2,), dtype=int64)
从矩阵中找出索引是(2,1,2)和(1,0,1)的数组，并组成新数组，然后返回