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)
上述代码是datasets.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.original_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(x, training=True)
            loss_value = self.model.compiled_loss(
                y,
                y_pred,
                sample_weight=w,
                regularization_losses=self.model.losses
            )

        gradients = tape.gradient(loss_value, self.model.trainable_variables)
        self.model.optimizer.apply_gradients(zip(gradients, self.model.trainable_variables))

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

        # 返回损失值和指标结果
        return loss_value.numpy(), metrics_results, None

    def evaluate(self, test_data):
        """评估模型在测试集上的性能。"""
        x_test, y_test = test_data[:]
        y_pred = self.model.predict(x_test, batch_size=128)
        loss_value = self.model.compiled_loss(y_test, y_pred, regularization_losses=self.model.losses)

        # 更新指标
        self.model.compiled_metrics.update_state(y_test, y_pred)
        metrics_results = [m.result().numpy() for m in self.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_pass_multiplier=self._backward_pass_multiplier,
            extra_samples=self._extra_samples
        )
上述是training.py文件
import sys
import numpy as np
from blinker import signal
from tensorflow.python.keras import backend as K
from tensorflow.python.keras.layers import Layer, multiply
from tensorflow.python.keras import Model, Input
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):
            yield [xi[s:s+bs] for xi in _tolist(x)], 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]

        try:
            # 获取倒数第一层或倒数第二层，取决于最后一层是否有可训练权重
            skip_one = not bool(model.layers[-1].trainable_weights)
            last_layer = -2 if skip_one else -1

            # 如果最后一层有可训练权重，说明无法自动提取预激活张量，需要警告
            if last_layer == -1:
                config = model.layers[-1].get_config()
                if config.get("activation", "linear") != "linear":
                    sys.stderr.write(self.FUSED_ACTIVATION_WARNING)

            return model.layers[last_layer]
        except:
            # 如果出错，可能是因为我们没有使用 gnorm importance
            return None

    def _augment_model(self, model, score, reweighting):
        # 提取模型的一些信息
        loss_fn = model.loss
        optimizer = model.optimizer.__class__.from_config(model.optimizer.get_config())
        output_shape = model.output_shape[1:]
        if isinstance(loss_fn, str) and loss_fn.startswith("sparse"):
            output_shape = output_shape[:-1] + (1,)

        # 确保一些内容看起来正常
        assert not isinstance(loss_fn, list)

        # 我们需要创建两个额外的输入：
        # 1. 真实标签 y_true
        # 2. 预测的得分 pred_score
        y_true = Input(shape=output_shape)
        pred_score = Input(shape=(reweighting.weight_size,))

        # 创建损失层和得分层
        loss_tensor = LossLayer(loss_fn)([y_true, model.output])
        score_tensor = _get_scoring_layer(
            score,
            y_true,
            model.output,
            loss_fn,
            self.layer,
            model
        )

        # 创建样本权重
        weights = reweighting.weight_layer()([score_tensor, pred_score])

        # 创建加权损失
        weighted_loss = multiply([loss_tensor, weights])
        for l in model.losses:
            weighted_loss += l
        weighted_loss_mean = K.mean(weighted_loss)

        # 创建指标层
        metrics = model.metrics or []
        metrics_tensors = [
            MetricLayer(metric)([y_true, model.output])
            for metric in metrics
        ]

        # 创建一个新模型，包括额外的输入和加权损失
        inputs = _tolist(model.input) + [y_true, pred_score]
        outputs = [weighted_loss]
        self.model = Model(inputs=inputs, outputs=outputs)

        # 编译新模型
        self.model.compile(
            optimizer=optimizer,
            loss=lambda yt, yp: weighted_loss_mean,  # 自定义损失
            metrics=metrics
        )

        # 保存原始模型的优化器和损失函数
        self.optimizer = optimizer
        self.loss_fn = loss_fn

    def evaluate_batch(self, x, y):
        # 在评估时，我们不需要样本权重，因此创建一个全为 1 的权重向量
        if len(y.shape) == 1:
            y = np.expand_dims(y, axis=1)
        batch_size = y.shape[0]
        dummy_weights = np.ones((batch_size, self.reweighting.weight_size))
        inputs = _tolist(x) + [y, dummy_weights]
        outputs = self.model.evaluate(
            x=inputs,
            y=np.zeros((batch_size,)),  # 占位符，因为我们的模型输出是损失
            verbose=0,
            return_dict=True
        )

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

        # 返回损失和其他指标
        return [outputs['loss']] + [outputs[m] for m in self.model.metrics_names[1:]]

    def score_batch(self, x, y):
        # 计算得分
        if len(y.shape) == 1:
            y = np.expand_dims(y, axis=1)
        batch_size = y.shape[0]
        dummy_weights = np.ones((batch_size, self.reweighting.weight_size))
        inputs = _tolist(x) + [y, dummy_weights]
        # 获取得分张量的输出
        get_score = K.function(
            inputs=self.model.inputs,
            outputs=self.model.get_layer('score_layer').output
        )
        scores = get_score(inputs)
        return scores.ravel()

    def train_batch(self, x, y, w):
        # 训练单个批次
        if len(y.shape) == 1:
            y = np.expand_dims(y, axis=1)
        batch_size = y.shape[0]
        inputs = _tolist(x) + [y, w]
        outputs = self.model.train_on_batch(
            x=inputs,
            y=np.zeros((batch_size,))  # 占位符，因为我们的模型输出是损失
        )

        # 获取损失和指标
        if isinstance(outputs, list):
            loss = outputs[0]
            metrics = outputs[1:]
        else:
            loss = outputs
            metrics = []

        # 暂时无法获取得分，因此返回 None
        scores = None

        # 返回损失、指标和得分
        result = (
            loss,
            metrics,
            scores
        )
        signal("is.training").send(result)

        return result

