使用自定义损失函数训练Keras时出现形状不兼容问题

在运行下述Keras代码时出现形状不兼容错误，使用环境为tensorflow==2.13.0，怀疑问题源于自定义损失函数或Keras底层。

import numpy as np
import pandas as pd
from tensorflow.keras.layers import Dense
from tensorflow.keras.models import Sequential

N = 1000
df = pd.DataFrame({
    'Feature1': np.random.normal(loc=0, scale=1, size=N),
    'Feature2': np.random.normal(loc=1, scale=2, size=N),
    'Label': np.random.choice([0, 1], size=N)
})

df_train = df.sample(frac = 0.80, random_state = 42)
df_test = df[~df.index.isin(df_train.index)]
print(f"df_train.shape = {df_train.shape}")
print(f"df_test.shape = {df_test.shape}")

X_train, y_train = df_train[['Feature1', 'Feature2']], df_train['Label']
X_test, y_test = df_test[['Feature1', 'Feature2']], df_test['Label']

def my_loss(data, y_pred):
    y_true = data[:, 0]
    amount = data[:, 1]
    amount_true = amount * y_true
    amount_pred = amount * y_pred
    error = amount_pred - amount_true
    return sum(error)

y_train_plus_amt = np.append(y_train.values.reshape(-1, 1),
    X_train['Feature1'].values.reshape(-1, 1), axis = 1)

M = Sequential()
M.add(Dense(16, input_shape=(X_train.shape[1],), activation = 'relu'))
M.compile(optimizer='adam', loss = my_loss, run_eagerly = True)
M.fit(X_train, y_train_plus_amt, epochs=10, batch_size=64)

报错信息：
Traceback (most recent call last):
File "", line 1, in
File "/Users/venv/lib/python3.9/site-packages/keras/src/utils/traceback_utils.py", line 70, in error_handler
raise e.with_traceback(filtered_tb) from None
File "", line 5, in my_loss
tensorflow.python.framework.errors_impl.InvalidArgumentError: {{function_node _wrapped__Mul_device/job:localhost/replica:0/task:0/device:CPU:0}} 形状不兼容: [64] vs. [64,16] [Op:Mul] name:

问题根源（非Keras深层问题，为代码设计错误） #

• 模型输出与损失函数形状不匹配：模型最后一层是Dense(16)，输出形状为(batch_size, 16)，但损失函数中amount的形状是(batch_size,)，两者相乘时无法广播，直接触发形状不兼容错误。你的任务是二分类，最后一层应改为Dense(1, activation='sigmoid')，输出单个0-1区间的概率值。
• 损失函数计算逻辑错误：
  1. 使用原生Python的sum()处理TensorFlow张量会引发运算异常，应改用tensorflow.reduce_sum()或tensorflow.reduce_mean()保证张量运算的正确性。
  2. 直接返回损失总和会导致梯度计算异常，通常应返回批次的平均损失或每个样本的损失值。
• 标签输入的形状处理缺失：传入fit的y_train_plus_amt包含标签和特征，但未确保其与模型输出的形状匹配，进一步加剧了形状冲突。

修正后的代码 #

import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow.keras.layers import Dense
from tensorflow.keras.models import Sequential

N = 1000
df = pd.DataFrame({
    'Feature1': np.random.normal(loc=0, scale=1, size=N),
    'Feature2': np.random.normal(loc=1, scale=2, size=N),
    'Label': np.random.choice([0, 1], size=N)
})

df_train = df.sample(frac = 0.80, random_state = 42)
df_test = df[~df.index.isin(df_train.index)]
print(f"df_train.shape = {df_train.shape}")
print(f"df_test.shape = {df_test.shape}")

X_train, y_train = df_train[['Feature1', 'Feature2']], df_train['Label']
X_test, y_test = df_test[['Feature1', 'Feature2']], df_test['Label']

def my_loss(data, y_pred):
    y_true = data[:, 0]
    # 重塑形状，与y_pred的(batch_size,1)匹配
    y_true = tf.reshape(y_true, shape=(-1, 1))
    amount = data[:, 1]
    amount = tf.reshape(amount, shape=(-1, 1))
    
    amount_true = amount * y_true
    amount_pred = amount * y_pred
    error = amount_pred - amount_true
    # 返回批次平均损失，保证梯度计算合理
    return tf.reduce_mean(tf.square(error))

# 用concatenate替代append，保证数组拼接的稳定性
y_train_plus_amt = np.concatenate([y_train.values.reshape(-1, 1),
                                   X_train['Feature1'].values.reshape(-1, 1)], axis=1)

M = Sequential()
M.add(Dense(16, input_shape=(X_train.shape[1],), activation='relu'))
# 二分类任务适配的输出层
M.add(Dense(1, activation='sigmoid'))
M.compile(optimizer='adam', loss=my_loss, run_eagerly=True)
M.fit(X_train, y_train_plus_amt, epochs=10, batch_size=64)

关键修正点说明 #

• 将模型最后一层改为Dense(1, activation='sigmoid')，输出单个概率值，匹配二分类任务需求，同时与损失函数中的张量形状兼容。
• 在损失函数中对y_true和amount进行形状重塑，确保与y_pred的(batch_size,1)形状一致，避免相乘时的形状冲突。
• 使用TensorFlow原生的reduce_mean替代Python原生sum，保证张量运算的正确性，同时返回合理的平均损失值。

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