从张量流中的张量中提取特定元素答案

【问题标题】：Extracting specific elements from a tensor in tensorflow从张量流中的张量中提取特定元素
【发布时间】：2017-01-23 10:15:13
【问题描述】：

我在 python 上使用 tensorflow 我有一个形状为 [?, 5, 37] 的数据张量和一个形状为 [?, 5] 的 idx 张量

我想从数据中提取元素并获得形状为 [?, 5] 的输出，这样：

output[i][j] = data[i][j][idx[i, j]] for all i in range(?) and j in range(5)

看起来 tf.gather_nd() 函数最符合我的需要，但我不知道如何使用它...

谢谢！

编辑：我设法用gather_nd 做到了，如下所示，但有更好的选择吗？（好像有点强硬）

    nRows = tf.shape(length_label)[0] ==> ?
    nCols = tf.constant(MAX_LENGTH_INPUT + 1, dtype=tf.int32) ==> 5
    m1 = tf.reshape(tf.tile(tf.range(nCols), [nRows]),
                                           shape=[nRows, nCols])
    m2 = tf.transpose(tf.reshape(tf.tile(tf.range(nRows), [nCols]),
                                            shape=[nCols, nRows]))
    indices = tf.pack([m2, m1, idx], axis=-1)
    # indices should be of shape [?, 5, 3] with indices[i,j]==[i,j,idx[i,j]]
    output = tf.gather_nd(data, indices=indices)

【问题讨论】：

您的解决方案在我看来不错。

标签： python tensorflow

【解决方案1】：

我设法用gather_nd 做到了，如下所示

nRows = tf.shape(length_label)[0] # ==> ?
nCols = tf.constant(MAX_LENGTH_INPUT + 1, dtype=tf.int32) # ==> 5
m1 = tf.reshape(tf.tile(tf.range(nCols), [nRows]),
                                       shape=[nRows, nCols])
m2 = tf.transpose(tf.reshape(tf.tile(tf.range(nRows), [nCols]),
                                        shape=[nCols, nRows]))
indices = tf.pack([m2, m1, idx], axis=-1)
# indices should be of shape [?, 5, 3] with indices[i,j]==[i,j,idx[i,j]]
output = tf.gather_nd(data, indices=indices)

【讨论】：