Make reduce_on_plateau more flexible
🚀 Feature
I want to have freedom to use a custom reduce on plateau type scheduler.
Motivation
Right now, only torch.optim.lr_scheduler.ReduceLROnPlateau is supported out of the box. See here. This means even if I specify {'reduce_on_plateau': True} in:
def configure_optimizers(self):
  optimizer = TrainConf.optimizer(self.parameters(), **TrainConf.optimizer_params)
  scheduler = TrainConf.scheduler(optimizer, **TrainConf.scheduler_params)
  return {
      'optimizer': optimizer, 
      'lr_scheduler': {
          'scheduler': scheduler, 
          'interval': TrainConf.scheduler_interval,
          'monitor': 'val_lwlrap',
          'reduce_on_plateau': isinstance(scheduler,
                                          (WarmupReduceLROnPlateau, # this is my custom class
                                           ReduceLROnPlateau)),
      }
  }
this will be overridden downstream.
Pitch
At the very least TrainerOptimizersMixin.configure_schedulers should check to see if the key was provided, and if so, don't overwrite it. So I propose changing
scheduler['reduce_on_plateau'] = isinstance(
    scheduler['scheduler'], optim.lr_scheduler.ReduceLROnPlateau
)
to
if 'reduce_on_plateau' not in scheduler:
  scheduler['reduce_on_plateau'] = isinstance(
      scheduler['scheduler'], optim.lr_scheduler.ReduceLROnPlateau
  )
But in the long run I'd suggest we allow the user to override a method of Trainer like this:
def step_scheduler(self):
  ...
That way we could access the variables needed and provide our own signature to the step method of the scheduler class.
Additional context
My particular custom scheduler looks like
class WarmupReduceLROnPlateau(ReduceLROnPlateau):
    """
    Subclassing torch.optim.lr_scheduler.ReduceLROnPlateau
    added warmup parameters
    """
    def __init__(self, optimizer, mode='min', factor=0.1, patience=10,
                 threshold=1e-4, threshold_mode='rel', cooldown=0,
                 min_lr=0, eps=1e-8, warmup_itrs=0, warmup_type='lin',
                 start_lr=1e-16, verbose=False):
        super().__init__(optimizer, mode=mode, factor=factor, patience=patience,
                 threshold=threshold, threshold_mode=threshold_mode,
                 cooldown=cooldown, min_lr=min_lr, eps=eps, verbose=verbose)
        self.warmup_itrs = warmup_itrs
        self.warmup_type = warmup_type
        self.start_lr = start_lr
        self.default_lrs = []
        self.itr = 0
        for param_group in optimizer.param_groups:
            self.default_lrs.append(param_group['lr'])

    def step(self, metrics):
        if self.itr < self.warmup_itrs:
            for i, param_group in enumerate(self.optimizer.param_groups):
                if self.warmup_type == 'exp':
                    new_lr = self.start_lr * \
                        (self.default_lrs[i] /
                         self.start_lr)**(self.itr/self.warmup_itrs)
                if self.warmup_type == 'lin':
                    new_lr = self.start_lr + \
                        (self.default_lrs[i] - self.start_lr) * \
                        (self.itr/self.warmup_itrs)
                param_group['lr'] = new_lr
        elif self.itr == self.warmup_itrs:
            param_group['lr'] = self.default_lrs[i]
        else:
            super.step(metrics)
        self.itr += 1

If this gets any traction I'd love to contribute a PR (would be my first to a widely used repo!)