import sys
import os
import pennylane as qml
from pennylane import numpy as np
import tensorflow as tf
from sklearn.datasets import load_wine
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.decomposition import PCA
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
from tensorflow.keras.layers import Dense
from importance_sampling.training import ConstantTimeImportanceTraining  # 确保该库已安装

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
plt.rcParams['axes.unicode_minus'] = False    # 解决负号显示问题

# 重定向 stdout 以捕获打印输出
class Tee(object):
    def __init__(self, *files):
        self.files = files
    def write(self, obj):
        for f in self.files:
            f.write(obj)
            f.flush()
    def flush(self):
        for f in self.files:
            f.flush()

output_file = open('QNNImprove', 'w', encoding='utf-8')
original_stdout = sys.stdout
sys.stdout = Tee(sys.stdout, output_file)

def save_plot(fig, filename):
    fig.savefig(filename)
    plt.close(fig)

def plot_history(history):
    # 绘制训练和验证准确率
    fig, ax = plt.subplots(figsize=(10, 5))
    ax.plot(history.history['accuracy'])
    ax.plot(history.history['val_accuracy'])
    ax.set_title('模型准确率')
    ax.set_ylabel('准确率')
    ax.set_xlabel('轮次')
    ax.legend(['训练', '验证'], loc='upper left')
    save_plot(fig, 'accuracy_plot.png')

    # 绘制训练和验证损失
    fig, ax = plt.subplots(figsize=(10, 5))
    ax.plot(history.history['loss'])
    ax.plot(history.history['val_loss'])
    ax.set_title('模型损失')
    ax.set_ylabel('损失')
    ax.set_xlabel('轮次')
    ax.legend(['训练', '验证'], loc='upper left')
    save_plot(fig, 'loss_plot.png')

def calculate_accuracy_fluctuation(history):
    val_accuracies = history.history['val_accuracy']
    mean_accuracy = np.mean(val_accuracies)
    fluctuation = np.sum(np.abs(val_accuracies - mean_accuracy))
    return mean_accuracy, fluctuation

# 定义设备，使用 8 个量子比特
n_qubits = 8  # 根据需要的特征数量调整量子比特数
dev = qml.device("default.qubit", wires=n_qubits)

def encode_features(inputs, n_qubits):
    """
    使用角度编码将输入特征编码到量子电路中。
    将特征映射到量子比特的旋转角度。
    """
    # 将特征映射到量子比特
    for i in range(n_qubits):
        qml.RY(inputs[i], wires=i)

@qml.qnode(dev, interface='tf', diff_method='backprop')
def quantum_neural_network(inputs, rot_weights):
    # 编码输入特征到量子比特
    encode_features(inputs, n_qubits)

    # 应用变分层
    for layer in range(rot_weights.shape[0]):
        for qubit in range(n_qubits):
            rot_angles = rot_weights[layer, qubit, :]  # 形状：(3,)
            qml.Rot(rot_angles[0], rot_angles[1], rot_angles[2], wires=qubit)
        # 修改纠缠方式，使用全连接纠缠
        for i in range(n_qubits):
            for j in range(i + 1, n_qubits):
                qml.CRX(np.pi / 2, wires=[i, j])

    # 测量期望值
    return [qml.expval(qml.PauliZ(i)) for i in range(n_qubits)]

# 定义用于绘制电路的 QNode
@qml.qnode(dev)
def quantum_neural_network_draw(inputs, rot_weights):
    # 编码输入特征到量子比特
    encode_features(inputs, n_qubits)

    # 应用变分层
    for layer in range(rot_weights.shape[0]):
        for qubit in range(n_qubits):
            rot_angles = rot_weights[layer, qubit, :]  # 形状：(3,)
            qml.Rot(rot_angles[0], rot_angles[1], rot_angles[2], wires=qubit)
        # 修改纠缠方式，使用全连接纠缠
        for i in range(n_qubits):
            for j in range(i + 1, n_qubits):
                qml.CRX(np.pi / 2, wires=[i, j])

    # 测量期望值
    return [qml.expval(qml.PauliZ(i)) for i in range(n_qubits)]

# 定义量子神经网络层
class QuantumNeuralNetworkLayer(tf.keras.layers.Layer):
    def __init__(self, n_qubits, n_layers, **kwargs):
        super(QuantumNeuralNetworkLayer, self).__init__(**kwargs)
        self.n_qubits = n_qubits
        self.n_layers = n_layers

        # 定义旋转权重
        init_weights = tf.random.uniform(
            shape=(n_layers, n_qubits, 3),
            minval=0,
            maxval=2 * np.pi,
            dtype=tf.float32
        )
        self.rot_weights = tf.Variable(
            initial_value=init_weights,
            trainable=True,
            name='rot_weights'
        )

    def call(self, inputs):
        def apply_qnode(x):
            x = tf.cast(x, dtype=tf.float32)
            result = quantum_neural_network(x, self.rot_weights)
            return tf.cast(result, dtype=tf.float32)

        outputs = tf.map_fn(
            apply_qnode,
            inputs,
            dtype=tf.float32
        )

        return outputs

# 定义量子神经网络模型
class QuantumNeuralNetwork(tf.keras.Model):
    def __init__(self, n_qubits, n_layers, num_classes, **kwargs):
        super(QuantumNeuralNetwork, self).__init__(**kwargs)
        self.n_qubits = n_qubits
        self.n_layers = n_layers
        self.num_classes = num_classes

        # 量子层
        self.quantum_layer = QuantumNeuralNetworkLayer(n_qubits, n_layers)

        # 经典输出层
        self.output_layer = tf.keras.layers.Dense(num_classes, activation='softmax')

    def call(self, inputs):
        q_output = self.quantum_layer(inputs)
        return self.output_layer(q_output)

if __name__ == "__main__":
    # **重要性采样始终启用**

    # 加载 Wine 数据集
    wine = load_wine()
    X = wine['data'].astype(np.float32)
    y = wine['target'].reshape(-1, 1)

    # One-hot 编码目标标签
    encoder = OneHotEncoder(sparse_output=False)
    y_encoded = encoder.fit_transform(y).astype(np.float32)
    num_classes = y_encoded.shape[1]

    # 先划分数据集
    x_train, x_test, y_train, y_test = train_test_split(
        X, y_encoded, test_size=0.3, random_state=42
    )

    # 特征标准化 (仅在训练集上拟合)
    scaler = StandardScaler()
    x_train_scaled = scaler.fit_transform(x_train)
    x_test_scaled = scaler.transform(x_test)

    # 使用 PCA 进行降维 (仅在训练集上拟合)
    n_features = n_qubits  # 将特征数量降至与量子比特数量一致
    pca = PCA(n_components=n_features)
    x_train_reduced = pca.fit_transform(x_train_scaled)
    x_test_reduced = pca.transform(x_test_scaled)

    # 将特征归一化到 [0, π] (基于训练集的最小值和最大值)
    x_min = x_train_reduced.min(axis=0)
    x_max = x_train_reduced.max(axis=0)
    x_train_normalized = (x_train_reduced - x_min) / (x_max - x_min) * np.pi
    x_test_normalized = (x_test_reduced - x_min) / (x_max - x_min) * np.pi

    x_train = x_train_normalized
    x_test = x_test_normalized

    learning_rate = 0.01  # 学习率
    epochs = 200  # 训练轮次
    batch_size = 16
    n_layers = 5  # 增加量子层数

    model = QuantumNeuralNetwork(
        n_qubits=n_qubits,
        n_layers=n_layers,
        num_classes=num_classes
    )

    # 使用默认的 'categorical_crossentropy'，不手动设置 'reduction'
    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
        loss='categorical_crossentropy',
        metrics=['accuracy']
    )

    def is_quantum_model(model):
        for layer in model.layers:
            if isinstance(layer, QuantumNeuralNetworkLayer):
                return True
        return False

    if is_quantum_model(model):
        print("模型包含量子神经网络（QNN）层。")
    else:
        print("模型不包含量子神经网络（QNN）层。")

    print("\n量子电路图：")
    sample_inputs = np.ones(n_qubits, dtype=np.float32) * (np.pi / 2)  # 使用 n_qubits 个特征
    sample_rot_weights = np.zeros((n_layers, n_qubits, 3), dtype=np.float32)  # rot_weights

    print(qml.draw(quantum_neural_network_draw)(sample_inputs, sample_rot_weights))  # 使用绘制专用的 QNode

    # 早停和学习率调整回调函数
    callbacks = [
        tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=True),
        tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=15)
    ]

    # **始终使用重要性采样进行训练**
    wrapped_model = ConstantTimeImportanceTraining(model, score="loss")
    history = wrapped_model.fit(
        x_train, y_train,
        batch_size=batch_size,
        epochs=epochs,
        verbose=1,
        validation_data=(x_test, y_test),
        callbacks=callbacks
    )

    score = model.evaluate(x_test, y_test, verbose=0)
    print('测试损失:', score[0])
    print('测试准确率:', score[1])

    mean_accuracy, fluctuation = calculate_accuracy_fluctuation(history)
    print('平均验证准确率:', mean_accuracy)
    print('准确率波动:', fluctuation)

    plot_history(history)

    # 关闭输出文件并恢复 stdout
    sys.stdout = original_stdout
    output_file.close()

    print("脚本执行完成。查看 'output.log' 获取打印输出，查看当前目录获取保存的图表。")
