fix the call for scheduler

Author: rohitgr7

pytorch_lightning/trainer/evaluation_loop.py

@@ -72,18 +72,6 @@ def get_evaluation_dataloaders(self, max_batches):
 
         return dataloaders, max_batches
 
-    def should_skip_evaluation(self, dataloaders, max_batches):
-        # skip when dataloaders aren't defined
-        if dataloaders is None:
-            return True
-
-        # enable disabling validation step with limit_val_batches = 0
-        should_skip = sum(max_batches) == 0
-        if should_skip:
-            return True
-
-        return False
-
     def on_evaluation_start(self, *args, **kwargs):
         if self.testing:
             self.trainer.call_hook('on_test_start', *args, **kwargs)

pytorch_lightning/trainer/trainer.py

@@ -548,7 +548,7 @@ def train(self):
             # hook
             self.train_loop.on_train_end()
 
-    def run_evaluation(self, test_mode: bool = False, max_batches=None):
+    def run_evaluation(self, test_mode: bool = False, max_batches=None, on_epoch=False):
 
         # used to know if we are logging for val, test + reset cached results
         self.logger_connector.set_stage(test_mode, reset=True)
@@ -560,7 +560,7 @@ def run_evaluation(self, test_mode: bool = False, max_batches=None):
         dataloaders, max_batches = self.evaluation_loop.get_evaluation_dataloaders(max_batches)
 
         # check if we want to skip this evaluation
-        if self.evaluation_loop.should_skip_evaluation(dataloaders, max_batches):
+        if sum(max_batches) == 0:
             return [], []
 
         # ref model
@@ -621,6 +621,10 @@ def run_evaluation(self, test_mode: bool = False, max_batches=None):
         # hook
         self.evaluation_loop.on_evaluation_epoch_end()
 
+        # update epoch-level lr_schedulers
+        if on_epoch:
+            self.optimizer_connector.update_learning_rates(interval='epoch')
+
         # hook
         self.evaluation_loop.on_evaluation_end()
 

pytorch_lightning/trainer/training_loop.py

@@ -626,19 +626,13 @@ def run_training_epoch(self):
             self.trainer.total_batch_idx += 1
 
             # stop epoch if we limited the number of training batches
-            if self._num_training_batches_reached():
+            if self._num_training_batches_reached(is_last_batch):
                 break
 
             # progress global step according to grads progress
             self.increment_accumulated_grad_global_step()
 
         # epoch end hook
-        should_check_val = self.should_check_val_fx(batch_idx, is_last_batch, on_epoch=True)
-        if should_check_val:
-            self.trainer.run_evaluation(test_mode=False)
-            # reset stage to train
-            self.trainer.logger_connector.set_stage("train")
-
         self.run_on_epoch_end_hook(epoch_output)
 
         # log epoch metrics
@@ -649,10 +643,19 @@ def run_training_epoch(self):
             self.num_optimizers
         )
 
-        # update LR schedulers
-        self.trainer.optimizer_connector.update_learning_rates(interval='epoch')
+        should_check_val = self.should_check_val_fx(batch_idx, is_last_batch, on_epoch=True)
+        if should_check_val:
+            self.trainer.run_evaluation(test_mode=False, on_epoch=True)
+            # reset stage to train
+            self.trainer.logger_connector.set_stage("train")
+
+        should_skip_eval = sum(self.trainer.num_val_batches) == 0
+        should_train_only_check = not self.trainer.enable_validation and should_skip_eval
+
+        if should_skip_eval or should_train_only_check:
+            # update epoch level lr_schedulers
+            self.trainer.optimizer_connector.update_learning_rates(interval='epoch')
 
-        should_train_only_check = not self.trainer.enable_validation and (sum(self.trainer.num_val_batches) == 0)
         self.check_checkpoint_callback(should_train_only_check)
         self.check_early_stopping_callback(should_train_only_check)
 
@@ -890,8 +893,8 @@ def increment_accumulated_grad_global_step(self):
     def _accumulated_batches_reached(self):
         return (self.trainer.batch_idx + 1) % self.trainer.accumulate_grad_batches == 0
 
-    def _num_training_batches_reached(self):
-        return (self.trainer.batch_idx + 1) == self.trainer.num_training_batches
+    def _num_training_batches_reached(self, is_last_batch=False):
+        return (self.trainer.batch_idx + 1) == self.trainer.num_training_batches or is_last_batch
 
     def should_accumulate(self):
         # checks if backward or backward + optimizer step (via closure)

tests/callbacks/test_callbacks.py

@@ -88,15 +88,15 @@ def test_trainer_callback_system(torch_save):
         call.on_before_zero_grad(trainer, model, trainer.optimizers[0]),
         call.on_batch_end(trainer, model),
         call.on_train_batch_end(trainer, model, ANY, ANY, 2, 0),
+        call.on_epoch_end(trainer, model),
+        call.on_train_epoch_end(trainer, model, ANY),
         call.on_validation_start(trainer, model),
         call.on_validation_epoch_start(trainer, model),
         call.on_validation_batch_start(trainer, model, ANY, 0, 0),
         call.on_validation_batch_end(trainer, model, ANY, ANY, 0, 0),
         call.on_validation_epoch_end(trainer, model),
         call.on_validation_end(trainer, model),
         call.on_save_checkpoint(trainer, model),
-        call.on_epoch_end(trainer, model),
-        call.on_train_epoch_end(trainer, model, ANY),
         call.on_train_end(trainer, model),
         call.on_fit_end(trainer, model),
         call.teardown(trainer, model, 'fit'),

tests/models/test_hooks.py

@@ -328,15 +328,15 @@ def on_test_model_train(self):
         'on_after_backward',
         'on_before_zero_grad',
         'on_train_batch_end',
+        'on_epoch_end',
+        'on_train_epoch_end',
         'on_validation_model_eval',
         'on_validation_epoch_start',
         'on_validation_batch_start',
         'on_validation_batch_end',
         'on_validation_epoch_end',
         'on_save_checkpoint',
         'on_validation_model_train',
-        'on_epoch_end',
-        'on_train_epoch_end',
         'on_train_end',
         'on_fit_end',
     ]