如何在训练自动编码器(回调)时将keras中的输入随机设置为零？

from tensorflow import keras
import numpy as np

class RandomZeroCallback(keras.callbacks.Callback):
    def __init__(self, zero_ratio):
        super(RandomZeroCallback, self).__init__()
        self.zero_ratio = zero_ratio

    def on_batch_begin(self, batch, logs=None):
        batch_size = len(self.model.inputs)
        input_shape = self.model.inputs[0].shape[1:]  # 获取输入的形状
        mask = np.random.choice([0, 1], size=(batch_size,) + input_shape, p=[self.zero_ratio, 1-self.zero_ratio])
        # 将输入随机设置为零
        for i in range(len(self.model.inputs)):
            self.model.inputs[i] = self.model.inputs[i] * mask[i]

from tensorflow import keras

# 创建自动编码器模型
autoencoder = keras.models.Sequential([...])

# 编译模型
autoencoder.compile(optimizer='adam', loss='mse')

# 创建回调函数实例
zero_callback = RandomZeroCallback(zero_ratio=0.5)

# 训练模型，并传入回调函数
autoencoder.fit(x_train, x_train, epochs=10, batch_size=32, callbacks=[zero_callback])