使用tensorflow-addons,第一层学习率设置为1e-4,后面3层为1e-2
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")
详细信息可以去看tensorflow-addons文档或者源码