Support torch.optim.lr_scheduler.ReduceLROnPlateau

pytorch_lightning/callbacks/__init__.py

@@ -1,7 +1,8 @@
-from .pt_callbacks import EarlyStopping, ModelCheckpoint, GradientAccumulationScheduler
+from .pt_callbacks import EarlyStopping, ModelCheckpoint, ReduceLROnPlateauScheduler, GradientAccumulationScheduler
 
 __all__ = [
     'EarlyStopping',
     'ModelCheckpoint',
+    'ReduceLROnPlateauScheduler',
     'GradientAccumulationScheduler',
 ]

pytorch_lightning/callbacks/pt_callbacks.py

@@ -145,6 +145,36 @@ def on_train_end(self, logs=None):
             print('Epoch %05d: early stopping' % (self.stopped_epoch + 1))
 
 
+class ReduceLROnPlateauScheduler(Callback):
+    """
+    Reduce learning rate when the monitored metric has stopped improving.
+    Wrapper for torch.optim.lr_schuduler.ReduceLROnPlateau learning rate
+    schedulers.
+
+    # Arguments
+        schedulers: list of torch.optim.lr_scheduler.ReduceLROnPlateau
+        monitor: quantity to be monitored.
+    """
+
+    def __init__(self, schedulers, monitor='val_loss'):
+        super(ReduceLROnPlateauScheduler, self).__init__()
+
+        self.monitor = monitor
+        self.schedulers = schedulers
+
+    def on_epoch_end(self, epoch, logs=None):
+        current = logs.get(self.monitor)
+        stop_training = False
+        if current is None:
+            print('ReduceLROnPlateau conditioned on metric `%s` '
+                  'which is not available. Available metrics are: %s' %
+                  (self.monitor, ','.join(list(logs.keys()))), RuntimeWarning)
+            exit(-1)
+
+        for scheduler in self.schedulers:
+            scheduler.step(current, epoch=epoch)
+
+
 class ModelCheckpoint(Callback):
     """Save the model after every epoch.
     `filepath` can contain named formatting options,

pytorch_lightning/trainer/trainer.py

@@ -21,7 +21,7 @@
 from pytorch_lightning.pt_overrides.override_data_parallel import (
     LightningDistributedDataParallel, LightningDataParallel)
 from pytorch_lightning.callbacks import GradientAccumulationScheduler, \
-    ModelCheckpoint, EarlyStopping
+    ReduceLROnPlateauScheduler, ModelCheckpoint, EarlyStopping
 from pytorch_lightning.utilities.debugging import MisconfigurationException
 import pdb
 from pytorch_lightning.trainer import ignored_warnings
@@ -190,6 +190,7 @@ def __init__(self,
         else:
             self.early_stop_callback = early_stop_callback
             self.enable_early_stop = True
+        self.lr_scheduler_callback = None
 
         # configure logger
         if logger is True:
@@ -793,12 +794,25 @@ def init_optimizers(self, optimizers):
         # two lists
         elif len(optimizers) == 2 and isinstance(optimizers[0], list):
             optimizers, lr_schedulers = optimizers
+            lr_schedulers = self.configure_schedulers(lr_schedulers)
             return optimizers, lr_schedulers
 
         # single list or tuple
         elif isinstance(optimizers, list) or isinstance(optimizers, tuple):
             return optimizers, []
 
+    def configure_schedulers(self, schedulers):
+        custom_schedulers = []
+        i = 0
+        while i < len(schedulers):
+            if isinstance(schedulers[i], torch.optim.lr_scheduler.ReduceLROnPlateau):
+                custom_schedulers.append(schedulers.pop(i))
+            i += 1
+        if custom_schedulers:
+            self.lr_scheduler_callback = ReduceLROnPlateauScheduler(custom_schedulers,
+                                                                    monitor='val_loss')
+        return schedulers
+
     def __single_gpu_train(self, model):
         # CHOOSE OPTIMIZER
         # allow for lr schedulers as well
@@ -1096,7 +1110,7 @@ def __train(self):
             # update LR schedulers
             if self.lr_schedulers is not None:
                 for lr_scheduler in self.lr_schedulers:
-                    lr_scheduler.step(self.current_epoch)
+                    lr_scheduler.step(epoch=self.current_epoch)
 
             # early stopping
             met_min_epochs = epoch_nb > self.min_nb_epochs
@@ -1544,6 +1558,11 @@ def __run_evaluation(self, test=False):
                 tqdm_metrics = self.__training_tqdm_dict
                 self.progress_bar.set_postfix(**tqdm_metrics)
 
+            # reduce learning rate based on metrics
+            if self.lr_scheduler_callback is not None and not test:
+                self.lr_scheduler_callback.on_epoch_end(epoch=self.current_epoch,
+                                                        logs=self.__training_tqdm_dict)
+
         # model checkpointing
         if self.proc_rank == 0 and self.checkpoint_callback is not None and not test:
             self.checkpoint_callback.on_epoch_end(epoch=self.current_epoch,