能否在TensorFlow/Keras中自定义加权均方误差损失函数训练模型？

你的需求本质是用样本加权的平方误差来训练模型，其中每个样本的平方误差权重就是数据集中的w列。下面我会一步步带你实现，包括训练和推理阶段的处理：

核心思路 #

Keras默认的损失函数只接收y_true和y_pred，但我们需要用到每个样本对应的w值，所以得把w作为模型的一个额外输入分支。训练时把w喂给模型参与损失计算，推理阶段则只保留特征输入（x1、x2、x3）的分支即可。

完整代码实现 #

1. 导入依赖并准备数据 #

import tensorflow as tf
from tensorflow.keras.layers import Input, Dense
from tensorflow.keras.models import Model
import numpy as np

# 用你提供的示例数据
data = np.array([
    [1, 2, 5, 3, 4],
    [4, 6, 6, 1, 2],
    [5, 2, 7, 3, 1],
    [1, 6, 8, 1, 8]
])
# 拆分特征、权重w、标签y
X = data[:, :3]       # x1, x2, x3
w = data[:, 3:4]      # 保持二维形状，匹配Keras输入层要求
y_true = data[:, 4:5]

2. 构建带自定义损失的模型 #

我们用add_loss方法来绑定自定义损失，这样能直接访问到w输入：

# 特征输入分支（推理阶段只用这个）
input_features = Input(shape=(3,), name="input_x")
# 权重w输入分支（仅训练时使用）
input_w = Input(shape=(1,), name="input_w")

# 模型预测部分（你可以根据需求增加更多层，比如Dropout、激活函数等）
hidden_layer = Dense(16, activation="relu")(input_features)
y_pred = Dense(1, name="prediction")(hidden_layer)

# 自定义损失函数：等价于numpy.dot(w, (y-y_pred)**2)
def weighted_mse(y_true, y_pred, w):
    squared_error = tf.square(y_true - y_pred)
    weighted_error = w * squared_error
    # numpy.dot在这里是对应元素相乘后求和，所以用tf.reduce_sum实现
    return tf.reduce_sum(weighted_error)

# 把损失添加到模型中
model = Model(inputs=[input_features, input_w], outputs=y_pred)
model.add_loss(weighted_mse(y_true, y_pred, input_w))
# 编译模型，选择合适的优化器
model.compile(optimizer="adam")

3. 训练模型 #

训练时需要同时传入特征X和权重w：

model.fit([X, w], epochs=100, verbose=1)

4. 推理阶段（无需w） #

我们可以从训练好的模型中提取出仅包含特征输入和预测输出的子模型，这样推理时就不用管w了：

# 生成推理专用模型
inference_model = Model(inputs=input_features, outputs=y_pred)

# 测试推理
test_X = np.array([[1, 2, 5], [4, 6, 6]])
predictions = inference_model.predict(test_X)
print("推理结果：", predictions)

关键细节说明 #

• 关于损失函数的等价性：numpy.dot(w, (y-y_pred)**2)当w和(y-y_pred)**2都是一维（或列向量）时，本质是对应元素相乘后求和，所以用tf.reduce_sum(w * squared_error)完全等价。
• 为什么不用sample_weight？Keras的sample_weight也能实现样本加权，但它更适合全局或外部传入的权重；而你的w是数据集的一部分，作为模型输入的方式更直观，也更灵活（比如如果后续需要让模型学习权重的话，这种结构也容易扩展）。
• 形状匹配：确保w的形状和y_true、y_pred一致，避免TensorFlow的广播错误，比如我们把w处理成二维列向量就是为了这个。

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