如何在keras自定义回调中访问tf.data.Dataset？

import tensorflow as tf
from tensorflow import keras

class CustomCallback(keras.callbacks.Callback):
    def __init__(self, dataset):
        super(CustomCallback, self).__init__()
        self.dataset = dataset

    def on_train_begin(self, logs=None):
        # 在训练开始时调用
        # 可以在这里访问self.dataset

    def on_epoch_end(self, epoch, logs=None):
        # 在每个epoch结束时调用
        # 可以在这里访问self.dataset

# 创建自定义回调对象
dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
callback = CustomCallback(dataset)

# 使用自定义回调进行训练
model.fit(dataset, epochs=10, callbacks=[callback])