[Keras] 현재의 Learning Rate를 확인하는 Custom Callback

Adagrad 이후의 Adam등에서는 Adaptive Learning Rate가 적용되면서 간단하게 하나의 Learning rate 값이 아닌 Weight에 따라 다른 속도로 decay되는 처리가 이루어진다.

 

하지만 초기에 준 inital value에 해당하는 Learning rate를 기반으로 전체적으로 적용되는 Learning rate가 있고, 여기에 이전 epoch들에서의 gradient의 변화 정도 등을 고려한 처리가 이루어지기 때문에 현재 epoch에서의 일종의 decayed learning rate가 존재한다고 봐도 된다.

 

예전의 Keras의 경우엔 optimizer.lr 등으로 현재의 learning rate를 구해왔으나, decay factor에 의해 epoch마다 값이 달라진 learning rate의 값은 optimizer._delayed_lr(tf.float32) 등을 사용해야 한다.

optimizer.lr은 초기에 할당한 constant 로 값이 epoch가 아무리 진행되어도 변하지 않는다.

ref: _decayed_lr 에 대한 definition

 https://github.com/tensorflow/tensorflow/blob/e5bf8de410005de06a7ff5393fafdf832ef1d4ad/tensorflow/python/keras/optimizer_v2/optimizer_v2.py#L679

---

현재 Learnint Rate를 출력하는 Custom Callback

다음의 class는 Custom Callback으로 epoch가 종료될 때, optimizer로부터 현재의 decayed learning rate를 구한다.

class LRLoggingCb(tf.keras.callbacks.Callback):
  def on_epoch_end(self, epoch, logs=None):
    init_lr = self.model.optimizer.lr
    current_lr = self.model.optimizer._decayed_lr(tf.float32)
    tf.summary.scalar('init_lr', data=init_lr, step=epoch)
    tf.summary.scalar('current_lr', data=current_lr, step=epoch)
    init_lr = tf.keras.backend.get_value(init_lr)
    print(f'\nEpoch {epoch+1}: Current learning rate is {current_lr:.4f}',
          '/ Init learning rate is {init_lr:.4f}.'

---

Model Save 와 Load시 주의할 점

주의할 것은 만약 training을 이어서 해야하는 경우, training이 다시 시작하는 시점에서 이어서 하고자 하는 epoch의 decayed learning rate를 따로 저장해두고 이를 optimizer.lr에 할당하고 fit을 호출해야한다는 점이다.

아래 그림은 이어서 훈련시키고자 하는 이전 decayed learning rate를 고려하지 않은 경우의 loss에 대한 learning curve이다.

• 이전에 10 epoch까지 kears.optimizers.Adam으로 훈련하고, 이를 tf2 format으로 저장한 이후 저장된 파일로부터 model architecture와 weight 등을 load하고 다시 10 epochs를 훈련했다. (즉, 총 epoch는 20임).
• 각 10epoch의 history 의 객체를 얻어 이를 concatenate시켜 20개의 loss 값들로 구성된 learning curve를 만들었다.

10 epoch까지 잘 내려가던 loss가 validation set(orange)의 경우와 training set(blue)의 경우 모두 11과 12 epoch에서 loss 가 증가하는 것을 확인할 수 있다. 이는 Learning rate가 이어져서 처리되지 못해 발생하는 현상이다.

---

기존의 10 epochs가 종료되는 시점의 decayed learning rate를 저장하고 이를 optimizer.lr 로 할당하고 나서 위의 처리를 수행하면 다음과 같은 learning curve를 얻는다.

위의 그림들은 tf.keras.Sequential을 통해 만든 모델을 이용한터라 decayed learning rate를 따로 변수로 저장하여 처리했다. (subclassing을 이용하여 만든다면 훨씬 편하게 처리가 가능하다.)

---

참고하면 좋은 자료들

https://stackoverflow.com/questions/58149839/learning-rate-of-custom-training-loop-for-tensorflow-2-0/61553047#61553047

Learning rate of custom training loop for tensorflow 2.0

 

https://github.com/tensorflow/models/issues/7687#issuecomment-688285977

How learning rate scheduler works with Custom training loop using tf.GradientTape() · Issue #7687 · tensorflow/models

https://gist.github.com/dsaint31x/a40aa2ef7986a009d4dd2aefd8b1c34b

ml_lr_save_load.ipynb