Keras에서 학습률 스케줄러를 사용하는 방법

from wandb.keras import WandbCallback

# Model Definition
model = ...

# Compile the model
model.compile(optimizer = <>, loss = <>, metrics = <>)

# Custom Scheduler Function
lr_start   = 1e-4
lr_max     = 0.000015 * replicas * batch_size
lr_min     = 1e-7
r_ramp_ep = 3
lr_sus_ep  = 0
lr_decay   = 0.7
   
def lrfn(epoch):
    if epoch < lr_ramp_ep:
        lr = (lr_max - lr_start) / lr_ramp_ep * epoch + lr_start
            
    elif epoch < lr_ramp_ep + lr_sus_ep:
        lr = lr_max
            
    else:
        lr = (lr_max - lr_min) * lr_decay**(epoch - lr_ramp_ep - lr_sus_ep) + lr_min
            
    return lr

# Using this Custom Function, create a Callback
lr_callback = tf.keras.callbacks.LearningRateScheduler(lrfn, verbose=True)

# Train the Model
model.fit(..., callbacks=[lr_callback, WandbCallback()])