Boolean mask Tensorflow, tf.boolean_mask - 保持原始张量的尺寸

Question

我有一个相同长度的 1dtensor和 1d 数组Booleans。

我想使用布尔数组作为张量的掩码，这样 True 保持张量中的原始元素值，而 False 将张量中的原始元素值设置为零。

例如

Tensor = [1,2,3,4,5]
Array = [True, False, False, False, True]

将布尔掩码应用于张量：

Desired result = [1, 0, 0, 0, 5]
Result with tf.boolean_mask = [1, 5]

tf.boolean_mask(tensor, array)但是，我尝试使用，这会减少生成的张量的维度，使其仅包含True元素，在上面的示例中为 2 维。

如何在保持张量的原始尺寸的同时将布尔掩码应用于张量？

Answer 1

您可以将布尔掩码转换为tf.int32张量并使用tf.multiply：

mask = tf.constant([True, False, False, False, True])
A = tf.range(1,6)

with tf.Session() as sess:
   res = sess.run( \
          tf.multiply(A, tf.cast(mask, tf.int32)) \
         ) # [1, 0, 0, 0, 5]