Update tests

Author: carmocca

tests/tests_pytorch/accelerators/test_ipu.py

@@ -47,15 +47,6 @@ def validation_step(self, batch, batch_idx):
     def test_step(self, batch, batch_idx):
         return self.step(batch)
 
-    def training_epoch_end(self, outputs) -> None:
-        pass
-
-    def validation_epoch_end(self, outputs) -> None:
-        pass
-
-    def test_epoch_end(self, outputs) -> None:
-        pass
-
 
 class IPUClassificationModel(ClassificationModel):
     def training_step(self, batch, batch_idx):

tests/tests_pytorch/accelerators/test_tpu.py

@@ -155,7 +155,6 @@ def on_train_end(self):
     model = ManualOptimizationModel()
     model_copy = deepcopy(model)
     model.training_step_end = None
-    model.training_epoch_end = None
 
     trainer = Trainer(
         max_epochs=1,

tests/tests_pytorch/callbacks/progress/test_tqdm_progress_bar.py

@@ -136,8 +136,6 @@ def predict_step(self, batch, batch_idx, dataloader_idx=None):
             return
 
     model = CustomModel()
-    model.validation_epoch_end = None
-    model.test_epoch_end = None
 
     # check the sanity dataloaders
     num_sanity_val_steps = 4

tests/tests_pytorch/callbacks/test_callback_hook_outputs.py

@@ -41,9 +41,6 @@ def on_validation_batch_end(self, outputs, batch, batch_idx: int, dataloader_idx
         def on_test_batch_end(self, outputs, batch, batch_idx: int, dataloader_idx: int) -> None:
             assert "x" in outputs
 
-        def training_epoch_end(self, outputs) -> None:
-            assert len(outputs) == self.trainer.num_training_batches
-
     model = TestModel()
 
     trainer = Trainer(
@@ -59,22 +56,3 @@ def training_epoch_end(self, outputs) -> None:
     assert any(isinstance(c, CB) for c in trainer.callbacks)
 
     trainer.fit(model)
-
-
-def test_free_memory_on_eval_outputs(tmpdir):
-    class CB(Callback):
-        def on_train_epoch_end(self, trainer, pl_module):
-            assert not trainer._evaluation_loop._outputs
-
-    model = BoringModel()
-
-    trainer = Trainer(
-        callbacks=CB(),
-        default_root_dir=tmpdir,
-        limit_train_batches=2,
-        limit_val_batches=2,
-        max_epochs=1,
-        enable_model_summary=False,
-    )
-
-    trainer.fit(model)

tests/tests_pytorch/callbacks/test_early_stopping.py

@@ -134,7 +134,7 @@ def test_early_stopping_patience(tmpdir, loss_values: list, patience: int, expec
     class ModelOverrideValidationReturn(BoringModel):
         validation_return_values = torch.tensor(loss_values)
 
-        def validation_epoch_end(self, outputs):
+        def on_validation_epoch_end(self):
             loss = self.validation_return_values[self.current_epoch]
             self.log("test_val_loss", loss)
 
@@ -164,7 +164,7 @@ def test_early_stopping_patience_train(
     class ModelOverrideTrainReturn(BoringModel):
         train_return_values = torch.tensor(loss_values)
 
-        def training_epoch_end(self, outputs):
+        def on_train_epoch_end(self):
             loss = self.train_return_values[self.current_epoch]
             self.log("train_loss", loss)
 
@@ -226,7 +226,7 @@ def test_early_stopping_no_val_step(tmpdir):
 )
 def test_early_stopping_thresholds(tmpdir, stopping_threshold, divergence_threshold, losses, expected_epoch):
     class CurrentModel(BoringModel):
-        def validation_epoch_end(self, outputs):
+        def on_validation_epoch_end(self):
             val_loss = losses[self.current_epoch]
             self.log("abc", val_loss)
 
@@ -252,7 +252,7 @@ def test_early_stopping_on_non_finite_monitor(tmpdir, stop_value):
     expected_stop_epoch = 2
 
     class CurrentModel(BoringModel):
-        def validation_epoch_end(self, outputs):
+        def on_validation_epoch_end(self):
             val_loss = losses[self.current_epoch]
             self.log("val_loss", val_loss)
 
@@ -352,12 +352,12 @@ def _epoch_end(self) -> None:
         self.log("abc", torch.tensor(loss))
         self.log("cba", torch.tensor(0))
 
-    def training_epoch_end(self, outputs):
+    def on_train_epoch_end(self):
         if not self.early_stop_on_train:
             return
         self._epoch_end()
 
-    def validation_epoch_end(self, outputs):
+    def on_validation_epoch_end(self):
         if self.early_stop_on_train:
             return
         self._epoch_end()

tests/tests_pytorch/callbacks/test_lr_monitor.py

@@ -248,7 +248,6 @@ def configure_optimizers(self):
             return [optimizer1, optimizer2], [lr_scheduler1, lr_scheduler2]
 
     model = CustomBoringModel()
-    model.training_epoch_end = None
 
     lr_monitor = LearningRateMonitor(logging_interval=logging_interval)
     log_every_n_steps = 2
@@ -306,7 +305,6 @@ def configure_optimizers(self):
             return [optimizer1, optimizer2]
 
     model = CustomBoringModel()
-    model.training_epoch_end = None
 
     lr_monitor = LearningRateMonitor(logging_interval=logging_interval)
     log_every_n_steps = 2
@@ -563,7 +561,6 @@ def finetune_function(self, pl_module, epoch: int, optimizer, opt_idx: int):
         enable_checkpointing=False,
     )
     model = TestModel()
-    model.training_epoch_end = None
     trainer.fit(model)
 
     expected = [0.1, 0.1, 0.1, 0.1, 0.1]

tests/tests_pytorch/checkpointing/test_checkpoint_callback_frequency.py

@@ -96,7 +96,7 @@ def training_step(self, batch, batch_idx):
             self.log("my_loss", batch_idx * (1 + local_rank), on_epoch=True)
             return super().training_step(batch, batch_idx)
 
-        def training_epoch_end(self, outputs) -> None:
+        def on_train_epoch_end(self):
             local_rank = int(os.getenv("LOCAL_RANK"))
             if self.trainer.is_global_zero:
                 self.log("my_loss_2", (1 + local_rank), on_epoch=True, rank_zero_only=True)

tests/tests_pytorch/checkpointing/test_model_checkpoint.py

@@ -59,9 +59,8 @@ def training_step(self, batch, batch_idx):
         self.log("early_stop_on", out["loss"])
         return out
 
-    def validation_epoch_end(self, outputs):
-        outs = torch.stack([x["x"] for x in outputs]).mean()
-        self.log("val_acc", outs)
+    def on_validation_epoch_end(self):
+        self.log("val_acc", torch.tensor(1.23))
 
 
 def mock_training_epoch_loop(trainer):
@@ -214,9 +213,8 @@ def validation_step(self, batch, batch_idx):
             self.log("val_log", log_value)
             return super().validation_step(batch, batch_idx)
 
-        def validation_epoch_end(self, outputs):
+        def on_validation_epoch_end(self):
             self.val_loop_count += 1
-            super().validation_epoch_end(outputs)
             self.scores.append(self.trainer.logged_metrics[monitor])
 
         def configure_optimizers(self):
@@ -829,7 +827,7 @@ def test_checkpointing_with_nan_as_first(tmpdir, mode):
     monitor += [5, 7, 8] if mode == "max" else [8, 7, 5]
 
     class CurrentModel(LogInTwoMethods):
-        def validation_epoch_end(self, outputs):
+        def on_validation_epoch_end(self):
             val_loss = monitor[self.current_epoch]
             self.log("abc", val_loss)
 
@@ -863,7 +861,6 @@ def validation_step(self, batch, batch_idx):
             self.log("val_loss", loss)
 
     model = ExtendedBoringModel()
-    model.validation_epoch_end = None
     trainer_kwargs = {
         "max_epochs": 1,
         "limit_train_batches": 2,
@@ -901,9 +898,6 @@ def validation_step(self, batch, batch_idx):
             self.log("val_loss", loss)
             return {"val_loss": loss}
 
-        def validation_epoch_end(self, *_):
-            ...
-
     def assert_trainer_init(trainer):
         assert trainer.global_step == 0
         assert trainer.current_epoch == 0

tests/tests_pytorch/checkpointing/test_trainer_checkpoint.py

@@ -41,7 +41,6 @@ def validation_step(self, batch, batch_idx):
             self.log("val_loss", loss, on_epoch=True, prog_bar=True)
 
     model = ExtendedBoringModel()
-    model.validation_epoch_end = None
     trainer = Trainer(
         default_root_dir=tmpdir,
         max_epochs=1,

tests/tests_pytorch/core/test_datamodules.py

@@ -163,10 +163,8 @@ def test_train_loop_only(tmpdir):
 
     model.validation_step = None
     model.validation_step_end = None
-    model.validation_epoch_end = None
     model.test_step = None
     model.test_step_end = None
-    model.test_epoch_end = None
 
     trainer = Trainer(default_root_dir=tmpdir, max_epochs=1, enable_model_summary=False)
 
@@ -185,7 +183,6 @@ def test_train_val_loop_only(tmpdir):
 
     model.validation_step = None
     model.validation_step_end = None
-    model.validation_epoch_end = None
 
     trainer = Trainer(default_root_dir=tmpdir, max_epochs=1, enable_model_summary=False)
 
@@ -278,10 +275,8 @@ def train_dataloader(self):
 
     model.validation_step = None
     model.validation_step_end = None
-    model.validation_epoch_end = None
     model.test_step = None
     model.test_step_end = None
-    model.test_epoch_end = None
 
     trainer = Trainer(default_root_dir=tmpdir, max_epochs=3, limit_train_batches=2, reload_dataloaders_every_n_epochs=2)
     trainer.fit(model, dm)