TensorFlow内置Loss对象可用但自定义/内置Loss函数不收敛问题排查

核心现象 #

• 自定义new_loss_orig逻辑上是MSE损失，但模型训练后完全不收敛，预测结果接近随机值
• 切换到内置tf.keras.losses.MeanSquaredError()对象时，模型能正常收敛并输出有效预测
• 尝试使用tf.keras.losses.mean_squared_error函数时仍出现不收敛问题，但打印梯度显示所有梯度非空且非零

可能原因分析 #

1. 标签与预测结果的形状不匹配 #

这是最可能的核心问题：

• 模型最后一层是Dense(1)，输出y_pred的形状为(batch_size, 1)
• 如果输入标签y_true的形状是(batch_size,)（一维张量），执行y_true - y_pred时会触发TensorFlow的广播机制，生成形状为(batch_size, batch_size)的张量
• 此时tf.reduce_mean(tf.square(...))会对所有batch_size*batch_size个元素求平均，导致损失值被缩小到正常MSE的1/batch_size量级，梯度更新幅度被严重削弱，模型看起来完全不收敛
• 内置MeanSquaredError()会自动处理形状匹配（比如自动扩展标签维度），避免广播错误

2. 损失函数的归约逻辑差异 #

• tf.keras.losses.mean_squared_error函数返回逐样本的MSE值（形状(batch_size,)），而自定义new_loss_orig`直接返回整个批量的平均损失（标量）
• 虽然在自定义训练循环中用tf.keras.metrics.Mean记录损失时结果数值一致，但如果形状不匹配导致中间计算错误，会直接影响梯度的方向和尺度

解决方法 #

步骤1：统一标签与预测结果的形状 #

在构建数据集时，将标签扩展为二维张量，确保和模型输出形状一致：

# 替换原数据集构建代码
train_dataset = tf.data.Dataset.from_tensor_slices((X_train_arr, tf.expand_dims(Y_train_arr, axis=-1))).shuffle(X_train_arr.shape[0]).batch(batch_size)
val_dataset = tf.data.Dataset.from_tensor_slices((X_val_arr, tf.expand_dims(Y_val_arr, axis=-1))).shuffle(X_val_arr.shape[0]).batch(batch_size)

步骤2：修复自定义损失函数的形状兼容 #

修改new_loss_orig和new_loss_mixup，确保输入形状不匹配时能自动调整：

def new_loss_orig(self, y_true, y_pred):
    # 统一维度：如果y_true是一维，扩展为二维
    if tf.rank(y_true) == 1:
        y_true = tf.expand_dims(y_true, axis=-1)
    mse = tf.reduce_mean(tf.square(y_true - y_pred))
    return mse

def new_loss_mixup(self, y_true, y_pred, long_run_pred):
    w_1 = 0.3
    w_2 = 0.7
    # 统一所有输入的维度
    if tf.rank(y_true) == 1:
        y_true = tf.expand_dims(y_true, axis=-1)
    if tf.rank(long_run_pred) == 1:
        long_run_pred = tf.expand_dims(long_run_pred, axis=-1)
    mse = tf.reduce_mean(w_1*tf.square(y_true - y_pred) + w_2*tf.square(long_run_pred-y_pred))
    return mse

步骤3：验证损失计算的正确性 #

在训练循环中添加打印语句，确认自定义损失和内置损失的数值一致：

with tf.GradientTape() as tape:
    y_batch_pred_train = model(x_batch_train, training=True)
    if self.new_loss == 'orig':
        loss_value = self.new_loss_orig(y_batch_train, y_batch_pred_train)
        # 对比内置损失
        builtin_loss = tf.keras.losses.MeanSquaredError()(y_batch_train, y_batch_pred_train)
        print(f"Custom loss: {loss_value.numpy()}, Builtin loss: {builtin_loss.numpy()}")

如果两者数值接近，说明损失计算逻辑正确；如果差异极大，继续检查形状匹配问题。

额外注意事项 #

• 确保new_loss_wmse函数实现完整（当前代码中return无返回值，会导致训练报错）
• 自定义训练循环中，优化器的apply_gradients依赖正确的梯度与权重映射，确保model.trainable_weights包含所有需要更新的参数

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