Scheduler should be initialized after Apex is configured

[alexandrebone]

Description

Using jointly a scheduler and apex-amp throws a UserWarning:

Seems like optimizer.step() has been overridden after learning rate scheduler initialization. Please, make sure to call optimizer.step() before lr_scheduler.step().

As suggested here (https://discuss.pytorch.org/t/cyclic-learning-rate-how-to-use/53796), the scheduler should be initialized after Apex is configured.

To reproduce

Run the MNIST example with use_amp=True.

Expected behavior

No scheduler-related user warning.

Workaround

The following hack seems to work:

    def on_train_start(self):
        self.trainer.lr_schedulers = [optim.lr_scheduler.CosineAnnealingLR(self.trainer.optimizers[0], T_max=10)]

[Borda]

Good catch, @alexandrebone mind you send a PR with the fix? 🤖

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

[dbpprt]

Any news on this? This bug is still active.

[Borda]

@SkafteNicki pls ^^

[SkafteNicki]

I will take a look at this. It will not be easy since we construct optimizers and schedulers at the same time currently. But maybe there is a workaround.

[SkafteNicki]

The problem seems to be that the scheduler equips the optimizer.step() method
with a number of fields, that are lost after the optimizer is passed to apex. These
can be reconstructed by calling init method of base scheduler class (which all
schedulers inherence from). Here is working example ('_with_counter' is one field
set by the scheduler) that prevents the UserWarning:

import torch

optimizer = torch.optim.Adam(torch.nn.Linear(10,10).parameters(), lr=0.1)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, 10, 0.1)

print(hasattr(scheduler.optimizer.step, '_with_counter'))  # >>> True

# This mimics what apex is doing internally
old_step = optimizer.step
def new_step(self, closure=None):
    retval = old_step()
    return retval        
optimizer.step = new_step

print(hasattr(scheduler.optimizer.step, '_with_counter'))  # >>> False

# Call the base lr scheduler class to re-equip the optimizer with 
# fields important for the scheduler 
scheduler.__class__.__mro__[-2].__init__(scheduler, optimizer)

print(hasattr(scheduler.optimizer.step, '_with_counter'))  # >>> True

optimizer.step()
scheduler.step()

@Borda is this fine?

[Borda]

it looks reasonable to me...
cc: @PyTorchLightning/core-contributors