TensorFlow中使用sequence mask如何保留原张量空间信息（不降维）？

你说得对，tf.boolean_mask确实会把符合掩码条件的元素提取出来并扁平化，直接丢失了原张量的空间结构。要保留原形状的话，我们可以通过元素级的掩码操作来实现，而不是提取元素。这里给你两种实用的方法：

方法1：用tf.where填充掩码外的位置 #

tf.where可以根据掩码条件，选择保留原张量的值，或者用指定值填充掩码外的区域，完美保留原形状：

import numpy as np
import tensorflow as tf
import tensorflow.contrib.eager as tfe
tfe.enable_eager_execution()

x = np.random.randint(100, size=(4,4))
indexes = tf.sequence_mask([1,2,2,4], 4)

# 把掩码为False的位置填充为0（你也可以换成NaN或者其他需要的填充值）
y = tf.where(indexes, x, tf.zeros_like(x))
print(y.numpy())

比如如果你的x是：

[[43, 12, 34, 56],
 [78, 68, 23, 90],
 [54, 46, 77, 88],
 [28, 15, 52, 3]]

处理后的y会是：

[[43, 0, 0, 0],
 [78, 68, 0, 0],
 [54, 46, 0, 0],
 [28, 15, 52, 3]]

完全保留了4x4的空间结构。

方法2：掩码转类型后与原张量相乘 #

如果你只想把掩码外的位置置为0，也可以把布尔掩码转换成和原张量同类型的整数张量，然后做元素级乘法：

# 将布尔掩码转为与x同类型的张量（True→1，False→0）
mask_int = tf.cast(indexes, x.dtype)
# 元素级相乘，掩码外的位置自动变为0
y = x * mask_int
print(y.numpy())

这个方法和上面的效果一致，代码更简洁，适合只需要置0的场景。

为什么不用tf.boolean_mask？ #

tf.boolean_mask的设计目的就是提取符合条件的元素并扁平化，所以它必然会改变张量的形状。如果要保留原空间信息，就需要用上述的元素级操作，直接在原张量上修改掩码区域的值。

内容的提问来源于stack exchange，提问作者burglarhobbit