上述是wine_QNNimprove.py文件
import sys
import numpy as np
import logging
from collections.abc import Iterable
import tensorflow as tf
from tensorflow.python.keras.callbacks import BaseLogger, CallbackList, History, ProgbarLogger
from tensorflow.python.keras.layers import Dropout
#from tensorflow.python.keras.utils import Sequence
from blinker import signal
from .datasets import InMemoryDataset, GeneratorDataset
from .model_wrappers import OracleWrapper
from .samplers import ConditionalStartSampler, VarianceReductionCondition, \
    AdaptiveAdditiveSmoothingSampler, AdditiveSmoothingSampler, ModelSampler, \
    ConstantVarianceSampler, ConstantTimeSampler, SCSGSampler
from .reweighting import BiasedReweightingPolicy, NoReweightingPolicy
from .utils.functional import ___, compose, partial

class _BaseImportanceTraining(object):
    def __init__(self, model, score="gnorm", layer=None):
        """重要性采样训练的基类。

        参数：
            model: 要训练的 Keras 模型。
        """
        # 检查模型是否兼容，并包装模型以支持重要性采样
        self._check_model(model)
        self.original_model = model
        self.model = OracleWrapper(
            model,
            self.reweighting,
            score=score,
            layer=layer
        )
        self.original_model.optimizer = self.model.optimizer
        signal("is.sample").connect(self._on_sample)
        signal("is.score").connect(self._on_scores)

    def _on_sample(self, event):
        self._latest_sample_event = event

    def _on_scores(self, scores):
        self._latest_scores = scores

    def _check_model(self, model):
        """检查模型是否使用了 Dropout 和 BatchNorm，并警告可能影响重要性计算。"""
        if any(
                isinstance(l, (tf.keras.layers.BatchNormalization, Dropout))
                for l in model.layers
        ):
            sys.stderr.write(("[NOTICE]: 您正在使用 BatchNormalization 和/或 Dropout。\n"
                              "这些层可能会影响重要性计算，建议您用 LayerNormalization 或 "
                              "在测试模式下的 BatchNormalization 以及 L2 正则化替代它们。\n"))

    @property
    def reweighting(self):
        """控制重要性采样时偏置的重加权策略。"""
        raise NotImplementedError()

    def sampler(self, dataset, batch_size, steps_per_epoch, epochs):
        """为给定的数据集创建一个新的采样器，以使用重要性采样进行采样。"""
        raise NotImplementedError()

    def fit(self, x, y, batch_size=32, epochs=1, verbose=1, callbacks=None,
            validation_split=0.0, validation_data=None, steps_per_epoch=None,
            on_sample=None, on_scores=None):
        """使用给定的数据创建一个 InMemoryDataset 实例，并使用重要性采样训练模型。

        参数：
            x: 训练数据的 Numpy 数组或 Numpy 数组列表。
            y: 目标数据的 Numpy 数组。
            batch_size: 每个梯度更新的样本数。
            epochs: 遍历整个训练集的次数。
            verbose: {0, >0}，是否使用 Keras 的进度条回调。
            callbacks: 在训练期间调用的 Keras 回调列表。
            validation_split: [0, 1) 之间的浮点数，用于验证的数据百分比。
            validation_data: tuple，(x_val, y_val)，用于模型验证的数据。
            steps_per_epoch: int 或 None，每个 epoch 的步数。
            on_sample: 接受 sampler, idxs, w, scores 的回调函数。
            on_scores: 接受 sampler 和 scores 的回调函数。
        返回：
            包含训练期间收集的信息的 Keras `History` 对象。
        """
        # 处理验证数据集
        if validation_data is not None:
            x_train, y_train = x, y
            x_test, y_test = validation_data
        elif validation_split > 0:
            assert validation_split < 1, "验证集比例必须小于 1"
            n = int(round(validation_split * len(x)))
            idxs = np.arange(len(x))
            np.random.shuffle(idxs)
            x_train, y_train = x[idxs[n:]], y[idxs[n:]]
            x_test, y_test = x[idxs[:n]], y[idxs[:n]]
        else:
            x_train, y_train = x, y
            x_test, y_test = np.empty(shape=(0, x.shape[1])), np.empty(shape=(0, y.shape[1]))

        # 创建数据集
        dataset = InMemoryDataset(
            x_train,
            y_train,
            x_test,
            y_test,
            categorical=False  # 不需要转换标签
        )

        return self.fit_dataset(
            dataset=dataset,
            batch_size=batch_size,
            epochs=epochs,
            steps_per_epoch=steps_per_epoch,
            verbose=verbose,
            callbacks=callbacks,
            on_sample=on_sample,
            on_scores=on_scores
        )

    def fit_generator(self, train, steps_per_epoch=None, batch_size=32,
                      epochs=1, verbose=1, callbacks=None,
                      validation_data=None, validation_steps=None,
                      on_sample=None, on_scores=None):
        """创建一个 GeneratorDataset 实例，并使用重要性采样训练模型。

        参数：
            train: 返回 (inputs, targets) 的生成器或可迭代对象。
            steps_per_epoch: 每个 epoch 的梯度更新次数。
            batch_size: 每个梯度更新的样本数。
            epochs: 训练轮数。
            verbose: {0, >0}，是否使用 Keras 的进度条回调。
            validation_data: 生成器、可迭代对象或 (inputs, targets) 元组。
            validation_steps: 如果 validation_data 是生成器，则需要指定此参数。
            on_sample: 接受 sampler, idxs, w, scores 的回调函数。
            on_scores: 接受 sampler 和 scores 的回调函数。
        """

        def infinite_iterator(iterable):
            while True:
                yield from iterable

        # 处理验证数据集
        if validation_data is not None:
            if isinstance(validation_data, (tuple, list)):
                test = validation_data
                test_len = None
            elif isinstance(validation_data, Iterable):
                test = infinite_iterator(validation_data)
                try:
                    test_len = len(validation_data) * batch_size
                except TypeError:
                    test_len = validation_steps * batch_size if validation_steps else None
            else:
                test = validation_data
                test_len = validation_steps * batch_size if validation_steps else None
        else:
            test = (np.empty(shape=(0, 1)), np.empty(shape=(0, 1)))
            test_len = None
            logging.info("未提供验证数据")

        # 设置训练数据和每个 epoch 的步数
        if isinstance(train, Iterable):
            if steps_per_epoch is None:
                try:
                    steps_per_epoch = len(train)
                except TypeError:
                    raise ValueError("使用没有 __len__ 方法的生成器时必须设置 steps_per_epoch")
            train = infinite_iterator(train)
        elif steps_per_epoch is None:
            raise ValueError("steps_per_epoch 未设置")

        dataset = GeneratorDataset(train, test, test_len)

        return self.fit_dataset(
            dataset=dataset,
            batch_size=batch_size,
            epochs=epochs,
            steps_per_epoch=steps_per_epoch,
            verbose=verbose,
            callbacks=callbacks,
            on_sample=on_sample,
            on_scores=on_scores
        )

    def fit_dataset(self, dataset, steps_per_epoch=None, batch_size=32,
                    epochs=1, verbose=1, callbacks=None, on_sample=None,
                    on_scores=None):
        """在给定的数据集上训练模型。

        参数：
            dataset: `BaseDataset` 的实例，提供要训练的数据。
            steps_per_epoch: int，每个 epoch 的梯度更新次数。
            batch_size: 每个梯度更新的样本数。
            epochs: 训练轮数。
            verbose: {0, >0}，是否使用 Keras 的进度条回调。
            callbacks: Keras 回调列表。
            on_sample: 接受 sampler, idxs, w, scores 的回调函数。
            on_scores: 接受 sampler 和 scores 的回调函数。
        """
        try:
            if len(dataset.train_data) < batch_size:
                raise ValueError("batch_size 大于训练集大小")
        except RuntimeError as e:
            assert "no size" in str(e)

        # 设置 steps_per_epoch
        if steps_per_epoch is None:
            steps_per_epoch = len(dataset.train_data) // batch_size

        # 创建回调列表
        self.history = History()
        callbacks = [BaseLogger()] + (callbacks or []) + [self.history]
        if verbose > 0:
            callbacks += [ProgbarLogger(count_mode="steps")]
        callbacks = CallbackList(callbacks)
        callbacks.set_model(self.original_model)
        callbacks.set_params({
            "epochs": epochs,
            "steps": steps_per_epoch,
            "verbose": verbose,
            "do_validation": len(dataset.test_data) > 0,
            "metrics": self.metrics_names + [
                "val_" + n for n in self.metrics_names
            ]
        })

        # 创建采样器
        sampler = self.sampler(dataset, batch_size, steps_per_epoch, epochs)

        # 开始训练循环
        epoch = 0
        self.original_model.stop_training = False
        callbacks.on_train_begin()
        while epoch < epochs:
            callbacks.on_epoch_begin(epoch)

            for step in range(steps_per_epoch):
                batch_logs = {"batch": step, "size": batch_size}
                callbacks.on_batch_begin(step, batch_logs)

                # 在这里进行重要性采样
                idxs, (x_batch, y_batch), w = sampler.sample(batch_size)

                # 将 x_batch 和 y_batch 转换为 TensorFlow 张量
                x_batch = tf.convert_to_tensor(x_batch)
                y_batch = tf.convert_to_tensor(y_batch)
                w = tf.convert_to_tensor(w, dtype=tf.float32)

                # 自定义训练步骤
                loss_value, metrics_results, scores = self.train_batch(x_batch, y_batch, w)

                # 更新采样器
                sampler.update(idxs, scores)

                # 记录指标
                batch_logs['loss'] = loss_value
                for metric_name, metric_value in zip(self.metrics_names[1:], metrics_results):
                    batch_logs[metric_name] = metric_value
                callbacks.on_batch_end(step, batch_logs)

                if on_scores is not None and hasattr(self, "_latest_scores"):
                    on_scores(
                        sampler,
                        self._latest_scores
                    )

                if on_sample is not None:
                    on_sample(
                        sampler,
                        self._latest_sample_event["idxs"],
                        self._latest_sample_event["w"],
                        self._latest_sample_event["predicted_scores"]
                    )

                if self.original_model.stop_training:
                    break

            # 在每个 epoch 结束时进行评估
            epoch_logs = {}
            if len(dataset.test_data) > 0:
                val_loss, val_metrics = self.evaluate(dataset.test_data)
                epoch_logs['val_loss'] = val_loss
                for name, value in zip(self.metrics_names[1:], val_metrics):
                    epoch_logs['val_' + name] = value
            callbacks.on_epoch_end(epoch, epoch_logs)
            if self.original_model.stop_training:
                break
            epoch += 1
        callbacks.on_train_end()

        return self.history

    @property
    def metrics_names(self):
        def name(x):
            if isinstance(x, str):
                return x
            if hasattr(x, "name"):
                return x.name
            if hasattr(x, "__name__"):
                return x.__name__
            return str(x)

        metrics = self.model.model.metrics or []
        return (
                ["loss"] +
                [name(m) for m in metrics]
        )

    def train_batch(self, x, y, w):
        with tf.GradientTape() as tape:
            y_pred = self.model.model(x, training=True)
            loss_fn = self.model.model.loss  # 已在编译时设置 reduction='AUTO'

            # 计算每个样本的损失，不再传递 reduction 参数
            per_sample_losses = self.model.model.compiled_loss(
                y,
                y_pred,
                sample_weight=None,
                regularization_losses=self.model.model.losses
                # 移除 reduction 参数
            )

            per_sample_losses = tf.cast(per_sample_losses, tf.float32)
            weights = tf.cast(w, tf.float32)

            # 计算加权损失
            weighted_losses = per_sample_losses * weights
            loss = tf.reduce_mean(weighted_losses)

            # 计算梯度并更新模型
            trainable_vars = self.model.model.trainable_variables
            gradients = tape.gradient(loss, trainable_vars)
            self.model.model.optimizer.apply_gradients(zip(gradients, trainable_vars))

            # 更新指标
            self.model.model.compiled_metrics.update_state(y, y_pred, sample_weight=weights)
            metrics = [metric.result().numpy() for metric in self.model.model.metrics]

        return loss.numpy(), metrics, per_sample_losses.numpy()

    def evaluate(self, test_data):
        print(f"test_data type: {type(test_data)}, length: {len(test_data)}")
        print(f"test_data contents: {test_data}")

        x_test, y_test = test_data[:2]
        # 只获取前两个元素
        x_test, y_test = test_data[:2]

        # 如果 x_test 或 y_test 是元组，提取其中的数组
        if isinstance(x_test, tuple) or isinstance(x_test, list):
            x_test = x_test[0]
        if isinstance(y_test, tuple) or isinstance(y_test, list):
            y_test = y_test[0]

        # 确保 x_test 和 y_test 是 NumPy 数组或张量
        x_test = np.asarray(x_test)
        y_test = np.asarray(y_test)

        y_pred = self.model.model.predict(x_test, batch_size=128)

        # 确保 y_test 和 y_pred 是张量
        y_test = tf.convert_to_tensor(y_test)
        y_pred = tf.convert_to_tensor(y_pred)

        loss_value = self.model.model.compiled_loss(y_test, y_pred, regularization_losses=self.model.model.losses)

        # 更新指标
        self.model.model.compiled_metrics.update_state(y_test, y_pred)
        metrics_results = [m.result().numpy() for m in self.model.model.metrics]

        return loss_value.numpy(), metrics_results

