在TensorFlow中实现用于回归问题的组间MSE差异自定义损失函数

import numpy as np
import tensorflow as tf
from sklearn.datasets import make_regression
from sklearn.model_selection import train_test_split

def custom_group_mse_loss_factory(group_batch_tensor):
    """
    创建一个自定义损失函数，计算两个组的MSE平方差。
    这个函数是一个工厂，它接收当前批次的组标识符张量，并返回一个
    标准的Keras损失函数 (y_true, y_pred) -> loss_value。

    Args:
        group_batch_tensor: 当前训练批次的组标识符张量 (例如，tf.Tensor, shape=(batch_size,))。
                            其中，组标识符为 0 或 1。
    Returns:
        一个可调用的损失函数，接受 y_true 和 y_pred 作为输入。
    """
    def loss(y_true, y_pred):
        # 确保预测值和真实值是扁平的张量
        y_pred = tf.reshape(y_pred, [-1])
        y_true = tf.reshape(y_true, [-1])

        # 根据组标识符创建布尔掩码
        mask_group1 = tf.equal(group_batch_tensor, 1)
        mask_group0 = tf.equal(group_batch_tensor, 0)

        # 使用掩码分离每个组的预测值和真实值
        y_pred_group1 = tf.boolean_mask(y_pred, mask_group1)
        y_pred_group0 = tf.boolean_mask(y_pred, mask_group0)
        y_true_group1 = tf.boolean_mask(y_true, mask_group1)
        y_true_group0 = tf.boolean_mask(y_true, mask_group0)

        # 确保数据类型一致
        y_pred_group1 = tf.cast(y_pred_group1, y_true.dtype)
        y_pred_group0 = tf.cast(y_pred_group0, y_true.dtype)

        # 计算每个组的MSE
        # 为了提高鲁棒性，处理批次中可能出现某个组为空的情况
        # 如果某个组为空，其MSE贡献为0，避免NaN
        mse_group1 = tf.cond(tf.size(y_true_group1) > 0, 
                             lambda: tf.reduce_mean(tf.square(y_true_group1 - y_pred_group1)), 
                             lambda: 0.0)
        mse_group0 = tf.cond(tf.size(y_true_group0) > 0, 
                             lambda: tf.reduce_mean(tf.square(y_true_group0 - y_pred_group0)), 
                             lambda: 0.0)

        # 返回两个组MSE的平方差，以获得更平滑的梯度
        return tf.square(mse_group1 - mse_group0)
    return loss

def train_with_early_stopping(model, loss_fn_factory,
                              X_train, y_train, g_train, X_val, y_val, g_val,
                              n_epoch=500, patience=10, batch_size=64):
    """
    使用自定义损失函数和早停策略训练模型。

    Args:
        model: 待训练的Keras模型。
        loss_fn_factory: 损失函数工厂，接收组标识符张量并返回损失函数。
        X_train, y_train, g_train: 训练集特征、标签和组标识符。
        X_val, y_val, g_val: 验证集特征、标签和组标识符。
        n_epoch: 最大训练轮数。
        patience: 早停的耐心值，即验证损失不再改善的轮数。
        batch_size: 训练批次大小。
    """
    best_val_loss = float('inf')
    wait = 0
    best_epoch = 0
    best_weights = None

    num_samples_train = X_train.shape[0]
    train_indices = np.arange(num_samples_train) # 用于数据混洗的索引数组

    for epoch in range(n_epoch):
        # 每个epoch开始时混洗训练数据
        np.random.shuffle(train_indices)
        X_train_shuffled = X_train[train_indices]
        y_train_shuffled = y_train[train_indices]
        g_train_shuffled = g_train[train