将布尔掩码应用于TensorFlow中张量的最后两个维度

Question

我正在将一堆 Numpy 计算移植到 TensorFlow。在我计算的一个阶段，我使用布尔掩码从一个大数组中提取和展平一个值子集。数组可以有多个维度，但布尔掩码仅作用于最后两个维度。在 Numpy 中，它看起来像这样：

mask = np.array([
    [False, True , True , True ],
    [True , False, True , True ],
    [True , True , False, False],
    [True , True , False, False]]

large_array_masked = large_array[..., mask]

我不知道如何在 TensorFlow 中做同样的事情。我试过了：

tf.boolean_mask(large_array, mask, axis = -2)

这不起作用，因为tf.boolean_mask() 似乎没有采用负轴参数。

作为一个丑陋的黑客，我尝试使用以下命令强制 mask 广播到与 large_array 相同的形状：

mask_broadcast = tf.logical_and(tf.fill(tf.shape(large_array), True), mask)
large_array_masked = tf.boolean_mask(large_array, mask_broadcast)

看来mask_broadcast 具有我想要的形状和值，但我得到了错误：

ValueError: Number of mask dimensions must be specified, even if some dimensions are None

这可能是因为large_array 是根据输入计算的，因此它的形状不是静态的。

有什么建议吗？

Answer 1

我想出了一个技巧来解决我的狭隘问题，所以我在这里发帖，但我接受了@Sorin 的答案，因为它可能更普遍适用。

为了解决tf.boolean_mask() 只能作用于初始索引的事实，我只是将索引向前滚动，应用遮罩，然后将它们回滚。在简化的形式中，它看起来像这样：

indices = tf.range(tf.rank(large_array))
large_array_rolled_forward = tf.transpose(
    large_array,
    tf.concat([indices[-2:], indices[:-2]], axis=0))
large_array_rolled_forward_masked = tf.boolean_mask(
    large_array_rolled_forward,
    mask)
new_indices = tf.range(tf.rank(large_array_rolled_forward_masked))
large_array_masked = tf.transpose(
    large_array_rolled_forward_masked,
    tf.concat([new_indices[1:], [0]], axis=0))

Answer 2

一般来说，我发现在 tensorflow 中您需要众所周知的形状。这是因为大多数运算都是矩阵乘法，并且矩阵是固定形状的。

如果你真的想这样做，你需要转换为稀疏张量，然后应用tf.sparse_retain。

我通常在 tensorflow 中使用的“等效”是将掩码与 large_array 相乘，使 False 值 (large_array_masked = large_array * mask) 变为 0。这样可以保持原始形状，以便更容易传递到密集层等......