class _UnbiasedImportanceTraining(_BaseImportanceTraining):
    def __init__(self, model, score="gnorm", layer=None):
        self._reweighting = BiasedReweightingPolicy(1.0)  # 无偏

        super(_UnbiasedImportanceTraining, self).__init__(model, score, layer)

    @property
    def reweighting(self):
        return self._reweighting

class ConstantTimeImportanceTraining(_UnbiasedImportanceTraining):
    """使用常时间重要性采样训练模型。

    参数：
        model: 要训练的 Keras 模型。
        score: {"gnorm", "loss", "full_gnorm"}，用于重要性采样的重要性度量。
        layer: None 或 int 或 Layer，计算 gnorm 时使用的层。
    """

    def __init__(self, model, backward_pass_multiplier=2.0, extra_samples=0.1,
                 score="gnorm", layer=None):
        self._backward_pass_multiplier = backward_pass_multiplier
        self._extra_samples = extra_samples

        super(ConstantTimeImportanceTraining, self).__init__(
            model,
            score,
            layer
        )

    def sampler(self, dataset, batch_size, steps_per_epoch, epochs):
        return ConstantTimeSampler(
            dataset,
            self.reweighting,
            self.model,
            backward_time=self._backward_pass_multiplier,
        )
上述是training.py文件

import numpy as np
from blinker import signal
from tensorflow.python.keras import backend as K
from tensorflow.python.keras.layers import Layer
from tensorflow.python.keras.metrics import categorical_accuracy
import tensorflow as tf
from .layers import GradientNormLayer, LossLayer, MetricLayer

class ModelWrapper(object):
    """ModelWrapper 的目的是包装一个神经网络，并添加一些额外的层，以生成得分、损失和样本权重，以执行重要性采样。"""

    def _iterate_batches(self, x, y, batch_size):
        bs = batch_size
        for s in range(0, len(y), bs):
            if isinstance(x, (list, tuple)):
                xs = [xi[s:s+bs] for xi in x]
            else:
                xs = x[s:s+bs]
            yield xs, y[s:s+bs]

    def evaluate(self, x, y, batch_size=128):
        result = np.mean(
            np.vstack([
                self.evaluate_batch(xi, yi)
                for xi, yi in self._iterate_batches(x, y, batch_size)
            ]),
            axis=0
        )

        signal("is.evaluation").send(result)
        return result

    def score(self, x, y, batch_size=128):
        bs = batch_size
        result = np.hstack([
            self.score_batch(xi, yi).T
            for xi, yi in self._iterate_batches(x, y, batch_size)
        ]).T

        signal("is.score").send(result)
        return result

    def set_lr(self, lr):
        """设置模型的学习率。"""

        try:
            K.set_value(
                self.optimizer.lr,
                lr
            )
        except AttributeError:
            try:
                K.set_value(
                    self.model.optimizer.lr,
                    lr
                )
            except AttributeError:
                raise NotImplementedError()

        try:
            K.set_value(
                self.small.optimizer.lr,
                lr
            )
        except AttributeError:
            pass

    def evaluate_batch(self, x, y):
        raise NotImplementedError()

    def score_batch(self, x, y):
        raise NotImplementedError()

    def train_batch(self, x, y, w):
        raise NotImplementedError()

