Learning rate and weight decay schedule in Tensorflow SGDW optimizer

user_13547857 0
Viewed 1931

I'm trying to reproduce part of this paper with TensorFlow, the problem is that the authors use SGD with weight decay, cutting the learning rate to 1/10 every 30 epochs.

TensorFlow documentation says that

when applying a decay to the learning rate, be sure to manually apply the decay to the weight_decay as well

So I tried with

schedule = keras.optimizers.schedules.ExponentialDecay(initial_learning_rate = 0.003,
                                                   decay_rate = 0.1,
                                                   decay_steps = steps_per_epoch*30,
                                                   staircase = True
)
optimizer = tfa.optimizers.SGDW(learning_rate = schedule,
                            weight_decay = schedule,
                            momentum = 0.9
)

(steps_per_epoch previously initialized)
As I would with Keras SGD, this however doesn't work and raises a "TypeError: Expected float32" for the decay_weight parameter. What's the correct way to achieve the target behaviour?

hyperparameters keras tensorflow machine-learning python
1 Answers

You are getting an error because you are using keras ExponentialDecay inside tensorflow add-on optimizer SGDW.

As per the paper hyper-parameters are

  1. weight decay of 0.001
  2. momentum of 0.9
  3. starting learning rate is 0.003 which is reduced by a factor of 10 after 30 epochs

So why not use LearningRateScheduler to reduce it by factor to 10 evey 30 epochs

Sample Code

model = tf.keras.Sequential([
    tf.keras.layers.Dense(8, input_shape=(10,)),
    tf.keras.layers.Dense(4, activation='relu'),
    tf.keras.layers.Dense(3, activation='softmax'),
])

X = np.random.randn(10,10)
y = np.random.randint(0,4,(10,3))

model.compile(
    optimizer=tfa.optimizers.SGDW(
        weight_decay=0.001,
        momentum=0.9,
        learning_rate=0.003),
      loss=tf.keras.losses.categorical_crossentropy)

def scheduler(epoch, lr):
  if epoch  % 30 == 0:
    lr = lr*0.1  
  return lr

callback = tf.keras.callbacks.LearningRateScheduler(scheduler)
model.fit(X, y, callbacks=[callback], epochs=100)
user_4239260
Related