From 70b8cb626304479dd257d8e31fe2649d453f7770 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Carlos=20Mochol=C3=AD?= <carlossmocholi@gmail.com>
Date: Mon, 30 Jan 2023 19:06:20 +0100
Subject: [PATCH] Update docs

---
 .../accelerators/accelerator_prepare.rst      |  16 ++-
 .../common/lightning_module.rst               | 135 ++++++------------
 docs/source-pytorch/extensions/logging.rst    |   4 +-
 .../model/manual_optimization.rst             |   2 +-
 docs/source-pytorch/starter/style_guide.rst   |   6 +-
 .../visualize/logging_advanced.rst            |   2 +-
 examples/pl_ipu/mnist_sample.py               |  20 +--
 7 files changed, 73 insertions(+), 112 deletions(-)

diff --git a/docs/source-pytorch/accelerators/accelerator_prepare.rst b/docs/source-pytorch/accelerators/accelerator_prepare.rst
index f736c57472d2c..f1da6867a0eee 100644
--- a/docs/source-pytorch/accelerators/accelerator_prepare.rst
+++ b/docs/source-pytorch/accelerators/accelerator_prepare.rst
@@ -105,19 +105,27 @@ Note if you use any built in metrics or custom metrics that use `TorchMetrics <h
 
 It is possible to perform some computation manually and log the reduced result on rank 0 as follows:
 
-.. testcode::
+.. code-block:: python
+
+    def __init__(self):
+        super().__init__()
+        self.outputs = []
+
 
     def test_step(self, batch, batch_idx):
         x, y = batch
         tensors = self(x)
+        self.outputs.append(tensors)
         return tensors
 
 
-    def test_epoch_end(self, outputs):
-        mean = torch.mean(self.all_gather(outputs))
+    def on_test_epoch_end(self):
+        mean = torch.mean(self.all_gather(self.outputs))
+        self.outputs.clear()  # free memory
 
         # When logging only on rank 0, don't forget to add
-        # ``rank_zero_only=True`` to avoid deadlocks on synchronization.
+        # `rank_zero_only=True` to avoid deadlocks on synchronization.
+        # caveat: monitoring this is unimplemented. see https://github.com/Lightning-AI/lightning/issues/15852
         if self.trainer.is_global_zero:
             self.log("my_reduced_metric", mean, rank_zero_only=True)
 
diff --git a/docs/source-pytorch/common/lightning_module.rst b/docs/source-pytorch/common/lightning_module.rst
index b365edf819768..e6b6df4be19b8 100644
--- a/docs/source-pytorch/common/lightning_module.rst
+++ b/docs/source-pytorch/common/lightning_module.rst
@@ -182,10 +182,8 @@ Under the hood, Lightning does the following (pseudocode):
     model.train()
     torch.set_grad_enabled(True)
 
-    outs = []
     for batch_idx, batch in enumerate(train_dataloader):
         loss = training_step(batch, batch_idx)
-        outs.append(loss.detach())
 
         # clear gradients
         optimizer.zero_grad()
@@ -214,7 +212,7 @@ If you want to calculate epoch-level metrics and log them, use :meth:`~pytorch_l
          self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)
          return loss
 
-The :meth:`~pytorch_lightning.core.module.LightningModule.log` object automatically reduces the
+The :meth:`~pytorch_lightning.core.module.LightningModule.log` method automatically reduces the
 requested metrics across a complete epoch and devices. Here's the pseudocode of what it does under the hood:
 
 .. code-block:: python
@@ -223,59 +221,44 @@ requested metrics across a complete epoch and devices. Here's the pseudocode of
     for batch_idx, batch in enumerate(train_dataloader):
         # forward
         loss = training_step(batch, batch_idx)
-        outs.append(loss)
+        outs.append(loss.detach())
 
         # clear gradients
         optimizer.zero_grad()
-
         # backward
         loss.backward()
-
         # update parameters
         optimizer.step()
 
-    epoch_metric = torch.mean(torch.stack([x for x in outs]))
+    # note: in reality, we do this incrementally, instead of keeping all outputs in memory
+    epoch_metric = torch.mean(torch.stack(outs))
 
 Train Epoch-level Operations
 ============================
 
 If you need to do something with all the outputs of each :meth:`~pytorch_lightning.core.module.LightningModule.training_step`,