class OracleWrapper(ModelWrapper):
    AVG_LOSS = 0
    LOSS = 1
    WEIGHTED_LOSS = 2
    SCORE = 3
    METRIC0 = 4

    FUSED_ACTIVATION_WARNING = ("[WARNING]: 最后一层具有融合激活函数，例如 Dense(..., activation=\"sigmoid\")。\n"
                                "为了自动提取预激活张量，请使用单独的激活层（参见示例）。\n")

    def __init__(self, model, reweighting, score="loss", layer=None):
        self.reweighting = reweighting
        self.layer = self._gnorm_layer(model, layer)
        self._augment_model(model, score, reweighting)

    def _gnorm_layer(self, model, layer):
        # 如果提供了层，则直接使用
        if isinstance(layer, Layer):
            return layer

        # 如果提供了层的索引，则提取该层
        if isinstance(layer, int):
            return model.layers[layer]

        # 对于量子神经网络，可能无法自动获取梯度范数层
        # 这里暂时返回 None，表示不使用特定的层
        return None

    def _augment_model(self, model, score, reweighting):
        # 保存原始模型
        self.original_model = model

        # 创建包装后的模型，不修改原始模型的结构
        self.model = model  # 使用原始模型

        # 保存必要的信息
        self.score_method = score
        self.reweighting = reweighting

        # 提取模型的一些信息
        self.loss_fn = model.loss
        self.optimizer = model.optimizer.__class__.from_config(model.optimizer.get_config())
        self.metrics = model.metrics or []

    def evaluate_batch(self, x, y):
        # 使用 model.evaluate
        outputs = self.model.evaluate(
            x=x,
            y=y,
            verbose=0,
            return_dict=True
        )
        loss = outputs['loss']
        metrics = [outputs[m] for m in self.model.metrics_names[1:]]

        signal("is.evaluate_batch").send(outputs)

        # 返回损失和其他指标
        return [loss] + metrics

    # 在 model_wrappers.py 文件中
    def score_batch(self, x, y):
        y_pred = self.model(x, training=False)
        loss_fn = self.model.loss

        # 如果 loss_fn 是字符串，使用 tf.keras.losses.get 将其转换为损失函数对象
        if isinstance(loss_fn, str):
            loss_fn = tf.keras.losses.get(loss_fn)
        elif isinstance(loss_fn, dict):
            loss_fn = tf.keras.losses.get(loss_fn)
        elif not callable(loss_fn):
            raise ValueError(f"Unsupported loss function type: {type(loss_fn)}")

        # 确保损失函数具有 reduction 参数，并设置为 NONE
        if isinstance(loss_fn, tf.keras.losses.Loss):
            loss_config = loss_fn.get_config()
            per_example_loss_fn = loss_fn.__class__(reduction=tf.keras.losses.Reduction.NONE, **loss_config)
            per_sample_losses = per_example_loss_fn(y, y_pred)
        else:
            per_sample_losses = loss_fn(y, y_pred)

        if self.score_method == "loss":
            scores = per_sample_losses
        else:
            scores = tf.ones_like(per_sample_losses)

        return scores.numpy()

    def train_batch(self, x, y, w):
        with tf.GradientTape() as tape:
            y_pred = self.model(x, training=True)
            # 计算每个样本的损失（不进行 reduction）
            per_sample_losses = self.model.compiled_loss(
                y,
                y_pred,
                sample_weight=None,
                regularization_losses=self.model.losses,
                reduction=tf.keras.losses.Reduction.NONE
            )
            # 计算得分
            if self.score_method == "loss":
                scores = per_sample_losses
            else:
                # 实现其他得分方式
                scores = tf.ones_like(per_sample_losses)

            # 计算样本权重
            weights = self.reweighting.weight_layer()([scores, w])

            # 计算加权损失
            weighted_losses = per_sample_losses * weights
            loss = tf.reduce_mean(weighted_losses)

        # 获取模型的可训练参数
        trainable_vars = self.model.trainable_variables

        # 计算梯度
        gradients = tape.gradient(loss, trainable_vars)

        # 更新模型参数
        self.model.optimizer.apply_gradients(zip(gradients, trainable_vars))

        # 更新指标
        self.model.compiled_metrics.update_state(y, y_pred, sample_weight=weights)
        metrics = [metric.result().numpy() for metric in self.model.metrics]

        # 返回损失、指标和得分
        return loss.numpy(), metrics, scores.numpy()

    def __call__(self, x, training=False):
        return self.model(x, training=training)
上述是model_wrappers.py文件
#
# Copyright (c) 2017 Idiap Research Institute, http://www.idiap.ch/
# Written by Angelos Katharopoulos <<PRESIDIO_ANONYMIZED_EMAIL_ADDRESS>>
#

from tensorflow.python.keras import backend as K
from tensorflow.python.keras.layers import Layer
import numpy as np


class ReweightingPolicy(object):
    """ReweightingPolicy defines how we weigh the samples given their
    importance.
    
    Each policy should provide
        1. A layer implementation for use with Keras models
        2. A python implementation for use with the samplers
    """
    def weight_layer(self):
        """Return a layer that accepts the scores and the return value of the
        sample_weights() method and produces new sample weights."""
        raise NotImplementedError()

    def sample_weights(self, idxs, scores):
        """Given the scores and the chosen indices return whatever is needed by
        the weight_layer to produce the sample weights."""
        raise NotImplementedError()

    @property
    def weight_size(self):
        """Return how many numbers per sample make up the sample weights"""
        raise NotImplementedError()


class AdjustedBiasedReweightingPolicy(ReweightingPolicy):
    """AdjustedBiasedReweightingPolicy adjusts the biased sample weights with
    the importance that has just been computed in the forward-backward pass.

    See AdjustedBiasedReweighting for details.
    """
    def __init__(self, k=1.0):
        self.k = k

    def weight_layer(self):
        return AdjustedBiasedReweighting(self.k)

    def sample_weights(self, idxs, scores):
        N = len(scores)
        S1 = scores[np.setdiff1d(np.arange(N), idxs)].sum()

        return np.tile([float(N), float(S1)], (len(idxs), 1))

    @property
    def weight_size(self):
        return 2


class BiasedReweightingPolicy(ReweightingPolicy):
    """BiasedReweightingPolicy computes the sample weights before  the
    forward-backward pass based on the sampling probabilities. It can introduce
    a bias that focuses on the hard examples when combined with the loss as an
    importance metric."""
    def __init__(self, k=1.0):
        self.k = k

    def weight_layer(self):
        return ExternalReweighting()

    def sample_weights(self, idxs, scores):
        N = len(scores)
        s = scores  # 因为 scores 对应选定的样本
        w = s.sum() / N / s
        w_hat = w ** self.k
        w_hat *= np.dot(w, s) / np.dot(w_hat, s)
        return w_hat

    @property
    def weight_size(self):
        return 1

class NoReweightingPolicy(ReweightingPolicy):
    """Set all sample weights to 1."""
    def weight_layer(self):
        return ExternalReweighting()

    def sample_weights(self, idxs, scores):
        return np.ones((len(idxs), 1))

    @property
    def weight_size(self):
        return 1


class CorrectingReweightingPolicy(ReweightingPolicy):
    """CorrectingReweightingPolicy aims to scale the sample weights according
    to the mistakes of the importance predictor

    Arguments
    ---------
    k: float
       The bias power used in all other reweighting schemes
    """
    def __init__(self, k=1.0):
        self.k = k
        self._biased_reweighting = BiasedReweightingPolicy(k)

    def weight_layer(self):
        return CorrectingReweighting()

    def sample_weights(self, idxs, scores):
        w = self._biased_reweighting.sample_weights(idxs, scores)

        return np.hstack([w, scores[idxs][:, np.newaxis]])

    @property
    def weight_size(self):
        return 2


class AdjustedBiasedReweighting(Layer):
    """Implement a Keras layer that using the sum of the weights and the number
    of samples it recomputes the weights.
    
    The specifics are the following:

    Given B = {i_1, i_2, ..., i_|B|} the mini-batch idexes, \\hat{s_i} the
    predicted score of sample i and k the max-loss bias constant.

        S_1 = \\sum_{i \\notin B} \\hat{s_i}
        S_2 = \\sum_{i \\in B} s_i
        a_i^{(1)} = \\frac{1}{N} (\\frac{S1 + S2}{s_i})^k \\forall i \\in B
        a_i^{(2)} = \\frac{1}{N} \\frac{S1 + S2}{s_i} \\forall i \\in B
        t = \\frac{\\sum_{i} a_i^{(2)} s_i }{\\sum_{i} a_i^{(1)} s_i}
        a_i = t a_i^{(1)} \\forall i \\in B
    """
    def __init__(self, k=1.0, **kwargs):
        self.k = k

        super(AdjustedBiasedReweighting, self).__init__(**kwargs)

    def build(self, input_shape):
        assert isinstance(input_shape, list)
        assert len(input_shape) == 2
        assert input_shape[0][1] == 1
        assert input_shape[1][1] == 2

        super(AdjustedBiasedReweighting, self).build(input_shape)

    def compute_output_shape(self, input_shape):
        return input_shape[0]

    def call(self, x):
        s, s_hat = x

        # Compute the variables defined in the class comment
        S2 = K.sum(s)
        S1 = s_hat[0, 1]
        N = s_hat[0, 0]

        # Compute the unbiased weights
        a2 = (S1 + S2) / N / s

        # Compute the biased weights and the scaling factor t
        a1 = K.pow(a2, self.k)
        sT = K.transpose(s)
        t = K.dot(sT, a2) / K.dot(sT, a1)

        return K.stop_gradient([a1 * t])[0]


class ExternalReweighting(Layer):
    """Use the provided input as sample weights"""
    def build(self, input_shape):
        super(ExternalReweighting, self).build(input_shape)

    def compute_output_shape(self, input_shape):
        return input_shape[0]

    def call(self, x):
        return K.stop_gradient(x[1])


