Layer-specific learning rate in Keras Model

Question

In a keras model, It's possible to set the learning rate for the model when compiling, like this,

model.compile(optimizer=Adam(learning_rate=0.001), loss=...)

Answer 1

You can use tfa.optimizers.MultiOptimizer<\/code> from the tensorflow_addons<\/code> package.

import tensorflow as tf
import tensorflow_addons as tfa

model = tf.keras.Sequential([
    tf.keras.Input(shape=(4,)),
    tf.keras.layers.Dense(8),
    tf.keras.layers.Dense(16),
    tf.keras.layers.Dense(32),
])
optimizers = [
    tf.keras.optimizers.Adam(learning_rate=1e-4),
    tf.keras.optimizers.Adam(learning_rate=1e-2)
]
optimizers_and_layers = [(optimizers[0], model.layers[0]), (optimizers[1], model.layers[1:])]
optimizer = tfa.optimizers.MultiOptimizer(optimizers_and_layers)
model.compile(optimizer=optimizer, loss="mse")