如何在TensorFlow 2中结合TimeSeriesGenerator实现掩码损失

1. 修正模型输入层定义 #

不要使用Input(tensor=mask)绑定固定张量，这会导致批次输入的shape不匹配。改为定义普通输入层，明确指定掩码的空间维度：

lookback = 7
# 时序数据输入：shape=(时间步长, 图像高度, 图像宽度, 通道数)
inputs = Input(shape=(lookback, 55, 50, 1))
# 掩码输入：shape=(图像高度, 图像宽度, 通道数)，模型会自动适配批次维度
input_mask = Input(shape=(55, 50, 1))

2. 优化自定义损失函数 #

确保仅计算掩码区域内的损失，避免非掩码区域稀释损失值：

def masked_MSE_loss(y_true, y_pred, mask):
    # 对真实值和预测值都应用掩码，过滤非目标区域
    y_true_masked = tf.math.multiply(y_true, mask)
    y_pred_masked = tf.math.multiply(y_pred, mask)
    # 计算掩码区域内的平方误差，再求平均（除以掩码区域元素数保证损失合理性）
    squared_error = tf.square(y_true_masked - y_pred_masked)
    masked_squared_error = tf.math.multiply(squared_error, mask)
    mse = tf.reduce_sum(masked_squared_error) / tf.reduce_sum(mask)
    return mse

3. 绑定模型损失与输入 #

将目标值作为额外输入传入模型，结合掩码和预测结果计算损失：

# 保留原有的ConvLSTM网络结构
convlstm1 = layers.ConvLSTM2D(filters=128, kernel_size=(3, 3), padding='same', activation='tanh', return_sequences=True)(inputs)
bathnorm1 = layers.BatchNormalization()(convlstm1)
convlstm2 = layers.ConvLSTM2D(filters=128, kernel_size=(3, 3), padding='same', activation='tanh', return_sequences=False)(bathnorm1)

convlstm3 = layers.ConvLSTM2D(filters=128, kernel_size=(3, 3), padding='same', activation='tanh', return_sequences=True)(inputs)
batchnorm2 = layers.BatchNormalization()(convlstm3)
convlstm4 = layers.ConvLSTM2D(filters=128, kernel_size=(3, 3), padding='same', activation='tanh', return_sequences=False)(batchnorm2)

concatenation = layers.concatenate([convlstm2, convlstm4])
outputs = layers.Conv2D(filters=1, kernel_size=1, padding="same", activation='tanh')(concatenation)

# 新增目标值输入层，用于损失计算
target_input = Input(shape=(55, 50, 1))
# 绑定自定义损失函数
model.add_loss(masked_MSE_loss(target_input, outputs, input_mask))

# 创建模型：输入包含时序数据、掩码、目标值；输出为预测结果
model = Model(inputs=[inputs, input_mask, target_input], outputs=outputs)
# 编译时无需指定loss，已通过add_loss绑定
model.compile(optimizer='adam', loss=None, metrics=['mae'])

4. 调整训练逻辑 #

将掩码按批次长度重复，匹配输入数据的样本数：

import numpy as np

# 若mask是tensor，先转为numpy数组
mask_np = mask.numpy() if isinstance(mask, tf.Tensor) else mask

for batch in train_gen:
    batch_input, batch_target = batch
    batch_size = batch_input.shape[0]
    # 生成与批次样本数匹配的掩码
    batch_mask = np.repeat(mask_np[np.newaxis, :, :, :], batch_size, axis=0)
    # 传入所有输入训练模型，y设为None（损失已绑定）
    model.fit(x=[batch_input, batch_mask, batch_target], y=None, epochs=1)

可选：自定义数据生成器简化训练 #

包装原TimeseriesGenerator，自动生成带掩码的批次数据：

class MaskedTimeseriesGenerator:
    def __init__(self, base_generator, mask):
        self.base_generator = base_generator
        self.mask = mask
        self.batch_size = base_generator.batch_size
        
    def __len__(self):
        return len(self.base_generator)
    
    def __getitem__(self, idx):
        batch_input, batch_target = self.base_generator[idx]
        # 生成批次对应的掩码
        batch_mask = np.repeat(self.mask[np.newaxis, :, :, :], self.batch_size, axis=0)
        # 返回输入列表和空目标（损失已绑定）
        return [batch_input, batch_mask, batch_target], None

# 使用自定义生成器训练
masked_train_gen = MaskedTimeseriesGenerator(train_gen, mask_np)
model.fit(masked_train_gen, epochs=100)

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