Save & load for scheduler

[graykode]

tensorflow uses optimizer and scheduler together to save scheduler when saving optimizer, but pytorch is independent. For example, tensorflow has a scheduler when calling an optimizer as follows.

import tensorflow as tf

if __name__ == '__main__':
    lr_schedule = tf.keras.optimizers.schedules.PolynomialDecay(
        initial_learning_rate=0.0,
        decay_steps=100,
        end_learning_rate=0.001,
    )
    optimizer = tf.keras.optimizers.Adam(
        learning_rate=lr_schedule, beta_1=0.99, beta_2=0.1, epsilon=1e-6
    )
    print(lr_schedule.get_config())
    print(optimizer.get_config())

{'initial_learning_rate': 0.0, 'decay_steps': 100, 'end_learning_rate': 0.001, 'power': 1.0, 'cycle': False, 'name': None}
{'name': 'Adam', 'learning_rate': {'class_name': 'PolynomialDecay', 'config': {'initial_learning_rate': 0.0, 'decay_steps': 100, 'end_learning_rate': 0.001, 'power': 1.0, 'cycle': False, 'name': None}}, 'decay': 0.0, 'beta_1': 0.99, 'beta_2': 0.1, 'epsilon': 1e-06, 'amsgrad': False}

To implement this in pytorch, modify the existing optimizer save and load code to be used as follows.

>>> optimizer_manager.save(optimizer, scheduler)
>>> optimizer_manager.load_with_scheduler(optimizer, scheduler, step=938)