def _tolist(x, acceptable_iterables=(list, tuple)):
    if not isinstance(x, acceptable_iterables):
        x = [x]
    return x

def _get_scoring_layer(score, y_true, y_pred, loss_fn="categorical_crossentropy",
                       layer=None, model=None):
    """获取一个计算每个 y_true 和 y_pred 对的得分的层"""
    assert score in ["loss", "gnorm", "full_gnorm", "acc"]

    if score == "loss":
        return LossLayer(loss_fn, name='score_layer')([
            y_true,
            y_pred
        ])
    elif score == "gnorm":
        return GradientNormLayer(
            layer.output,
            loss_fn,
            fast=False,
            name='score_layer'
        )([
            y_true,
            y_pred
        ])
    elif score == "full_gnorm":
        return GradientNormLayer(
            model.trainable_weights,
            loss_fn,
            fast=False,
            name='score_layer'
        )([
            y_true,
            y_pred
        ])
    elif score == "acc":
        return MetricLayer(categorical_accuracy, name='score_layer')([
            y_true,
            y_pred
        ])

class ModelWrapperDecorator(ModelWrapper):
    def __init__(self, model_wrapper, implemented_attributes=set()):
        self.model_wrapper = model_wrapper
        self.implemented_attributes = (
            implemented_attributes | set(["model_wrapper"])
        )

    def __getattribute__(self, name):
        _getattr = object.__getattribute__
        implemented_attributes = _getattr(self, "implemented_attributes")
        if name in implemented_attributes:
            return _getattr(self, name)
        else:
            model_wrapper = _getattr(self, "model_wrapper")
            return getattr(model_wrapper, name)
上述是model_wrappers.py文件，你需要仔细阅读上述我给你的几个文件中的代码，主要理清InMemoryDataset是如何使用的，然后检查所有数据处理的接口是否有问题，如果有问题需要进行修改并且给出对应的代码，注意所有修改都不能变动代码的原文逻辑，主要修改的地方在于原本的代码是经典神经网络，而现在将经典神经网络修改为了量子神经网络，所以适配出现了问题。