TensorFlow中Masking Layer引发维度压缩ValueError问题求助

问题背景 #

在Python 3.8.15、TensorFlow 2.11环境下，运行序列到序列模型时，添加layers.Masking层后调用model.fit()会抛出ValueError，注释掉Masking层则可正常运行。错误核心信息为：

ValueError: Can not squeeze dim[1], expected a dimension of 1, got 10 for '{{node mean_squared_error/weighted_loss/Squeeze}} = SqueezeT=DT_FLOAT, squeeze_dims=[-1]' with input shapes: [32,10].

错误原因 #

• 模型输出形状：由于LSTM设置了return_sequences=True，加上最后一层Dense(1)，模型输出形状为(32, 10, 1)（样本数、时间步、输出维度）
• 标签形状：labels = np.mean(inputs, axis=2)生成的形状是(32, 10)，比模型输出少了最后一维
• 掩码逻辑触发：启用Masking层时，Keras损失函数会自动应用掩码处理逻辑，尝试对样本权重的最后一维进行挤压操作，但此时样本权重的形状是(32,10)，没有可挤压的维度（需要维度为1才能执行挤压），因此触发错误；不使用Masking时，损失函数不会触发该掩码处理逻辑，所以未报错。

解决方案 #

只需调整标签的形状，增加一个维度，使其与模型输出的维度完全匹配，具体有两种实现方式：

方式1：生成标签时直接扩展维度 #

修改标签生成代码，通过[..., np.newaxis]或tf.expand_dims添加最后一维：

# 方案1：用numpy扩展维度
labels = np.mean(inputs, axis=2)[..., np.newaxis]
# 方案2：用TensorFlow扩展维度
labels = tf.expand_dims(np.mean(inputs, axis=2), axis=-1)

方式2：在模型末尾添加Reshape层（可选） #

如果不想修改标签，也可以在Dense层后添加Reshape层，将模型输出调整为与标签一致的形状：

layers.Dense(1, activation='linear', name='dense2'),
layers.Reshape((timesteps,))

修改后的完整代码示例 #

import tensorflow as tf
import numpy as np
from tensorflow.keras import layers

samples, timesteps, features = 32, 10, 8
inputs = np.random.random([samples, timesteps, features]).astype(np.float32)
inputs[:, 3, :] = 0.
inputs[:, 5, :] = 0.
print(inputs.shape)

# 调整labels形状，增加最后一维
labels = np.mean(inputs, axis=2)[..., np.newaxis]
print(labels.shape)  # 现在形状为(32,10,1)，与模型输出匹配

model = tf.keras.models.Sequential([
    layers.Input(shape=(timesteps, features), name='input1'),
    layers.Masking(mask_value=0.0, name='mask1'),
    layers.Bidirectional(layers.LSTM(32, return_sequences=True, name='lstm1'), name='bilstm1'),
    layers.Dropout(0.2, name='dropout1'),
    layers.Dense(1, activation='linear', name='dense2')
])

output = model(inputs)
print(output.shape)  # 形状为(32,10,1)

model.compile(
    loss        = tf.keras.losses.mean_squared_error,
    optimizer   = tf.keras.optimizers.Adam(),
    metrics     = [
        tf.keras.metrics.mean_absolute_error,
    ],
    run_eagerly = False,
)

print(model.summary())

history = model.fit(
    x      = inputs,
    y      = labels,
    epochs = 30
)

内容的提问来源于stack exchange，提问作者Daniel von Eschwege