如何在Keras中进行分片赋值(使用TensorFlow作为后端)

import tensorflow as tf
from tensorflow import keras

model = keras.Sequential()
model.add(keras.layers.Dense(units=64, activation='relu', input_shape=(input_dim,)))
model.add(keras.layers.Dense(units=10, activation='softmax'))

model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

class ShardAssignmentCallback(keras.callbacks.Callback):
    def __init__(self, shard_values):
        super(ShardAssignmentCallback, self).__init__()
        self.shard_values = shard_values

    def on_train_batch_begin(self, batch, logs=None):
        shard_index = batch % len(self.shard_values)
        shard_value = self.shard_values[shard_index]
        tf.keras.backend.set_value(self.model.optimizer.lr, shard_value)

shard_values = [0.001, 0.01, 0.1]  # 分片赋值的值列表
shard_assignment_callback = ShardAssignmentCallback(shard_values)

model.fit(x_train, y_train, batch_size=32, epochs=10, callbacks=[shard_assignment_callback])