python - 在 tensorflow 中使用 bool 张量选择张量的子集

Question

我有两个 2 阶张量 arr1 和 arr2，形状为 m 乘以 n。张量 arr2 是 bool 值；它的每一行中恰好有一个条目是 True。我想提取一个长度为 m 的新 rank-1 张量 arr3，其中 arr3 的第 i 个条目等于 arr1 的第 i 行的条目对应于 arr2 的第 i 行是等于 True。

在 numpy 中，我可以这样做:

arr1 = np.array([[1,2],
                 [3,4]])
arr2 = np.array([[0,1],
                 [1,0]], dtype="bool")
arr3 = arr1[arr2]

我可以在 tensorflow 中做类似的事情吗？我知道我可以 eval() 我的张量，然后使用 numpy 函数，但这似乎效率不高。

这question建议使用 tf.gather 和 tf.select，但它不像我的问题那样处理折叠输出的维度。

Answer 1

您可以使用 tf.boolean_mask 吗？

from __future__ import print_function

import tensorflow as tf

with tf.Session() as sess:
    arr1 = tf.constant([[1,2],
                        [3,4]])
    arr2 = tf.constant([[False, True],
                        [True, False]])
    print(sess.run(tf.boolean_mask(arr1, arr2)))

应该给出:[2, 3]