[Bugfix] Fixed epoch level schedulers not being called when val_check_interval < 1.0 (Lightning-AI#6075)

* fix bug

* fix tests

* changelog

* fix pep8

* fix tests

* fix and add some tests

* add test for rlop

* chlog

* Update CHANGELOG.md

Co-authored-by: rohitgr7 <[email protected]>

Author: 2 people

CHANGELOG.md

@@ -46,6 +46,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Expose DeepSpeed loss parameters to allow users to fix loss instability ([#6115](https://github.com/PyTorchLightning/pytorch-lightning/pull/6115))
 
 
+- Fixed epoch level schedulers not being called when `val_check_interval < 1.0` ([#6075](https://github.com/PyTorchLightning/pytorch-lightning/pull/6075))
+
+
 ## [1.2.1] - 2021-02-23
 
 ### Fixed

pytorch_lightning/trainer/connectors/optimizer_connector.py

@@ -66,13 +66,21 @@ def update_learning_rates(self, interval: str, monitor_metrics=None):
                         continue
                 # update LR
                 old_lr = lr_scheduler['scheduler'].optimizer.param_groups[0]['lr']
+
                 if lr_scheduler['reduce_on_plateau']:
                     lr_scheduler['scheduler'].step(monitor_val)
                 else:
                     lr_scheduler['scheduler'].step()
+
                 new_lr = lr_scheduler['scheduler'].optimizer.param_groups[0]['lr']
 
                 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/trainer/training_loop.py

@@ -478,6 +478,7 @@ def run_training_epoch(self):
         train_dataloader = self.trainer.data_connector.get_profiled_train_dataloader(train_dataloader)
         dataloader_idx = 0
         should_check_val = False
+        val_loop_called = False
 
         for batch_idx, (batch, is_last_batch) in train_dataloader:
 
@@ -513,6 +514,7 @@ def run_training_epoch(self):
             should_check_val = self.should_check_val_fx(batch_idx, is_last_batch)
             if should_check_val:
                 self.trainer.run_evaluation()
+                val_loop_called = True
 
                 # reset stage to train
                 self.trainer._running_stage = RunningStage.TRAINING
@@ -558,21 +560,23 @@ def run_training_epoch(self):
         )
 
         should_check_val = self.should_check_val_fx(batch_idx, is_last_batch, on_epoch=True)
-        if should_check_val:
-            self.trainer.run_evaluation(on_epoch=True)
-
-            # reset stage to train
-            self.trainer._running_stage = RunningStage.TRAINING
-
         should_skip_eval = self.trainer.evaluation_loop.should_skip_evaluation(self.trainer.num_val_batches)
         should_train_only = self.trainer.disable_validation or should_skip_eval
 
-        if should_train_only:
-            # update epoch level lr_schedulers
+        # update epoch level lr_schedulers if no val loop outside train loop is triggered
+        if (val_loop_called and not should_check_val) or should_train_only:
             self.trainer.optimizer_connector.update_learning_rates(interval='epoch')
+
+        if should_train_only:
             self.check_checkpoint_callback(True)
             self.check_early_stopping_callback(True)
 
+        if should_check_val:
+            self.trainer.run_evaluation(on_epoch=True)
+
+            # reset stage to train
+            self.trainer._running_stage = RunningStage.TRAINING
+
         # increment the global step once
         # progress global step according to grads progress
         self.increment_accumulated_grad_global_step()
@@ -818,7 +822,7 @@ def should_check_val_fx(self, batch_idx, is_last_batch, on_epoch=False):
         is_val_check_epoch = (self.trainer.current_epoch + 1) % self.trainer.check_val_every_n_epoch == 0
         can_check_val = self.trainer.enable_validation and is_val_check_epoch
         is_last_batch_for_infinite_dataset = is_last_batch and self.trainer.val_check_batch == float("inf")
-        epoch_end_val_check = self.trainer.val_check_batch == self.trainer.num_training_batches
+        epoch_end_val_check = (batch_idx + 1) % self.trainer.num_training_batches == 0
 
         should_check_val = ((is_val_check_batch and epoch_end_val_check) or self.trainer.should_stop
                             or is_last_batch_for_infinite_dataset

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_correct_score_and_checkpoint(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):
 
@@ -74,21 +77,28 @@ def __init__(self):
             self.val_logs = torch.randn(max_epochs, limit_val_batches)
 
         def training_step(self, batch, batch_idx):
-            out = super().training_step(batch, batch_idx)
             log_value = self.train_log_epochs[self.current_epoch, batch_idx]
             self.log('train_log', log_value, on_epoch=True)
-            return out
+            return super().training_step(batch, batch_idx)
 
         def validation_step(self, batch, batch_idx):
-            out = super().validation_step(batch, batch_idx)
             log_value = self.val_logs[self.current_epoch, batch_idx]
             self.log('val_log', log_value)
             self.log('epoch', self.current_epoch, on_epoch=True)
-            return out
+            return super().validation_step(batch, batch_idx)
 
         def configure_optimizers(self):
-            optimizer = torch.optim.SGD(self.layer.parameters(), lr=0.2)
-            lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1)
+            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}'
@@ -109,11 +119,15 @@ def configure_optimizers(self):
         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(lr_scheduler_debug) == max_epochs
 
     for epoch in range(max_epochs):
         score = scores[epoch]
@@ -130,9 +144,118 @@ def configure_optimizers(self):
         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], 4 == 0.2 * (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,reduce_lr_on_plateau",
+    [
+        (0.25, True),
+        (0.25, False),
+        (0.33, False),
+    ],
+)
+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
+
+    class CustomBoringModel(BoringModel):
+
+        def __init__(self):
+            super().__init__()
+            self.val_logs = torch.randn(per_epoch_call_count * max_epochs, limit_val_batches)
+            self.val_loop_count = 0
+
+        def validation_step(self, batch, batch_idx):
+            log_value = self.val_logs[self.val_loop_count, batch_idx]
+            self.log('val_log', log_value)
+            self.log('epoch', self.current_epoch, on_epoch=True)
+            return super().validation_step(batch, batch_idx)
+
+        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)
+
+    model = CustomBoringModel()
+
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        callbacks=[checkpoint],
+        limit_train_batches=limit_train_batches,
+        limit_val_batches=limit_val_batches,
+        max_epochs=max_epochs,
+        val_check_interval=val_check_interval,
+        progress_bar_refresh_rate=0,
+        num_sanity_val_steps=0,
+    )
+    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(lr_scheduler_debug) == max_epochs
+
+    for epoch in range(max_epochs):
+        for ix in range(per_epoch_call_count):
+            global_ix = ix + per_epoch_call_count * epoch
+            score = scores[global_ix]
+            expected_score = getattr(model, f'{monitor}s')[global_ix].mean().item()
+            expected_filename = f'{monitor}={score:.4f}-epoch={epoch}.ckpt'
+            assert math.isclose(score, expected_score, rel_tol=1e-4)
+
+            chk = pl_load(os.path.join(checkpoint.dirpath, expected_filename))
+            assert chk['epoch'] == epoch + 1
+            assert chk['global_step'] == per_epoch_steps * (global_ix + 1)
+
+            mc_specific_data = chk['callbacks'][type(checkpoint)]
+            assert mc_specific_data['dirpath'] == checkpoint.dirpath
+            assert mc_specific_data['monitor'] == monitor
+            assert mc_specific_data['current_score'] == score
+
+            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))
+
+        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])

tests/trainer/optimization/test_optimizers.py

@@ -34,6 +34,7 @@ def test_optimizer_with_scheduling(tmpdir):
         max_epochs=1,
         limit_val_batches=0.1,
         limit_train_batches=0.2,
+        val_check_interval=0.5
     )
     trainer.fit(model)
     assert trainer.state == TrainerState.FINISHED, f"Training failed with {trainer.state}"
@@ -164,15 +165,15 @@ def test_reducelronplateau_scheduling(tmpdir):
     model.configure_optimizers = lambda: {
         'optimizer': optimizer,
         'lr_scheduler': torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer),
-        'monitor': 'early_stop_on',
+        'monitor': 'val_acc',
     }
     trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True)
     trainer.fit(model)
     assert trainer.state == TrainerState.FINISHED, f"Training failed with {trainer.state}"
     lr_scheduler = trainer.lr_schedulers[0]
     assert lr_scheduler == dict(
         scheduler=lr_scheduler['scheduler'],
-        monitor='early_stop_on',
+        monitor='val_acc',
         interval='epoch',
         frequency=1,
         reduce_on_plateau=True,