class CorrectingReweighting(Layer):
    """Use the provided weights and the score to correct sample weights that
    were computed with a very wrong predicted score"""
    def __init__(self, min_decrease=0, max_increase=2, **kwargs):
        self.min_decrease = min_decrease
        self.max_increase = max_increase

        super(CorrectingReweighting, self).__init__(**kwargs)

    def build(self, input_shape):
        super(CorrectingReweighting, self).build(input_shape)

    def compute_output_shape(self, input_shape):
        return input_shape[0]

    def call(self, x):
        s, x1 = x
        a = x1[:, :1]
        s_hat = x1[:, 1:2]

        # Rescale the weights, making sure we mostly scale down
        a_hat = a * K.clip(s_hat / s, self.min_decrease, self.max_increase)

        # Scale again so that the reported loss is comparable to the other ones
        t = 1
        #sT = K.transpose(s)
        #t = K.dot(sT, a) / K.dot(sT, a_hat)

        return K.stop_gradient([a_hat * t])[0]


UNWEIGHTED = NoReweightingPolicy()
UNBIASED = BiasedReweightingPolicy(k=1.0)
上述是reweighting.py文件
def _get_dataset_length(dset, default=1):
    """返回数据集的训练数据长度，如果无法获取，则返回默认值。"""
    try:
        return len(dset.train_data)
    except RuntimeError:
        return default

import numpy as np
from blinker import signal

class BaseSampler(object):
    def __init__(self, dataset, reweighting):
        self.dataset = dataset
        self.reweighting = reweighting

    def _slice_data(self, x, y, idxs):
        if isinstance(x, (list, tuple)):
            return [xi[idxs] for xi in x], y[idxs]
        else:
            return x[idxs], y[idxs]

    def _send_messages(self, idxs, xy, w, predicted_scores):
        signal("is.sample").send({
            "idxs": idxs,
            "xy": xy,
            "w": w,
            "predicted_scores": predicted_scores
        })

    def _get_samples_with_scores(self, batch_size):
        raise NotImplementedError()

    def sample(self, batch_size):
        # 获取样本和分数
        idxs1, scores, xy = self._get_samples_with_scores(batch_size)

        if scores is not None:
            # 使用 argsort 对分数进行排序（降序）
            idxs_sorted = np.argsort(-scores)
            # 调用子类的采样方法
            selected_indices, weights = self._select_samples(idxs_sorted, scores, batch_size)
        else:
            # 如果没有分数，则随机采样
            selected_indices = np.random.choice(len(idxs1), batch_size, replace=False)
            weights = np.ones(len(selected_indices))

        # 获取选中的数据
        if xy is None:
            # 如果没有预先加载数据，则从数据集中提取
            x_selected, y_selected = self._slice_data(self.dataset._x_train, self.dataset._y_train, selected_indices)
        else:
            x, y = xy
            x_selected, y_selected = self._slice_data(x, y, selected_indices)

        # 处理分数
        scores_selected = scores[selected_indices] if scores is not None else np.ones(batch_size)
        self._send_messages(selected_indices, (x_selected, y_selected), weights, scores_selected)
        return selected_indices, (x_selected, y_selected), weights

    def update(self, idxs, results):
        pass

    def _select_samples(self, idxs_sorted, scores, batch_size):
        raise NotImplementedError()

class ConstantTimeSampler(BaseSampler):
    def __init__(self, dataset, reweighting, model, B=32, batch_size=16, K_ratio=0.5,
                 backward_time=2.0, tau_th=1.0, ratio=0.5, min_a=0.2):
        """
        Args:
            dataset: The dataset to sample from.
            reweighting: Reweighting policy.
            model: The model wrapper with score method.
            B: Number of samples to uniformly sample in the first stage.
            batch_size: Total batch size for training.
            K_ratio: Ratio of samples selected by the optimizer (0 < K_ratio < 1).
            backward_time: Determines the ratio for importance sampling.
            tau_th: Variance reduction threshold.
            ratio: Ratio to split between high-score and random samples.
            min_a: Minimum number of random samples.
        """
        self.model = model
        self.N = len(dataset._x_train)  # 使用 _x_train 获取数据集大小
        self.B = min(B, self.N)  # 确保 B 不超过 N
        self.batch_size = batch_size
        self.K_ratio = K_ratio  # K_ratio = K / batch_size
        self.backward_time = backward_time
        self.tau_th = tau_th
        self.ratio = ratio
        self.min_a = min_a
        self.usage_counts = np.zeros(self.N, dtype=int)
        self.condition = VarianceReductionCondition(tau_th)
        super(ConstantTimeSampler, self).__init__(dataset, reweighting)

    def _select_samples(self, idxs_sorted, scores, batch_size):
        # 两阶段采样
        # 第一阶段：从 B 个样本中选出 K 个高分样本
        K = int(self.K_ratio * self.batch_size)
        L = self.batch_size - K

        # 选择前 K 个高分样本
        high_score_indices = idxs_sorted[:K]

        # 从剩余的 B - K 个样本中随机选取 L 个样本
        remaining_indices = idxs_sorted[K:]
        if len(remaining_indices) >= L:
            random_indices = np.random.choice(remaining_indices, L, replace=False)
        else:
            # 如果剩余样本不足 L，则允许重复采样
            random_indices = np.random.choice(remaining_indices, L, replace=True)

        # 合并选择的样本索引
        final_indices_in_B = np.concatenate([high_score_indices, random_indices])

        # 计算权重以确保无偏性
        # 第一阶段被选中概率： P_first = B / N
        P_first = self.B / self.N

        # 第二阶段被选中概率：
        # 对于前 K 个样本： P_second = 1
        # 对于随机选取的 L 个样本： P_second = L / (B - K)
        P_second = np.zeros(self.batch_size)
        P_second[:K] = 1.0
        P_second[K:] = L / (self.B - K)

        # 总的被选中概率
        P_select = P_first * P_second

        # 目标分布 P_target(x) = 1 / N
        P_target = 1.0 / self.N

        # 重要性权重 w(x) = P_target(x) / P_select(x)
        weights = (P_target / P_select).reshape(-1, 1)

        # 权重归一化，使其平均值为1
        weights /= np.mean(weights)

        return final_indices_in_B, weights

    def _get_samples_with_scores(self, batch_size):
        # 第一阶段：均匀随机采样 B 个样本
        sampled_indices = np.random.choice(self.N, self.B, replace=False)
        x_sampled, y_sampled = self._slice_data(self.dataset._x_train, self.dataset._y_train, sampled_indices)

        # 计算 B 个样本的分数
        scores_sampled = self.model.score(x_sampled, y_sampled, batch_size=self.B)
        scores_sampled = np.nan_to_num(scores_sampled, nan=0.0, posinf=1.0, neginf=0.0)
        scores_sampled += 1e-8  # 避免除零

        return sampled_indices, scores_sampled, (x_sampled, y_sampled)

    def sample(self, batch_size):
        print(f"Entering sample method")
        print(f"Requested batch_size: {batch_size}, self.N: {self.N}, self.B: {self.B}")

        # 确保 batch_size 不超过数据集大小
        if batch_size > self.N:
            batch_size = self.N

        # 获取样本和分数
        idxs1, scores_sampled, xy = self._get_samples_with_scores(batch_size)

        if scores_sampled is not None:
            # 使用 argsort 对分数进行排序（降序）
            sorted_indices = np.argsort(-scores_sampled)  # Descending order
            # 调用 _select_samples 方法获取选中的样本索引和权重
            selected_indices, weights = self._select_samples(sorted_indices, scores_sampled, batch_size)
        else:
            # 如果没有分数，则随机采样
            selected_indices = np.random.choice(len(idxs1), batch_size, replace=False)
            weights = np.ones(len(selected_indices))

        # 获取选中的样本数据
        if xy is None:
            # 如果没有预先加载数据，则从数据集中提取
            x_selected, y_selected = self._slice_data(self.dataset._x_train, self.dataset._y_train, selected_indices)
        else:
            x, y = xy
            x_selected, y_selected = self._slice_data(x, y, selected_indices)

        # 处理分数
        scores_selected = scores_sampled[selected_indices] if scores_sampled is not None else np.ones(batch_size)
        self._send_messages(selected_indices, (x_selected, y_selected), weights, scores_selected)
        return selected_indices, (x_selected, y_selected), weights
上述是samplers.py文件
import numpy as np
import tensorflow as tf

