Keras模型中如何将含掩码的样本传入自定义损失函数？

这个问题其实很常见，当你需要在损失函数里用到每个样本专属的额外数据（比如这里的掩码）时，核心思路是把掩码作为模型的一个输入张量，这样每个batch里的样本掩码都能被传递到损失计算的环节中。下面我给你两种清晰的实现方式，都是基于Keras函数式API的（多输入场景下函数式API最灵活）：

方法一：将掩码和真实标签都作为模型输入，直接在模型内定义损失 #

这种方法逻辑清晰，所有输入都明确声明，损失计算过程一目了然。

步骤如下：

1. 定义所有输入张量：包括你的3个标量、2个矩阵（其中一个是掩码），以及真实标签Y。
2. 构建模型主体部分，把输入处理后得到预测值Y_hat。
3. 按照你的公式计算自定义损失：逐点相乘Y_hat和掩码，减去Y后平方，最后取均值。
4. 使用add_loss()把损失添加到模型中，编译时无需指定loss参数。

代码示例：

import tensorflow as tf
from tensorflow import keras

# 1. 定义所有输入：根据你的实际形状调整维度
scalar_input = keras.Input(shape=(3,), name="sample_scalars")
matrix_input = keras.Input(shape=(64, 64), name="sample_matrix")  # 假设矩阵是64x64
mask_input = keras.Input(shape=(64, 64), name="sample_mask")  # 掩码和矩阵形状一致
y_true_input = keras.Input(shape=(64, 64), name="true_labels")

# 2. 构建模型主体（这里只是示例，替换成你自己的模型结构）
# 先把标量和展平后的矩阵拼接
flattened_matrix = keras.layers.Flatten()(matrix_input)
combined_features = keras.layers.concatenate([scalar_input, flattened_matrix])

# 经过几层全连接层得到预测结果，最后reshape成和矩阵一样的形状
x = keras.layers.Dense(128, activation="relu")(combined_features)
x = keras.layers.Dense(256, activation="relu")(x)
y_hat = keras.layers.Dense(64*64, activation="linear")(x)
y_hat = keras.layers.Reshape((64, 64))(y_hat)

# 3. 计算自定义损失
masked_loss = tf.reduce_mean(tf.square(y_hat * mask_input - y_true_input))

# 4. 创建模型并添加损失
model = keras.Model(
    inputs=[scalar_input, matrix_input, mask_input, y_true_input],
    outputs=y_hat
)
model.add_loss(masked_loss)

# 编译模型（不需要指定loss参数）
model.compile(optimizer=keras.optimizers.Adam(learning_rate=1e-4))

# 训练时，传入所有输入数据（假设你已经准备好了对应的数组）
# X_scalars: (num_samples, 3)
# X_matrix: (num_samples, 64, 64)
# X_masks: (num_samples, 64, 64)
# Y_true: (num_samples, 64, 64)
model.fit(
    [X_scalars, X_matrix, X_masks, Y_true],
    epochs=20,
    batch_size=32,
    validation_split=0.1
)

方法二：将掩码与真实标签打包，在损失函数中拆分 #

如果你不想把掩码单独作为输入，也可以把真实标签Y和掩码打包成一个张量，然后在自定义损失函数里拆分它们。这种方法适合你想保持模型输入数量较少的情况。

代码示例：

import tensorflow as tf
from tensorflow import keras

# 定义输入
scalar_input = keras.Input(shape=(3,), name="sample_scalars")
matrix_input = keras.Input(shape=(64, 64), name="sample_matrix")
# 把Y和掩码拼接成一个形状为(64, 128)的张量（前64列是Y，后64列是掩码）
y_mask_input = keras.Input(shape=(64, 128), name="y_and_mask")

# 构建模型主体（和方法一一致）
flattened_matrix = keras.layers.Flatten()(matrix_input)
combined_features = keras.layers.concatenate([scalar_input, flattened_matrix])
x = keras.layers.Dense(128, activation="relu")(combined_features)
x = keras.layers.Dense(256, activation="relu")(x)
y_hat = keras.layers.Dense(64*64, activation="linear")(x)
y_hat = keras.layers.Reshape((64, 64))(y_hat)

# 自定义损失函数：拆分Y和掩码
def masked_mse_loss(y_mask, y_pred):
    y_true = y_mask[..., :64]  # 取前64列作为真实标签
    mask = y_mask[..., 64:]   # 取后64列作为掩码
    return tf.reduce_mean(tf.square(y_pred * mask - y_true))

# 创建模型并编译
model = keras.Model(inputs=[scalar_input, matrix_input, y_mask_input], outputs=y_hat)
model.compile(optimizer=keras.optimizers.Adam(learning_rate=1e-4), loss=masked_mse_loss)

# 训练前需要把Y和掩码拼接：假设Y_true和X_masks都是(64,64)
Y_mask_combined = tf.concat([Y_true, X_masks], axis=-1)

# 开始训练
model.fit(
    [X_scalars, X_matrix, Y_mask_combined],
    epochs=20,
    batch_size=32,
    validation_split=0.1
)

关键注意点 #

• 两种方法的核心都是让每个样本的掩码能被损失函数访问到，因为掩码是样本专属的，不能作为全局变量硬编码在损失函数里。
• 函数式API是处理多输入场景的最佳选择，Sequential模型无法直接支持这种多输入的情况。
• 逐点乘法在TensorFlow中用*或者tf.multiply()都可以，两者是等价的。

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