区分Keras的学习率

Question

我希望为Keras的每一层应用不同的学习率（在Fastai中也是如此）。我发现，要达到这一目的，只需在keras code block found here中修改此行self.optimizer.apply_gradients(zip(gradients, trainable_vars))，方法是将梯度乘以相应的学习率（并将全局学习率设置为1）。

但是，此方法仅适用于SGD和其他简单的优化器，因为动量之类的东西会使这种简单的梯度乘法变得不正确。

理想情况下，这是我想要实现的方式：

# example data
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train = x_train.reshape(-1, 784).astype("float32") / 255.0
x_test = x_test.reshape(-1, 784).astype("float32") / 255.0

lrs = [1e-2, 1e-1]
model = tf.keras.Sequential()
model.add(tf.keras.layers.Dense(20, input_dim=784, activation='relu'))
model.add(tf.keras.layers.Dense(1))
model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=lrs), # DOESNT WORK
    loss="mean_squared_error",
    metrics=["mean_absolute_error"],
)

关于如何实现此目标的任何想法。另一个想法是使优化程序的数量与学习率相同，并在自定义模型的train_step函数中对其进行更新。