class BaseDataset(object):
    class DataProxy(object):
        def __init__(self, dataset, subset):
            self.data = getattr(dataset, f"_{subset}_data")
            self.size = getattr(dataset, f"_{subset}_size")

        def __getitem__(self, idxs):
            return self.data(idxs)

        def __len__(self):
            return self.size()

    def _train_data(self, idxs=slice(None)):
        """返回训练数据作为 (x, y)"""
        raise NotImplementedError()

    def _train_size(self):
        """返回训练数据的大小"""
        raise NotImplementedError()

    def _test_data(self, idxs=slice(None)):
        """返回测试数据作为 (x, y)"""
        raise NotImplementedError()

    def _test_size(self):
        """返回测试数据的大小"""
        raise NotImplementedError()

    def _slice_data(self, x, y, idxs):
        if isinstance(x, (list, tuple)):
            return [xi[idxs] for xi in x], y[idxs]
        else:
            return x[idxs], y[idxs]

    def _extract_shape(self, x):
        if isinstance(x, (list, tuple)):
            return [xi.shape[1:] for xi in x]
        return x.shape[1:]

    @property
    def train_data(self):
        return self.DataProxy(self, "train")

    @property
    def test_data(self):
        return self.DataProxy(self, "test")

    @property
    def shape(self):
        """返回样本的形状"""
        raise NotImplementedError()

    @property
    def output_size(self):
        """返回输出的大小"""
        raise NotImplementedError()

    @property
    def output_shape(self):
        """返回输出的形状（序列模型可能会有所不同）"""
        return (self.output_size,)

class InMemoryDataset(BaseDataset):
    """一个包含四个numpy数组（x_train, y_train, x_test, y_test）的内存中的数据集。
    x可以是一个numpy数组列表。"""

    def __init__(self, X_train, y_train, X_test, y_test, categorical=True, preprocess_x=None):
        # 如果提供了预处理函数，则应用预处理
        if preprocess_x is not None:
            self._x_train = preprocess_x(X_train)
            self._x_test = preprocess_x(X_test)
        else:
            self._x_train = X_train
            self._x_test = X_test

        # 决定是否将目标转换为one-hot编码
        if categorical:
            if len(y_train.shape) == 1 or y_train.shape[1] == 1:
                # 如果y_train是向量或有一列，则转换为one-hot
                self._y_train = tf.keras.utils.to_categorical(y_train)
                self._y_test = tf.keras.utils.to_categorical(y_test)
            else:
                # 假设y_train已经是one-hot编码的
                self._y_train = y_train
                self._y_test = y_test
            self._output_size = self._y_train.shape[1]
        # 对于稀疏分类任务
        elif issubclass(y_train.dtype.type, np.integer):
            self._y_train = y_train
            self._y_test = y_test
            self._output_size = np.max(self._y_train) + 1  # 假设类别索引从0开始
        # 非分类任务，直接复制标签
        else:
            self._y_train = y_train
            self._y_test = y_test
            if len(self._y_train.shape) == 1:
                self._output_size = 1
            else:
                self._output_size = self._y_train.shape[1]

    def _train_data(self, idxs=slice(None)):
        return self._slice_data(self._x_train, self._y_train, idxs)

    def _train_size(self):
        return self._len(self._x_train)

    def _test_data(self, idxs=slice(None)):
        return self._slice_data(self._x_test, self._y_test, idxs)

    def _test_size(self):
        return self._len(self._x_test)

    def _len(self, x):
        if isinstance(x, (list, tuple)):
            return len(x[0])
        else:
            return len(x)

    @property
    def shape(self):
        return self._extract_shape(self._x_train)

    @property
    def output_size(self):
        return self._output_size

    @classmethod
    def from_loadable(cls, dataset, categorical=True, preprocess_x=None):
        (x_train, y_train), (x_test, y_test) = dataset.load_data()
        return cls(x_train, y_train, x_test, y_test, categorical=categorical, preprocess_x=preprocess_x)

class GeneratorDataset(BaseDataset):
    """一个基于生成器的数据集，能够处理无法完全加载到内存中的大型数据集"""

    def __init__(self, train_gen, test_data, test_size=None):
        self.train_gen = train_gen
        self._x_test, self._y_test = test_data
        self._test_size_value = test_size

        # 假设生成器会生成相同形状的数据
        sample = next(train_gen)
        self._x_sample, self._y_sample = sample

    def _train_data(self, idxs=slice(None)):
        # 生成器无法通过索引获取特定的数据，因此这里返回整个生成器
        return self.train_gen

    def _train_size(self):
        # 生成器的数据大小未知，返回 None
        return None

    def _test_data(self, idxs=slice(None)):
        return self._slice_data(self._x_test, self._y_test, idxs)

    def _test_size(self):
        if self._test_size_value is not None:
            return self._test_size_value
        else:
            if isinstance(self._x_test, (list, tuple)):
                return len(self._x_test[0])
            else:
                return len(self._x_test)

    @property
    def shape(self):
        if isinstance(self._x_sample, (list, tuple)):
            return [xi.shape[1:] for xi in self._x_sample]
        else:
            return self._x_sample.shape[1:]

    @property
    def output_size(self):
        if self._y_sample.ndim == 1:
            return 1
        return self._y_sample.shape[1]
上述是datatsets.py文件
#
# Copyright (c) 2017 Idiap Research Institute, http://www.idiap.ch/
# Written by Angelos Katharopoulos <<PRESIDIO_ANONYMIZED_EMAIL_ADDRESS>>
#

from functools import reduce

from tensorflow.python.keras import backend as K
from tensorflow.python.keras import losses
from tensorflow.python.keras.layers import Layer


def _per_sample_loss(loss_function, mask, x):
    """Compute the per sample loss supporting masking and returning always
    tensor of shape (batch_size, 1)

    Arguments
    ---------
        loss_function: callable
        mask: boolean tensor or None
        x: list/tuple of inputs to the loss
    """
    # Compute the loss
    loss = loss_function(*x)

    # Apply masking if needed
    if mask is not None:
        if not isinstance(mask, list):
            mask = [mask]
        mask = [K.cast(m, K.floatx()) for m in mask if m is not None]
        mask = reduce(lambda a, b: a*b, mask)
        mask_dims = len(K.int_shape(mask))
        mask_mean = K.mean(mask, axis=list(range(1, mask_dims)), keepdims=True)
        loss *= mask
        loss /= mask_mean

    # If the loss has more than 1 dimensions then aggregate the last dimension
    dims = len(K.int_shape(loss))
    if dims > 1:
        loss = K.mean(loss, axis=list(range(1, dims)))

    return K.expand_dims(loss)


class LossLayer(Layer):
    """LossLayer outputs the loss per sample
    
    # Arguments
        loss: The loss function to use to combine the model output and the
              target
    """
    def __init__(self, loss, **kwargs):
        self.supports_masking = True
        self.loss = losses.get(loss)

        super(LossLayer, self).__init__(**kwargs)

    def compute_mask(self, inputs, input_mask):
        return None

    def build(self, input_shape):
        pass # Nothing to do

        super(LossLayer, self).build(input_shape)

    def compute_output_shape(self, input_shape):
        # We need two inputs X and y
        assert len(input_shape) == 2

        # (None, 1) because all losses should be scalar
        return (input_shape[0][0], 1)

    def call(self, x, mask=None):
        return _per_sample_loss(self.loss, mask, x)


class GradientNormLayer(Layer):
    """GradientNormLayer aims to output the gradient norm given a list of
    parameters (whose gradient to compute) and a loss function to combine the
    two inputs.
    
    # Arguments
        parameter_list: A list of Keras variables to compute the gradient
                        norm for
        loss: The loss function to use to combine the model output and the
              target into a scalar and then compute the gradient norm
        fast: If set to True it means we know that the gradient with respect to
              each sample only affects one part of the parameter list so we can
              use the batch mode to compute the gradient
    """
    def __init__(self, parameter_list, loss, fast=False, **kwargs):
        self.supports_masking = True
        self.parameter_list = parameter_list
        self.loss = losses.get(loss)
        self.fast = fast

        super(GradientNormLayer, self).__init__(**kwargs)

    def compute_mask(self, inputs, input_mask):
        return None

    def build(self, input_shape):
        pass # Nothing to do

        super(GradientNormLayer, self).build(input_shape)

    def compute_output_shape(self, input_shape):
        # We get two inputs
        assert len(input_shape) == 2

        return (input_shape[0][0], 1)

    def call(self, x, mask=None):
        # x should be an output and a target
        assert len(x) == 2

        losses = _per_sample_loss(self.loss, mask, x)
        if self.fast:
            grads = K.sqrt(K.sum([
                self._sum_per_sample(K.square(g))
                for g in K.gradients(losses, self.parameter_list)
            ]))
        else:
            nb_samples = K.shape(losses)[0]
            grads = K.map_fn(
                lambda i: self._grad_norm(losses[i]),
                K.arange(0, nb_samples),
                dtype=K.floatx()
            )

        return K.reshape(grads, (-1, 1))

    def _sum_per_sample(self, x):
        """Sum across all the dimensions except the batch dim"""
        # Instead we might be able to use x.ndims but there have been problems
        # with ndims and Keras so I think len(int_shape()) is more reliable
        dims = len(K.int_shape(x))
        return K.sum(x, axis=list(range(1, dims)))

    def _grad_norm(self, loss):
        grads = K.gradients(loss, self.parameter_list)
        return K.sqrt(
            sum([
                K.sum(K.square(g))
                for g in grads
            ])
        )