-override the :meth:`~pytorch_lightning.core.module.LightningModule.training_epoch_end` method.
-
-.. code-block:: python
-
-     def training_step(self, batch, batch_idx):
-         x, y = batch
-         y_hat = self.model(x)
-         loss = F.cross_entropy(y_hat, y)
-         preds = ...
-         return {"loss": loss, "other_stuff": preds}
-
-
-     def training_epoch_end(self, training_step_outputs):
-         all_preds = torch.stack(training_step_outputs)
-         ...
-
-The matching pseudocode is:
+override the :meth:`~pytorch_lightning.core.module.LightningModule.on_training_epoch_end` method.
 
 .. code-block:: python
 
-    outs = []
-    for batch_idx, batch in enumerate(train_dataloader):
-        # forward
-        loss = training_step(batch, batch_idx)
-        outs.append(loss)
+    def __init__(self):
+        super().__init__()
+        self.training_step_outputs = []
 
-        # clear gradients
-        optimizer.zero_grad()
 
-        # backward
-        loss.backward()
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = self.model(x)
+        loss = F.cross_entropy(y_hat, y)
+        preds = ...
+        self.training_step_outputs.append(preds)
+        return loss
 
-        # update parameters
-        optimizer.step()
 
-    training_epoch_end(outs)
+    def on_train_epoch_end(self):
+        all_preds = torch.stack(self.training_step_outputs)
+        self.training_step_outputs.clear()  # free memory
+        ...
 
 Training with DataParallel
 ==========================
@@ -309,15 +292,10 @@ method which will have outputs from all the devices and you can accumulate to ge
          return (losses[0] + losses[1]) / 2
 
 
-     def training_epoch_end(self, training_step_outputs):
-         for out in training_step_outputs:
-             ...
-
 Here is the Lightning training pseudo-code for DP:
 
 .. code-block:: python
 
-    outs = []
     for batch_idx, train_batch in enumerate(train_dataloader):
         batches = split_batch(train_batch)
         dp_outs = []
@@ -327,12 +305,7 @@ Here is the Lightning training pseudo-code for DP:
             dp_outs.append(dp_out)
 
         # 2
-        out = training_step_end(dp_outs)
-        outs.append(out)
-
-    # do something with the outputs for all batches
-    # 3
-    training_epoch_end(outs)
+        training_step_end(dp_outs)
 
 ------------------
 
@@ -400,21 +373,30 @@ Validation Epoch-level Metrics
 ==============================
 
 If you need to do something with all the outputs of each :meth:`~pytorch_lightning.core.module.LightningModule.validation_step`,
-override the :meth:`~pytorch_lightning.core.module.LightningModule.validation_epoch_end` method. Note that this method is called before :meth:`~pytorch_lightning.core.module.LightningModule.training_epoch_end`.
+override the :meth:`~pytorch_lightning.core.module.LightningModule.on_validation_epoch_end` method.
+Note that this method is called before :meth:`~pytorch_lightning.core.module.LightningModule.on_train_epoch_end`.
 
 .. code-block:: python
 
-     def validation_step(self, batch, batch_idx):
-         x, y = batch
-         y_hat = self.model(x)
-         loss = F.cross_entropy(y_hat, y)
-         pred = ...
-         return pred
+    def __init__(self):
+        super().__init__()
+        self.validation_step_outputs = []
 
 
-     def validation_epoch_end(self, validation_step_outputs):
-         all_preds = torch.stack(validation_step_outputs)
-         ...
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = self.model(x)
+        loss = F.cross_entropy(y_hat, y)
+        pred = ...
+        self.validation_step_outputs.append(pred)
+        return pred
+
+
+    def on_validation_epoch_end(self):
+        all_preds = torch.stack(self.validation_step_outputs)
+        self.validation_step_outputs.clear()  # free memory
+        ...
+
 
 Validating with DataParallel
 ============================
@@ -448,15 +430,10 @@ method which will have outputs from all the devices and you can accumulate to ge
          return (losses[0] + losses[1]) / 2
 
 
-     def validation_epoch_end(self, validation_step_outputs):
-         for out in validation_step_outputs:
-             ...
-
 Here is the Lightning validation pseudo-code for DP:
 
 .. code-block:: python
 
-    outs = []
     for batch in dataloader:
         batches = split_batch(batch)
         dp_outs = []
@@ -466,12 +443,7 @@ Here is the Lightning validation pseudo-code for DP:
             dp_outs.append(dp_out)
 
         # 2
