如何使用张量流中的给定索引对张量进行切片？答案

【问题标题】：How to slice a tensor using given indices in tensorflow?如何使用张量流中的给定索引对张量进行切片？
【发布时间】：2019-04-13 19:50:01
【问题描述】：

我有一个带有概率的张量。这是一个形状为 (?, 30) 的动态张量，我选择这 30 个值中概率最大的索引为：

    best_probability = tf.argmax(probability, axis = 1)

现在张量 best_probability 的维度是 (?,)。现在我想从另一个称为数据的张量中选择具有这些索引的值，该张量具有维度 (?, 30, 1024, 3)。基本上使用 best_probability 张量从 30 个值中选择一个具有最佳概率的值。

最终输出的尺寸应为 (?, 1024, 3)。

PS：- 我尝试了gather_nd，但它需要索引 best_probability 张量，例如 [[0, 9], [1, 10], [2, 15], [3, 25]]。为此，我编写了以下 sn-p。

 selected_data = tf.stack(tf.range(probability.shape[0]),
                                tf.argmax(probability, axis = 1))

这不起作用，因为我正在处理动态张量。有没有办法解决这个问题。

【问题讨论】：

【解决方案1】：

我能够使用 tf.batch_gather 和 tf.reshape 解决这个问题

selected_data = tf.reshape(tf.batch_gather(data, best_probability),
                           (-1, data.shape[2],data.shape[3]))

【讨论】：