使用 tensorflow 通过索引张量从值张量中获取值答案

【问题标题】：use tensorflow to get value from a value tensor by index tensor使用 tensorflow 通过索引张量从值张量中获取值
【发布时间】：2020-09-26 09:25:07
【问题描述】：

有一个索引张量是这样的：[[1,2,3],[1,2,3]] (shape is batch * length)

并且有一个像这样的值张量：（形状是批次*长度*深度）

[[[0.9,0.9,0.1,0.1],[0.9,0.1,0.8,0.1],[0.9,0.1,0.1,0.6]],
[[0.1,0.9,0.8,1],[1,2,0.8,0.1],[0.1,0.1,2,0.6]]].

如何使用 tensorflow 获得 [[0.9,0.8,0.6],[0.9,0.8,0.6]]？

【问题讨论】：

tf.math.reduce_max(tensor, axis=-1)?
我遇到的问题是获取一行值并将其用作另一个矩阵的最后一维中的索引，对应的值。这不是关于最大值

标签： tensorflow

【解决方案1】：

我不确定这是不是最好的解决方案，但它确实有效：tf.gather_nd(values, tf.expand_dims(index, -1), batch_dims=2)

例如：

>>> index = tf.constant([[1,2,3],[1,2,3]])
>>> values = tf.constant([[[0.9,0.9,0.1,0.1],[0.9,0.1,0.8,0.1],[0.9,0.1,0.1,0.6]],[[0.1,0.9,0.8,1],[1,2,0.8,0.1],[0.1,0.1,2,0.6]]])
>>> result = tf.gather_nd(values, tf.expand_dims(index, -1), batch_dims=2)
>>> result.eval()
array([[0.9, 0.8, 0.6],
       [0.9, 0.8, 0.6]], dtype=float32)

【讨论】：