-        out = validation_step_end(dp_outs)
-        outs.append(out)
-
-    # do something with the outputs for all batches
-    # 3
-    validation_epoch_end(outs)
+        validation_step_end(dp_outs)
 
 ----------------
 
@@ -924,12 +896,6 @@ test_step_end
 .. automethod:: pytorch_lightning.core.module.LightningModule.test_step_end
     :noindex:
 
-test_epoch_end
-~~~~~~~~~~~~~~
-
-.. automethod:: pytorch_lightning.core.module.LightningModule.test_epoch_end
-    :noindex:
-
 to_onnx
 ~~~~~~~
 
@@ -954,11 +920,6 @@ training_step_end
 .. automethod:: pytorch_lightning.core.module.LightningModule.training_step_end
     :noindex:
 
-training_epoch_end
-~~~~~~~~~~~~~~~~~~
-.. automethod:: pytorch_lightning.core.module.LightningModule.training_epoch_end
-    :noindex:
-
 unfreeze
 ~~~~~~~~
 
@@ -983,12 +944,6 @@ validation_step_end
 .. automethod:: pytorch_lightning.core.module.LightningModule.validation_step_end
     :noindex:
 
-validation_epoch_end
-~~~~~~~~~~~~~~~~~~~~
-
-.. automethod:: pytorch_lightning.core.module.LightningModule.validation_epoch_end
-    :noindex:
-
 -----------
 
 Properties
@@ -1247,7 +1202,8 @@ for more information.
             transfer_batch_to_device()
             on_after_batch_transfer()
 
-            training_step()
+            out = training_step()
+            training_step_end(out)
 
             on_before_zero_grad()
             optimizer_zero_grad()
@@ -1264,8 +1220,6 @@ for more information.
 
             if should_check_val:
                 val_loop()
-        # end training epoch
-        training_epoch_end()
 
         on_train_epoch_end()
 
@@ -1277,7 +1231,6 @@ for more information.
         on_validation_start()
         on_validation_epoch_start()
 
-        val_outs = []
         for batch_idx, batch in enumerate(val_dataloader()):
             on_validation_batch_start(batch, batch_idx)
 
@@ -1286,11 +1239,9 @@ for more information.
             batch = on_after_batch_transfer(batch)
 
             out = validation_step(batch, batch_idx)
+            out = validation_step_end(out)
 
             on_validation_batch_end(batch, batch_idx)
-            val_outs.append(out)
-
-        validation_epoch_end(val_outs)
 
         on_validation_epoch_end()
         on_validation_end()
diff --git a/docs/source-pytorch/extensions/logging.rst b/docs/source-pytorch/extensions/logging.rst
index 70b0f7c09e346..a0c36ba4807ad 100644
--- a/docs/source-pytorch/extensions/logging.rst
+++ b/docs/source-pytorch/extensions/logging.rst
@@ -152,7 +152,7 @@ The :meth:`~pytorch_lightning.core.module.LightningModule.log` method has a few
    * - Hook
      - on_step
      - on_epoch
-   * - on_train_start, on_train_epoch_start, on_train_epoch_end, training_epoch_end
+   * - on_train_start, on_train_epoch_start, on_train_epoch_end
      - False
      - True
    * - on_before_backward, on_after_backward, on_before_optimizer_step, on_before_zero_grad
@@ -161,7 +161,7 @@ The :meth:`~pytorch_lightning.core.module.LightningModule.log` method has a few
    * - on_train_batch_start, on_train_batch_end, training_step, training_step_end
      - True
      - False
-   * - on_validation_start, on_validation_epoch_start, on_validation_epoch_end, validation_epoch_end
+   * - on_validation_start, on_validation_epoch_start, on_validation_epoch_end
      - False
      - True
    * - on_validation_batch_start, on_validation_batch_end, validation_step, validation_step_end
diff --git a/docs/source-pytorch/model/manual_optimization.rst b/docs/source-pytorch/model/manual_optimization.rst
index aad1c1aa00263..498521ff873de 100644
--- a/docs/source-pytorch/model/manual_optimization.rst
+++ b/docs/source-pytorch/model/manual_optimization.rst
@@ -278,7 +278,7 @@ If you want to call schedulers that require a metric value after each epoch, con
         self.automatic_optimization = False
 
 
-    def training_epoch_end(self, outputs):
+    def on_train_epoch_end(self):
         sch = self.lr_schedulers()
 
         # If the selected scheduler is a ReduceLROnPlateau scheduler.
