Layer-specific learning rate in Keras Model

user_15760624 0
Viewed 897

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=...)

This sets the same learning rate for all the layers in the model, but how do I set different learning rates for each layer of my model?
like this,
layer 1 : 0.001
layer 2 : 0.05
layer 3 : 0.02
etc.
How do I do this in keras? or in tf.keras?

keras tensorflow machine-learning python
1 Answers

You can use tfa.optimizers.MultiOptimizer from the tensorflow_addons package.

See directly from the docs:

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")

Note "Each optimizer will optimize only the weights associated with its paired layer."

user_95553880
Related