add test for rlop

Author: rohitgr7

pytorch_lightning/trainer/connectors/optimizer_connector.py

@@ -81,5 +81,11 @@ def update_learning_rates(self, interval: str, monitor_metrics=None):
 
                 if self.trainer.dev_debugger.enabled:
                     self.trainer.dev_debugger.track_lr_schedulers_update(
-                        self.trainer.batch_idx, interval, scheduler_idx, old_lr, new_lr, monitor_key=monitor_key
+                        self.trainer.batch_idx,
+                        interval,
+                        scheduler_idx,
+                        old_lr,
+                        new_lr,
+                        monitor_key=monitor_key,
+                        monitor_val=monitor_val
                     )

pytorch_lightning/utilities/debugging.py

@@ -121,13 +121,16 @@ def track_train_loss_history(self, batch_idx, loss):
         self.saved_train_losses.append(loss_dict)
 
     @enabled_only
-    def track_lr_schedulers_update(self, batch_idx, interval, scheduler_idx, old_lr, new_lr, monitor_key=None):
+    def track_lr_schedulers_update(
+        self, batch_idx, interval, scheduler_idx, old_lr, new_lr, monitor_key=None, monitor_val=None
+    ):
         loss_dict = {
             'batch_idx': batch_idx,
             'interval': interval,
             'scheduler_idx': scheduler_idx,
             'epoch': self.trainer.current_epoch,
             'monitor_key': monitor_key,
+            'monitor_val': monitor_val,
             'old_lr': old_lr,
             'new_lr': new_lr
         }

tests/checkpointing/test_model_checkpoint.py

@@ -26,6 +26,7 @@
 import torch
 import yaml
 from omegaconf import Container, OmegaConf
+from torch import optim
 
 import pytorch_lightning as pl
 import tests.helpers.utils as tutils
@@ -47,8 +48,8 @@ def training_step(self, batch, batch_idx):
 
     def validation_epoch_end(self, outputs):
         outs = torch.stack([x['x'] for x in outputs]).mean()
-        self.log('epoch', self.current_epoch, on_epoch=True)
-        self.log('val_acc', outs, on_epoch=True)
+        self.log('epoch', self.current_epoch)
+        self.log('val_acc', outs)
 
 
 @mock.patch.dict(os.environ, {"PL_DEV_DEBUG": "1"})
@@ -57,14 +58,16 @@ def validation_epoch_end(self, outputs):
     [('base', "base", 'val_log'), ('base', "base", 'train_log_epoch'), (None, "base", 'train_log_epoch'),
      ("base", None, 'train_log_epoch')],
 )
-def test_model_checkpoint_score_and_ckpt(tmpdir, validation_step, val_dataloaders, monitor):
+@pytest.mark.parametrize('reduce_lr_on_plateau', [False, True])
+def test_model_checkpoint_score_and_ckpt(tmpdir, validation_step, val_dataloaders, monitor, reduce_lr_on_plateau):
     """
     Test that when a model checkpoint is saved, it saves with
     the correct score appended to ckpt_path and checkpoint data
     """
     max_epochs = 3
     limit_train_batches = 5
     limit_val_batches = 7
+    lr = 1e-1
 
     class CustomBoringModel(BoringModel):
 
@@ -84,6 +87,20 @@ def validation_step(self, batch, batch_idx):
             self.log('epoch', self.current_epoch, on_epoch=True)
             return super().validation_step(batch, batch_idx)
 
+        def configure_optimizers(self):
+            optimizer = optim.SGD(self.parameters(), lr=lr)
+
+            if reduce_lr_on_plateau:
+                lr_scheduler = {
+                    'scheduler': optim.lr_scheduler.ReduceLROnPlateau(optimizer),
+                    'monitor': monitor,
+                    'strict': True,
+                }
+            else:
+                lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1)
+
+            return [optimizer], [lr_scheduler]
+
     filename = '{' + f'{monitor}' + ':.4f}-{epoch}'
     checkpoint = ModelCheckpoint(dirpath=tmpdir, filename=filename, monitor=monitor, save_top_k=-1)
 
@@ -102,12 +119,15 @@ def validation_step(self, batch, batch_idx):
         max_epochs=max_epochs,
         progress_bar_refresh_rate=0,
     )
-    trainer.fit(model)
+    results = trainer.fit(model)
+    assert results
+    assert trainer.state == TrainerState.FINISHED, f"Training failed with {trainer.state}"
 
     ckpt_files = list(Path(tmpdir).glob('*.ckpt'))
     scores = [metric[monitor] for metric in trainer.dev_debugger.logged_metrics if monitor in metric]
+    lr_scheduler_debug = trainer.dev_debugger.saved_lr_scheduler_updates
     assert len(ckpt_files) == len(scores) == max_epochs
-    assert len(trainer.dev_debugger.saved_lr_scheduler_updates) == max_epochs
+    assert len(lr_scheduler_debug) == max_epochs
 
     for epoch in range(max_epochs):
         score = scores[epoch]
@@ -124,27 +144,33 @@ def validation_step(self, batch, batch_idx):
         assert mc_specific_data['monitor'] == monitor
         assert mc_specific_data['current_score'] == score
 