上述是scores.py文件
#
# Copyright (c) 2017 Idiap Research Institute, http://www.idiap.ch/
# Written by Angelos Katharopoulos <<PRESIDIO_ANONYMIZED_EMAIL_ADDRESS>>
#

from tensorflow.python.keras import backend as K
from tensorflow.python.keras.layers import Layer
from tensorflow.python.keras import initializers

from ..utils.tf_custom import tf


class _BaseNormalization(Layer):
    """Implement utility functions for the normalization layers."""
    def _moments(self, x, axes):
        return (
            K.mean(x, axis=axes, keepdims=True),
            K.var(x, axis=axes, keepdims=True)
        )


class BatchRenormalization(_BaseNormalization):
    """Batch renormalization layer (Sergey Ioffe, 2017).

    # Arguments
        momentum: Momentum for the moving average
        epsilon: Added to variance to avoid divide by 0
        rmax: Maximum correction for the variance
        dmax: Maximum correction for the bias
    """
    def __init__(self, momentum=0.99, epsilon=1e-3, rmax_0=1., rmax_inf=3.,
                 dmax_0=0., dmax_inf=5., rmax_duration=40000,
                 dmax_duration=25000, **kwargs):
        super(BatchRenormalization, self).__init__(**kwargs)

        self.momentum = momentum
        self.epsilon = epsilon
        self.rmax_0 = rmax_0
        self.rmax_inf = rmax_inf
        self.rmax_dur = rmax_duration
        self.dmax_0 = dmax_0
        self.dmax_inf = dmax_inf
        self.dmax_dur = dmax_duration

    def build(self, input_shape):
        dim = input_shape[-1]
        if dim is None:
            raise ValueError(("The normalization axis should have a "
                              "defined dimension"))
        self.dim = dim

        # Trainable part
        self.gamma = self.add_weight(
            shape=(dim,),
            name="gamma",
            initializer=initializers.get("ones")
        )
        self.beta = self.add_weight(
            shape=(dim,),
            name="beta",
            initializer=initializers.get("zeros")
        )

        # Statistics
        self.moving_mean = self.add_weight(
            shape=(dim,),
            name="moving_mean",
            initializer=initializers.get("zeros"),
            trainable=False
        )
        self.moving_sigma = self.add_weight(
            shape=(dim,),
            name="moving_sigma",
            initializer=initializers.get("ones"),
            trainable=False
        )

        # rmax, dmax and steps
        self.steps = self.add_weight(
            shape=tuple(),
            name="steps",
            initializer=initializers.get("zeros"),
            trainable=False
        )
        self.rmax = self.add_weight(
            shape=tuple(),
            name="rmax",
            initializer=initializers.Constant(self.rmax_0),
            trainable=False
        )
        self.dmax = self.add_weight(
            shape=tuple(),
            name="dmax",
            initializer=initializers.Constant(self.dmax_0),
            trainable=False
        )

        self.built = True

    def _moments(self, x):
        axes = range(len(K.int_shape(x))-1)
        if K.backend() == "tensorflow":
            return tf.nn.moments(x, axes)
        else:
            # TODO: Maybe the following can be optimized a bit?
            mean = K.mean(K.reshape(x, (-1, self.dim)), axis=0)
            var = K.var(K.reshape(x, (-1, self.dim)), axis=0)

            return mean, var

    def _clip(self, x, x_min, x_max):
        if K.backend() == "tensorflow":
            return tf.clip_by_value(x, x_min, x_max)
        else:
            return K.maximum(K.minimum(x, x_max), x_min)

    def call(self, inputs, training=None):
        x = inputs
        assert not isinstance(x, list)

        # Compute the minibatch statistics
        mean, var = self._moments(x)
        sigma = K.sqrt(var + self.epsilon)

        # If in training phase set rmax, dmax large so that we use the moving
        # averages to do the normalization
        rmax = K.in_train_phase(self.rmax, K.constant(1e5), training)
        dmax = K.in_train_phase(self.dmax, K.constant(1e5), training)

        # Compute the corrections based on rmax, dmax
        r = K.stop_gradient(self._clip(
            sigma/self.moving_sigma,
            1./rmax,
            rmax
        ))
        d = K.stop_gradient(self._clip(
            (mean - self.moving_mean)/self.moving_sigma,
            -dmax,
            dmax
        ))

        # Actually do the normalization and the rescaling
        xnorm = ((x-mean)/sigma)*r + d
        y = self.gamma * xnorm + self.beta

        # Add the moving average updates
        self.add_update([
            K.moving_average_update(self.moving_mean, mean, self.momentum),
            K.moving_average_update(self.moving_sigma, sigma, self.momentum)
        ], x)

        # Add the r, d updates
        rmax_prog = K.minimum(1., self.steps/self.rmax_dur)
        dmax_prog = K.minimum(1., self.steps/self.dmax_dur)
        self.add_update([
            K.update_add(self.steps, 1),
            K.update(
                self.rmax,
                self.rmax_0 + rmax_prog*(self.rmax_inf-self.rmax_0)
            ),
            K.update(
                self.dmax,
                self.dmax_0 + dmax_prog*(self.dmax_inf-self.dmax_0)
            )
        ])

        # Fix the output's uses learning phase
        y._uses_learning_phase = rmax._uses_learning_phase

        return y


class LayerNormalization(_BaseNormalization):
    """LayerNormalization is a determenistic normalization layer to replace
    BN's stochasticity.

    # Arguments
        axes: list of axes that won't be aggregated over
    """
    def __init__(self, axes=None, bias_axes=[-1], epsilon=1e-3, **kwargs):
        super(LayerNormalization, self).__init__(**kwargs)
        self.axes = axes
        self.bias_axes = bias_axes
        self.epsilon = epsilon

    def build(self, input_shape):
        # Get the number of dimensions and the axes that won't be aggregated
        # over
        ndims = len(input_shape)
        axes = self.axes or []
        bias_axes = self.bias_axes or [-1]

        # Figure out the shape of the statistics
        gamma_shape = [1]*ndims
        beta_shape = [1]*ndims
        for ax in axes:
            gamma_shape[ax] = input_shape[ax]
        for ax in bias_axes:
            beta_shape[ax] = input_shape[ax]

        # Figure out the axes we will aggregate over accounting for negative
        # axes
        self.reduction_axes = [
            ax for ax in range(ndims)
            if ax > 0 and (ax+ndims) % ndims not in axes
        ]

        # Create trainable variables
        self.gamma = self.add_weight(
            shape=gamma_shape,
            name="gamma",
            initializer=initializers.get("ones")
        )
        self.beta = self.add_weight(
            shape=beta_shape,
            name="beta",
            initializer=initializers.get("zeros")
        )

        self.built = True

    def call(self, inputs):
        x = inputs
        assert not isinstance(x, list)

        # Compute the per sample statistics
        mean, var = self._moments(x, self.reduction_axes)
        std = K.sqrt(var + self.epsilon)

        return self.gamma*(x-mean)/std + self.beta


class StatsBatchNorm(_BaseNormalization):
    """Use the accumulated statistics for batch norm instead of computing them
    for each minibatch.

    # Arguments
        momentum: Momentum for the moving average
        epsilon: Added to variance to avoid divide by 0
    """
    def __init__(self, momentum=0.99, epsilon=1e-3, update_stats=False,
                 **kwargs):
        super(StatsBatchNorm, self).__init__(**kwargs)

        self.momentum = momentum
        self.epsilon = epsilon
        self.update_stats = update_stats

    def build(self, input_shape):
        dim = input_shape[-1]
        if dim is None:
            raise ValueError(("The normalization axis should have a "
                              "defined dimension"))
        self.dim = dim

        # Trainable part
        self.gamma = self.add_weight(
            shape=(dim,),
            name="gamma",
            initializer=initializers.get("ones")
        )
        self.beta = self.add_weight(
            shape=(dim,),
            name="beta",
            initializer=initializers.get("zeros")
        )

        # Statistics
        self.moving_mean = self.add_weight(
            shape=(dim,),
            name="moving_mean",
            initializer=initializers.get("zeros"),
            trainable=False
        )
        self.moving_variance = self.add_weight(
            shape=(dim,),
            name="moving_variance",
            initializer=initializers.get("ones"),
            trainable=False
        )

        self.built = True

    def call(self, inputs, training=None):
        x = inputs
        assert not isinstance(x, list)

        # Do the normalization and the rescaling
        xnorm = K.batch_normalization(
            x,
            self.moving_mean,
            self.moving_variance,
            self.beta,
            self.gamma,
            epsilon=self.epsilon
        )

        # Compute and update the minibatch statistics
        if self.update_stats:
            mean, var = self._moments(x, axes=range(len(K.int_shape(x))-1))
            self.add_update([
                K.moving_average_update(self.moving_mean, mean, self.momentum),
                K.moving_average_update(self.moving_variance, var, self.momentum)
            ], x)

        return xnorm


