Keras中三维响应下均方误差失效的问题排查

问题场景回顾 #

你的Keras序列模型代码如下：

layer_width= 32

model= keras.Sequential()
# keras batch_size is the first dimension of data, while we are using the second dimension to create the mini-batches
model.add(keras.Input(shape= (batch_size, n_feat), batch_size= 72))
model.add(layers.Masking(mask_value= -1))
model.add(layers.LSTM(layer_width, activation= 'relu', recurrent_dropout= .5,))
model.add(layers.Dense(n_y))
model.compile(optimizer='adam', loss='mean_squared_error')

当n_y=1、传入shape为(72,256,1)的目标数据时训练正常，但n_y=6、传入shape为(72,256,6)的目标数据时，触发错误：

InvalidArgumentError: Graph execution error:
Can not squeeze dim[2], expected a dimension of 1, got 6

核心错误原因 #

这个错误的本质是模型输出形状与目标数据形状不匹配，具体逻辑如下：

1. LSTM层默认参数return_sequences=False，这意味着它只会返回输入序列最后一个时间步的输出，而非每个时间步的输出。结合你的模型，LSTM层的输出形状为(72, 32)（72是batch_size，32是layer_width），经过Dense(6)后，最终模型输出形状为(72, 6)。
2. 但你传入的目标数据y_train形状是(72,256,6)，包含了256个时间步的标签。此时损失函数mean_squared_error在计算时，会尝试对齐模型输出与目标的维度，由于两者维度数量不一致（模型输出是2维，目标是3维），内部会尝试挤压目标的多余维度；但目标的第三维度是6而非1，无法被挤压，因此触发维度不匹配的错误。
3. 当n_y=1时，目标数据形状是(72,256,1)，损失函数可以自动挤压第三维度（维度大小为1），将目标转为(72,256)，再通过广播机制和模型输出(72,1)进行计算，因此不会报错。

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