Keras掩码报错：无法压缩维度[1]，自定义损失适配问题

1. 报错原因与损失函数修改 #

报错原因 #

当Embedding层开启mask_zero=True时，Keras会自动生成掩码并传递给后续GRU层，用于跳过填充的0值时间步。但你的自定义损失函数直接取y_pred[:, -1]（序列最后一个时间步的预测），此时Keras损失计算框架会尝试结合掩码对损失做加权处理，而你的损失输出维度（仅最后一步损失）与框架预期的全序列损失维度不匹配，最终触发维度挤压错误。

修改后的损失函数 #

需要结合掩码定位每个样本真实的最后一个有效时间步（而非直接取序列末尾，因为填充序列的末尾可能是无效的0值），修改后的损失函数如下：

def BCE_Last_Event(y_true, y_pred):
    # 获取传递到输出的掩码
    mask = tf.keras.backend.in_train_phase(
        y_pred._keras_mask, None
    )
    if mask is None:
        # 无掩码时直接取最后一步
        y_last_pred = tf.expand_dims(y_pred[:, -1], -1)
    else:
        # 计算每个样本最后一个有效时间步的索引
        mask_float = tf.cast(mask, tf.float32)
        last_indices = tf.math.reduce_sum(mask_float, axis=1) - 1
        last_indices = tf.cast(last_indices, tf.int32)
        # 生成用于提取的索引矩阵
        batch_size = tf.shape(y_pred)[0]
        batch_indices = tf.range(batch_size)
        gather_indices = tf.stack([batch_indices, last_indices], axis=1)
        # 提取最后有效步的预测值
        y_last_pred = tf.expand_dims(tf.gather_nd(y_pred, gather_indices), -1)
    
    y_last_true = y_true
    return tf.keras.metrics.binary_crossentropy(y_last_true, y_last_pred, axis=-1)

注：GRU层会自动接收Embedding层的掩码，后续输出y_pred会携带_keras_mask属性，确保掩码能传递到损失函数。

2. 含-999填充的数值向量与Embedding输出拼接后的掩码处理 #

步骤1：构建统一掩码 #

对-999填充的数值向量手动构建掩码，再与Embedding层的掩码合并，确保GRU只处理两个特征都有效的时间步：

# 数值输入定义
num_input = Input(shape=(None,), name="num_feat")
# 构建数值特征掩码：非-999的位置为有效
num_mask = tf.not_equal(num_input, -999)
# 获取Embedding层生成的掩码
emb_mask = emb._keras_mask
# 合并掩码：两个特征都有效时才视为有效时间步
combined_mask = tf.logical_and(emb_mask, num_mask)

# 扩展数值特征维度，匹配Embedding输出维度
num_emb = tf.expand_dims(num_input, -1)
# 拼接Embedding输出与数值特征
concat_input = tf.concat([emb, num_emb], axis=-1)

# 将统一掩码传递给GRU层
gru = GRU(32,
          activation="tanh",
          dropout=0,
          recurrent_dropout=0,
          go_backwards=False,
          return_sequences=True,
          name="gru_cat")(concat_input, mask=combined_mask)

步骤2：损失函数适配 #

将问题1修改后的损失函数中的mask替换为上述combined_mask即可，逻辑一致：通过统一掩码找到每个样本的最后有效时间步，提取对应预测值计算损失。

3. 推理时掩码的行为 #

掩码在推理时会生效，但不会输出空值或重复前一时间步分数：

• GRU层在推理时会利用掩码跳过填充的时间步，计算到最后一个有效时间步后，后续填充步的GRU状态会保持最后有效步的状态。
• 由于模型设置return_sequences=True，输出序列长度与输入一致，填充位置的输出是基于GRU最后有效状态计算的结果，并非空值或直接复制前一步输出。
• 若不需要填充位置的输出，可在推理后根据掩码过滤掉无效时间步的结果。

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