如何在TensorFlow自定义损失函数中获取InputLayer输入向量指定元素

问题原因 #

Sequential模型的损失函数仅默认接收y_true和y_pred两个参数，无法直接获取输入层的张量，因此需要改用函数式API构建模型，灵活传递输入张量到损失函数中。

解决方案1：Lambda包装损失函数（推荐） #

通过函数式API构建模型，用Lambda将输入层张量传入自定义损失函数：

import tensorflow as tf

# 定义输入层与模型结构
input_layer = tf.keras.layers.Input(shape=(x.shape[1],))
first_layer = tf.keras.layers.Dense(128, activation="relu")(input_layer)
output_layer = tf.keras.layers.Dense(y.shape[1], activation="sigmoid")(first_layer)

# 自定义损失函数：接收输入张量、真实标签、预测值
def custom_loss(input_tensor, y_true, y_pred):
    # 计算基础MSE损失
    mse_loss = tf.keras.metrics.mean_squared_error(y_true, y_pred)
    # 获取输入向量的第10个元素（索引9，0起始）
    tenth_element = input_tensor[:, 9]
    # 组合损失（可根据需求调整运算逻辑）
    total_loss = mse_loss + tenth_element
    return total_loss

# 构建函数式模型
model = tf.keras.Model(inputs=input_layer, outputs=output_layer)

# 编译模型：用Lambda将输入层张量传入损失函数
model.compile(
    optimizer="adam",
    loss=lambda y_true, y_pred: custom_loss(input_layer.output, y_true, y_pred),
    metrics=["mae"],
    run_eagerly=True  # 调试时开启，正式训练可移除
)

# 训练模型
model.fit(x=train_input, y=train_output, epochs=5, validation_split=0.5)

解决方案2：模型内部定义损失（Keras 3推荐写法） #

通过add_loss在模型内部直接定义损失逻辑，无需额外包装：

import tensorflow as tf

# 定义输入层与模型结构
input_layer = tf.keras.layers.Input(shape=(x.shape[1],))
first_layer = tf.keras.layers.Dense(128, activation="relu")(input_layer)
output_layer = tf.keras.layers.Dense(y.shape[1], activation="sigmoid")(first_layer)

# 定义真实标签的输入张量
y_true_input = tf.keras.layers.Input(shape=(y.shape[1],))

# 计算损失
mse_loss = tf.keras.metrics.mean_squared_error(y_true_input, output_layer)
tenth_element = input_layer[:, 9]
total_loss = mse_loss + tenth_element

# 构建模型：输入包含原始输入与真实标签
model = tf.keras.Model(inputs=[input_layer, y_true_input], outputs=output_layer)
# 添加损失函数
model.add_loss(total_loss)
# 添加MAE指标
model.add_metric(
    tf.keras.metrics.mean_absolute_error(y_true_input, output_layer),
    name="mae"
)

# 编译模型
model.compile(optimizer="adam", run_eagerly=True)

# 训练：输入为[训练输入, 训练标签]，y参数设为None
model.fit(x=[train_input, train_output], y=None, epochs=5, validation_split=0.5)

关键注意事项 #

• 输入向量的第10个元素对应索引9（Python为0起始索引）。
• 使用input_tensor[:, 9]保证每个样本的第10个元素与MSE损失的形状匹配，避免广播错误。
• run_eagerly=True仅用于调试，正式训练时移除可提升性能。

内容的提问来源于stack exchange，提问作者G. C.