diff --git a/docs/source-pytorch/starter/style_guide.rst b/docs/source-pytorch/starter/style_guide.rst
index f4f8be8e39a07..77823f798ed46 100644
--- a/docs/source-pytorch/starter/style_guide.rst
+++ b/docs/source-pytorch/starter/style_guide.rst
@@ -149,19 +149,19 @@ In practice, the code looks like this:
 
         def training_step_end(...):
 
-        def training_epoch_end(...):
+        def on_train_epoch_end(...):
 
         def validation_step(...):
 
         def validation_step_end(...):
 
-        def validation_epoch_end(...):
+        def on_validation_epoch_end(...):
 
         def test_step(...):
 
         def test_step_end(...):
 
-        def test_epoch_end(...):
+        def on_test_epoch_end(...):
 
         def configure_optimizers(...):
 
diff --git a/docs/source-pytorch/visualize/logging_advanced.rst b/docs/source-pytorch/visualize/logging_advanced.rst
index 788deffdb3fea..bdf3f2f5e1997 100644
--- a/docs/source-pytorch/visualize/logging_advanced.rst
+++ b/docs/source-pytorch/visualize/logging_advanced.rst
@@ -358,7 +358,7 @@ In LightningModule
    * - on_after_backward, on_before_backward, on_before_optimizer_step, optimizer_step, configure_gradient_clipping, on_before_zero_grad, training_step, training_step_end
      - True
      - False
-   * - training_epoch_end, test_epoch_end, test_step, test_step_end, validation_epoch_end, validation_step, validation_step_end
+   * - test_step, test_step_end, validation_step, validation_step_end
      - False
      - True
 
diff --git a/examples/pl_ipu/mnist_sample.py b/examples/pl_ipu/mnist_sample.py
index 56835afec5f64..e0d430579f5fe 100644
--- a/examples/pl_ipu/mnist_sample.py
+++ b/examples/pl_ipu/mnist_sample.py
@@ -27,6 +27,9 @@ def __init__(self, hidden_dim: int = 128, learning_rate: float = 0.0001):
         self.l1 = torch.nn.Linear(28 * 28, self.hparams.hidden_dim)
         self.l2 = torch.nn.Linear(self.hparams.hidden_dim, 10)
 
+        self.val_outptus = []
+        self.test_outputs = []
+
     def forward(self, x):
         x = x.view(x.size(0), -1)
         x = torch.relu(self.l1(x))
@@ -42,16 +45,15 @@ def training_step(self, batch, batch_idx):
     def validation_step(self, batch, batch_idx):
         x, y = batch
         probs = self(x)
-        # we currently return the accuracy as the validation_step/test_step is run on the IPU devices.
-        # Outputs from the step functions are sent to the host device, where we calculate the metrics in
-        # validation_epoch_end and test_epoch_end for the test_step.
         acc = self.accuracy(probs, y)
+        self.val_outputs.append(acc)
         return acc
 
     def test_step(self, batch, batch_idx):
         x, y = batch
         logits = self(x)
         acc = self.accuracy(logits, y)
+        self.test_outputs.append(acc)
         return acc
 
     def accuracy(self, logits, y):
@@ -61,13 +63,15 @@ def accuracy(self, logits, y):
         acc = torch.sum(torch.eq(torch.argmax(logits, -1), y).to(torch.float32)) / len(y)
         return acc
 
-    def validation_epoch_end(self, outputs) -> None:
+    def on_validation_epoch_end(self) -> None:
         # since the training step/validation step and test step are run on the IPU device
         # we must log the average loss outside the step functions.
-        self.log("val_acc", torch.stack(outputs).mean(), prog_bar=True)
+        self.log("val_acc", torch.stack(self.val_outptus).mean(), prog_bar=True)
+        self.val_outptus.clear()
 
-    def test_epoch_end(self, outputs) -> None:
-        self.log("test_acc", torch.stack(outputs).mean())
+    def on_test_epoch_end(self) -> None:
+        self.log("test_acc", torch.stack(self.test_outputs).mean())
+        self.test_outputs.clear()
 
     def configure_optimizers(self):
         return torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
@@ -75,9 +79,7 @@ def configure_optimizers(self):
 
 if __name__ == "__main__":
     dm = MNISTDataModule(batch_size=32)
-
     model = LitClassifier()
-
     trainer = Trainer(max_epochs=2, accelerator="ipu", devices=8)
 
     trainer.fit(model, datamodule=dm)