tf.gather_nd 从一维索引数组中获取二维数组的值答案

【问题标题】：tf.gather_nd to get values of 2d array from 1d array of indicestf.gather_nd 从一维索引数组中获取二维数组的值
【发布时间】：2020-11-14 03:10:40
【问题描述】：

我有一个二维值数组和一个一维索引数组。我想使用tf.gather_nd 或其他一些 numpy 或 tensorflow 命令来获取 2D 数组的每个元素的每个索引的值，如下所示：

2D_array = np.arange(27).reshape((9, 3))
array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11],
       [12, 13, 14],
       [15, 16, 17],
       [18, 19, 20],
       [21, 22, 23],
       [24, 25, 26]])

1D_array = np.random.randint(low=0, high=3, size=9)
array([2, 2, 0, 2, 0, 1, 0, 1, 0])

output
[2, 5, 6, 11, 12, 16, 18, 22, 24]

我目前拥有的是这个烂摊子：

output = tf.gather_nd(a, list(zip(np.arange(len(b)), b)))
<tf.Tensor: shape=(9,), dtype=int32, numpy=array([ 2,  5,  6, 11, 12, 16, 18, 22, 24])>

非常感谢避免使用 zip 和 list 的更简洁的解决方案，因为我认为我目前没有完全正确地使用 tf.gather_nd。

【问题讨论】：

标签： python arrays numpy tensorflow

【解决方案1】：

你的代码中的1D_array 应该是一个二维数组

inds = np.random.randint(low=0, high=3, size=9)
inds = np.dstack((np.arange(len(inds)),inds))

有了这个二维索引数组，你可以直接应用tf.gather_nd

import tensorflow as tf 
import numpy as np

tmp = np.arange(27).reshape((9, 3))
inds = np.random.randint(low=0, high=3, size=9)
inds = np.dstack((np.arange(len(inds)),inds))

output = tf.gather_nd(tmp,inds)

with tf.Session() as sess:
    output_val = sess.run(output)

【讨论】：