从keras dropout层提取dropout蒙版？

import tensorflow as tf
from tensorflow.keras.layers import Layer

class CustomDropout(Layer):
    def __init__(self, rate, **kwargs):
        super(CustomDropout, self).__init__(**kwargs)
        self.rate = rate

    def call(self, inputs, training=None):
        if training:
            mask = tf.random.uniform(shape=tf.shape(inputs)) >= self.rate
            return inputs * tf.cast(mask, dtype=inputs.dtype), mask
        return inputs

# 使用自定义Dropout层
dropout_layer = CustomDropout(rate=0.5)
outputs, dropout_mask = dropout_layer(inputs, training=True)

class DropoutMaskCallback(tf.keras.callbacks.Callback):
    def __init__(self, dropout_layer):
        super(DropoutMaskCallback, self).__init__()
        self.dropout_layer = dropout_layer

    def on_train_batch_end(self, batch, logs=None):
        mask = self.dropout_layer.mask.numpy()
        # 处理蒙印，例如保存或分析

# 假设dropout_layer是模型中的一个Dropout层
callback = DropoutMaskCallback(dropout_layer)
model.fit(x_train, y_train, callbacks=[callback])