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函数中更新它们。

Answer 1

现在使用tensorflow插件可以做到这一点。

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