TensorFlow中利用tf.where和tf.gather_nd提取满足条件的Tensor元素

TensorFlow中条件数值提取的核心思路

在TensorFlow中，若需提取满足特定条件的数值并组成新的Tensor，首先应使用tf.where()获取符合条件元素的索引。需要注意的是，这与NumPy不同，在NumPy中可直接通过索引引用提取数值，但在TensorFlow中该方式不可行。为此，TensorFlow提供了tf.gather()和tf.gather_nd()函数来实现这一功能。

具体实现代码如下：

import tensorflow as tf
sess = tf.Session()
def get_tensor():
  x = tf.random_uniform((5, 4))
  ind = tf.where(x>0.5)
  y = tf.gather_nd(x, ind)
  return x, ind, y

在上述代码示例中，依次输出了原始的Tensor x、满足特定条件（此处设定为大于0.5）的数值索引，以及由这些索引对应的数值组成的新Tensor。请执行以下步骤，并观察三个Tensor对应的具体数值：

x, ind, y = get_tensor()
x_, ind_, y_ = sess.run([x, ind, y])

运行后可得到如下结果：

从结果可以看出，上述过程成功将Tensor x中大于0.5的数值提取出来，并组合成了一个新的Tensor y。

若将代码中的tf.gather_nd替换为tf.gather，结果会有何不同？由于结果展示较为复杂，此处省略具体输出。需要注意的是，tf.gather适用于索引为一维的情况。在本例中，索引ind为二维，若错误选用tf.gather，则对应的x、ind和y的维度分别如下：

x.shape = (5, 4)
ind.shape = (9, 2)
y.shape = (9, 2, 4)

以上关于在TensorFlow中提取满足特定条件数值组成新Tensor的方法分享完毕，希望能对读者有所帮助，也欢迎大家继续支持本站。