Discriminative Learning Rates for Keras

user_2530674 0
Viewed 526

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.

keras tensorflow
1 Answers

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"],
)
user_56451650
Related