PlateauLRScheduler resume

[ademyanchuk]

Hi, one more question. If I resume training and use PlateauLRScheduler, I should save and load its state dict, right?

[rwightman]

@ademyanchuk yeah, the plateau scheduler is a bit of a special case, I don't handle it properly in my train script but it should be saved and restored to work on resume.

I tend to use plateau on 'smaller' transfer learning tasks w/ adam or other adaptive opts where resume is less of a concern so I've lived with it (for now).

[ademyanchuk]

Yes. I understand, thank you. Using it for kaggle now, doing some experiments. And need to work on a bigger resolution now. So 24 hours is not enough. Thanks though, your code is enjoyed to read, use and learn from.
Cheers!

[ademyanchuk]

It was not that simple to figure out and make it work [the resuming part], so I put it here for future-self or anyone who will be interested of doing the same.

1. As Ross mentioned, you need to save scheduler state and resume scheduler from this state next time
2. Don't forget to use epoch + 1 in step method like in Ross's train.py script. It will zero out your lr if you give it epoch=0 argument with this block of code

if epoch <= self.warmup_t:
  lrs = [self.warmup_lr_init + epoch * s for s in self.warmup_steps]
  super().update_groups(lrs)

3. The most tricky part. On the resuming stage, make sure not to provide warmup_t argument. Because, if it is > 0, and you provide lr_min and warmup_lr_init, the lr of optimizer is really messed up, and you get lr_min instead of resumed optimizer lr. I have not completely understood how it is messing up with all thouse lrs. Important part, I also have to use some ugly workaround and delete optimizer.param_groups[0]["initial_lr"] before scheduler initialization in this stage.

This is how I get it to work the way I expected it to work. I may not figure out a clean and simple solution, perhaps @rwightman can shed some light.

All the best :)

[rwightman]

@ademyanchuk thanks for the notes, yeah, these are all ugly little bits that I haven't smoothed out... in general I'd say the LR schedulers are one aspect of PyTorch that I've never liked the design for, despite their attempts to improve them. I tried to improve on some aspects here, but still work to do. I'm working on some new training 'cogs' and at some point I will re-tackle these schedulers.