Discriminative Learning Rates for Keras - tensorflow

Question

I wish to apply different learning rates for each layer (as is done in Fastai) for Keras. All I have found in coming close to this is by modifying this line self.optimizer.apply_gradients(zip(gradients, trainable_vars)) in the keras code block found here by multiplying the gradients with its corresponding learning rate (and set the global learning rate to 1).
However, this method would only work with SGD and other simple optimizers as things such as momentum would distort this simple multiplication of the gradient.
Here is ideally how I would like to implement it:
# 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"],
)
Any thoughts on how I could implement this. Another idea is to have the same number of optimizers as learning rates, and to update them in the train_step function inside the custom model.

Answer 1

This is possible using tensorflow addons now.
pip install tensorflow-addons
import tensorflow_addons as tfa
optimizers = [Adam(learning_rate=1e-2),
Adam(learning_rate=1e-1)]
#specifying the optimizers and layers in which it will operate
optimizers_and_layers = [
(optimizers[0], model.layers[0]),
(optimizers[1], model.layers[1]),
]
# Using Multi Optimizer from Tensorflow Addons
opt = tfa.optimizers.MultiOptimizer(optimizers_and_layers)
model.compile(
optimizer=opt,
loss="mean_squared_error",
metrics=["mean_absolute_error"],
)