-        lr_scheduler_specific_data = chk['lr_schedulers'][0]
-        assert lr_scheduler_specific_data['_step_count'] == epoch + 2
-        assert lr_scheduler_specific_data['_last_lr'][0] == 0.1 * (0.1**(epoch + 1))
+        if not reduce_lr_on_plateau:
+            lr_scheduler_specific_data = chk['lr_schedulers'][0]
+            assert lr_scheduler_specific_data['_step_count'] == epoch + 2
+            assert lr_scheduler_specific_data['_last_lr'][0] == lr * (lr**(epoch + 1))
+
+        assert lr_scheduler_debug[epoch]['monitor_val'] == (score if reduce_lr_on_plateau else None)
+        assert lr_scheduler_debug[epoch]['monitor_key'] == (monitor if reduce_lr_on_plateau else None)
 
 
 @mock.patch.dict(os.environ, {"PL_DEV_DEBUG": "1"})
 @pytest.mark.parametrize(
-    "val_check_interval,lr_sched_step_count_inc",
+    "val_check_interval,reduce_lr_on_plateau",
     [
-        (0.25, 1),
-        (0.33, 0),
+        (0.25, True),
+        (0.25, False),
+        (0.33, False),
     ],
 )
-def test_model_checkpoint_score_and_ckpt_val_check_interval(tmpdir, val_check_interval, lr_sched_step_count_inc):
+def test_model_checkpoint_score_and_ckpt_val_check_interval(tmpdir, val_check_interval, reduce_lr_on_plateau):
     """
     Test that when a model checkpoint is saved, it saves with the correct
     score appended to ckpt_path and checkpoint data with val_check_interval
     """
     max_epochs = 3
     limit_train_batches = 12
     limit_val_batches = 7
+    lr = 1e-1
     monitor = 'val_log'
     per_epoch_steps = int(limit_train_batches * val_check_interval)
     per_epoch_call_count = limit_train_batches // per_epoch_steps
@@ -166,6 +192,20 @@ def validation_epoch_end(self, outputs):
             self.val_loop_count += 1
             super().validation_epoch_end(outputs)
 
+        def configure_optimizers(self):
+            optimizer = optim.SGD(self.parameters(), lr=lr)
+
+            if reduce_lr_on_plateau:
+                lr_scheduler = {
+                    'scheduler': optim.lr_scheduler.ReduceLROnPlateau(optimizer),
+                    'monitor': monitor,
+                    'strict': True,
+                }
+            else:
+                lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1)
+
+            return [optimizer], [lr_scheduler]
+
     filename = '{' + f'{monitor}' + ':.4f}-{epoch}'
     checkpoint = ModelCheckpoint(dirpath=tmpdir, filename=filename, monitor=monitor, save_top_k=-1)
 
@@ -181,12 +221,15 @@ def validation_epoch_end(self, outputs):
         progress_bar_refresh_rate=0,
         num_sanity_val_steps=0,
     )
-    trainer.fit(model)
+    results = trainer.fit(model)
+    assert results
+    assert trainer.state == TrainerState.FINISHED, f"Training failed with {trainer.state}"
 
     ckpt_files = list(Path(tmpdir).glob('*.ckpt'))
     scores = [metric[monitor] for metric in trainer.dev_debugger.logged_metrics if monitor in metric]
+    lr_scheduler_debug = trainer.dev_debugger.saved_lr_scheduler_updates
     assert len(ckpt_files) == len(scores) == per_epoch_call_count * max_epochs
-    assert len(trainer.dev_debugger.saved_lr_scheduler_updates) == max_epochs
+    assert len(lr_scheduler_debug) == max_epochs
 
     for epoch in range(max_epochs):
         for ix in range(per_epoch_call_count):
@@ -205,11 +248,14 @@ def validation_epoch_end(self, outputs):
             assert mc_specific_data['monitor'] == monitor
             assert mc_specific_data['current_score'] == score
 
-            lr_scheduler_specific_data = chk['lr_schedulers'][0]
+            if not reduce_lr_on_plateau:
+                lr_scheduler_specific_data = chk['lr_schedulers'][0]
+                did_update = 1 if ix + 1 == per_epoch_call_count else 0
+                assert lr_scheduler_specific_data['_step_count'] == epoch + 1 + did_update
+                assert lr_scheduler_specific_data['_last_lr'][0] == lr * (lr**(epoch + did_update))
 
-            did_update = 1 if ix + 1 == per_epoch_call_count else 0
-            assert lr_scheduler_specific_data['_step_count'] == epoch + 1 + did_update
-            assert lr_scheduler_specific_data['_last_lr'][0] == 0.1 * (0.1**(epoch + did_update))
+        assert lr_scheduler_debug[epoch]['monitor_val'] == (score if reduce_lr_on_plateau else None)
+        assert lr_scheduler_debug[epoch]['monitor_key'] == (monitor if reduce_lr_on_plateau else None)
 
 
 @pytest.mark.parametrize("save_top_k", [-1, 0, 1, 2])