class GroupNormalization(_BaseNormalization):
    """GroupNormalization is an improvement to LayerNormalization presented in
    https://arxiv.org/abs/1803.08494.

    #Arguments
        G: The channels per group in the normalization
        epsilon: Added to the variance to avoid division with 0
    """
    def __init__(self, G=32, epsilon=1e-3, **kwargs):
        super(GroupNormalization, self).__init__(**kwargs)

        self.G = G
        self.epsilon = epsilon

    def build(self, input_shape):
        # Get the number of dimensions and the channel axis
        ndims = len(input_shape)
        channel_axis = -1 if K.image_data_format() == "channels_last" else -3
        shape = [1]*ndims
        shape[channel_axis] = input_shape[channel_axis]

        # Make sure everything is in order
        assert None not in shape, "The channel axis must be defined"
        assert shape[channel_axis] % self.G == 0, ("The channels must be "
                                                   "divisible by the number "
                                                   "of groups")

        # Create the trainable weights
        self.gamma = self.add_weight(
            shape=shape,
            name="gamma",
            initializer=initializers.get("ones")
        )
        self.beta = self.add_weight(
            shape=shape,
            name="beta",
            initializer=initializers.get("zeros")
        )

        # Keep the channel axis for later use
        self.channel_axis = channel_axis
        self.ndims = ndims

        # We 're done
        self.built = True

    def call(self, inputs):
        x = inputs
        assert not isinstance(x, list)

        # Get the shapes
        G = self.G
        channel_axis = self.channel_axis
        ndims = self.ndims
        original_shape = K.shape(x)
        shape = [
            original_shape[i]
            for i in range(ndims)
        ]
        if channel_axis == -1:
            shape[channel_axis] //= G
            shape.append(G)
            axes = sorted([
                ndims + channel_axis - i
                for i in range(3)
            ])
        else:
            shape[channel_axis] //= G
            shape.insert(channel_axis, G)
            axes = sorted([
                ndims + channel_axis + i
                for i in range(3)
            ])

        # Do the group norm
        x = K.reshape(x, shape)
        mean, var = self._moments(x, axes)
        std = K.sqrt(var + self.epsilon)
        x = (x - mean)/std
        x = K.reshape(x, original_shape)

        return self.gamma*x + self.beta
上述是normalization.py文件
#
# Copyright (c) 2017 Idiap Research Institute, http://www.idiap.ch/
# Written by Angelos Katharopoulos <<PRESIDIO_ANONYMIZED_EMAIL_ADDRESS>>
#

from functools import reduce

from tensorflow.python.keras import backend as K
from tensorflow.python.keras.layers import Layer
from tensorflow.python.keras.metrics import categorical_accuracy, binary_accuracy, \
    get as get_metric, sparse_categorical_accuracy

from ..utils.functional import compose


class MetricLayer(Layer):
    """Create a layer that computes a metric taking into account masks."""
    def __init__(self, metric_func, **kwargs):
        self.supports_masking = True
        self.metric_func = metric_func

        super(MetricLayer, self).__init__(**kwargs)

    def compute_mask(self, inputs, input_mask):
        return None

    def build(self, input_shape):
        # Special care for accuracy because keras treats it specially
        try:
            if "acc" in self.metric_func:
                self.metric_func = self._generic_accuracy
        except TypeError:
            pass # metric_func is not a string
        self.metric_func = compose(K.expand_dims, get_metric(self.metric_func))

        super(MetricLayer, self).build(input_shape)

    def compute_output_shape(self, input_shape):
        # We need two inputs y_true, y_pred
        assert len(input_shape) == 2

        return (input_shape[0][0], 1)

    def call(self, inputs, mask=None):
        # Compute the metric
        metric = self.metric_func(*inputs)
        if K.int_shape(metric)[-1] == 1:
            metric = K.squeeze(metric, axis=-1)
        if len(K.int_shape(metric)) == 0:
            metric = K.ones(K.shape(inputs[0])[0]) * metric

        # Apply the mask if needed
        if mask is not None:
            if not isinstance(mask, list):
                mask = [mask]
            mask = [K.cast(m, K.floatx()) for m in mask if m is not None]
            mask = reduce(lambda a, b: a*b, mask)
            metric *= mask
            metric /= K.mean(mask, axis=-1, keepdims=True)

        # Make sure that the tensor returned is (None, 1)
        dims = len(K.int_shape(metric))
        if dims > 1:
            metric = K.mean(metric, axis=list(range(1, dims)))

        return K.expand_dims(metric)

    @staticmethod
    def _generic_accuracy(y_true, y_pred):
        if K.int_shape(y_pred)[1] == 1:
            return binary_accuracy(y_true, y_pred)
        if K.int_shape(y_true)[-1] == 1:
            return sparse_categorical_accuracy(y_true, y_pred)

        return categorical_accuracy(y_true, y_pred)


class TripletLossLayer(Layer):
    """A bit of an unorthodox layer that implements the triplet loss with L2
    normalization.

    It receives 1 vector which is the concatenation of the three
    representations and performs the following operations.
    x = concat(x_a, x_p, x_n)
    N = x.shape[1]/3
    return ||x[:N] - x[2*N:]||_2^2 - ||x[:N] - x[N:2*N]||_2^2
    """
    def __init__(self, **kwargs):
        super(TripletLossLayer, self).__init__(**kwargs)

    def build(self, input_shape):
        assert not isinstance(input_shape, list)
        self.N = input_shape[1] // 3
        self.built = True

    def compute_output_shape(self, input_shape):
        assert not isinstance(input_shape, list)
        return (input_shape[0], 1)

    def call(self, x):
        N = self.N

        xa = x[:, :N]
        xp = x[:, N:2*N]
        xn = x[:, 2*N:]

        xa = xa / K.sqrt(K.sum(xa**2, axis=1, keepdims=True))
        xp = xp / K.sqrt(K.sum(xp**2, axis=1, keepdims=True))
        xn = xn / K.sqrt(K.sum(xn**2, axis=1, keepdims=True))

        dn = K.sum(K.square(xa - xn), axis=1, keepdims=True)
        dp = K.sum(K.square(xa - xp), axis=1, keepdims=True)

        return dn - dp
上述是metrics.py文件
absl-py                      2.1.0
appdirs                      1.4.4
astunparse                   1.6.3
autograd                     1.7.0
blinker                      1.8.2
cachetools                   5.5.0
certifi                      2024.8.30
charset-normalizer           3.3.2
contourpy                    1.3.0
cycler                       0.12.1
dill                         0.3.9
docplex                      2.28.240
flatbuffers                  24.3.25
fonttools                    4.54.1
gast                         0.4.0
google-auth                  2.35.0
google-auth-oauthlib         0.4.6
google-pasta                 0.2.0
grpcio                       1.66.2
h5py                         3.12.1
idna                         3.10
importlib_metadata           8.5.0
importlib_resources          6.4.5
joblib                       1.4.2
keras                        2.10.0
Keras-Preprocessing          1.1.2
kiwisolver                   1.4.7
libclang                     18.1.1
Markdown                     3.7
MarkupSafe                   2.1.5
matplotlib                   3.9.2
mpmath                       1.3.0
networkx                     3.2.1
numpy                        1.26.4
oauthlib                     3.2.2
opt_einsum                   3.4.0
packaging                    24.1
pbr                          6.1.0
PennyLane                    0.15.0
pillow                       10.4.0
pip                          24.2
protobuf                     3.19.6
psutil                       6.0.0
pyasn1                       0.6.1
pyasn1_modules               0.4.1
pydot                        3.0.2
pydotplus                    2.0.2
pyparsing                    3.1.4
python-dateutil              2.9.0.post0
qiskit                       1.0.2
qiskit-aer                   0.14.0.1
qiskit-algorithms            0.3.0
qiskit-optimization          0.6.1
requests                     2.32.3
requests-oauthlib            2.0.0
rsa                          4.9
rustworkx                    0.15.1
scikit-learn                 1.5.2
scipy                        1.13.1
semantic-version             2.6.0
setuptools                   75.1.0
six                          1.16.0
stevedore                    5.3.0
symengine                    0.11.0
sympy                        1.13.3
tensorboard                  2.10.1
tensorboard-data-server      0.6.1
tensorboard-plugin-wit       1.8.1
tensorflow                   2.10.0
tensorflow-estimator         2.10.0
tensorflow-io-gcs-filesystem 0.31.0
termcolor                    2.4.0
threadpoolctl                3.5.0
toml                         0.10.2
typing_extensions            4.12.2
urllib3                      2.2.3
Werkzeug                     3.0.4
wheel                        0.44.0
wrapt                        1.16.0
zipp                         3.20.2
上述是对应的运行环境如果我想要按照你刚才所说的Katharopoulos, A., & Fleuret, F. (2018). Not All Samples Are Created Equal: Deep Learning with Importance Sampling. International Conference on Machine Learning (ICML).这篇文献将损失修改为梯度模的近似来作为得分的评价，该如何修改代码，应该修改哪些部分？给出修改部分的对应完整代码