Merge branch 'master' into fix/setup_ddp_hook

Author: SeanNaren

CHANGELOG.md

@@ -40,6 +40,10 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added no return warning to predict ([#6139](https://github.com/PyTorchLightning/pytorch-lightning/pull/6139))
 
 
+- Added `outputs` parameter to callback's `on_validation_epoch_end` & `on_test_epoch_end` hooks ([#6120](https://github.com/PyTorchLightning/pytorch-lightning/pull/6120))
+
+
+
 ### Changed
 
 - Renamed `pytorch_lightning.callbacks.swa` to `pytorch_lightning.callbacks.stochastic_weight_avg` ([#6259](https://github.com/PyTorchLightning/pytorch-lightning/pull/6259))
@@ -68,6 +72,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Deprecated metrics in favor of `torchmetrics` ([#6505](https://github.com/PyTorchLightning/pytorch-lightning/pull/6505),
 
     [#6530](https://github.com/PyTorchLightning/pytorch-lightning/pull/6530),
+  
+    [#6547](https://github.com/PyTorchLightning/pytorch-lightning/pull/6547),
 
 )
 

docs/source/benchmarking/performance.rst

@@ -181,3 +181,19 @@ Most UNIX-based operating systems provide direct access to tmpfs through a mount
     .. code-block:: python
 
         datamodule = MyDataModule(data_root="/dev/shm/my_data")
+
+
+Zero Grad ``set_to_none=True``
+------------------------------
+
+In order to modestly improve performance, once can override :meth:`~pytorch_lightning.core.lightning.LightningModule.optimizer_zero_grad`.
+
+For a more detailed explanation of pros / cons of this technique,
+read `this <https://pytorch.org/docs/master/optim.html#torch.optim.Optimizer.zero_grad>`_ documentation by the PyTorch team.
+
+.. testcode::
+
+    class Model(LightningModule):
+
+        def optimizer_zero_grad(self, epoch, batch_idx, optimizer, optimizer_idx):
+            optimizer.zero_grad(set_to_none=True)

pytorch_lightning/callbacks/base.py

@@ -17,7 +17,7 @@
 """
 
 import abc
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional
 
 from pytorch_lightning.core.lightning import LightningModule
 
@@ -81,23 +81,23 @@ def on_train_epoch_start(self, trainer, pl_module: LightningModule) -> None:
         """Called when the train epoch begins."""
         pass
 
-    def on_train_epoch_end(self, trainer, pl_module: LightningModule, outputs: Any) -> None:
+    def on_train_epoch_end(self, trainer, pl_module: LightningModule, outputs: List[Any]) -> None:
         """Called when the train epoch ends."""
         pass
 
     def on_validation_epoch_start(self, trainer, pl_module: LightningModule) -> None:
         """Called when the val epoch begins."""
         pass
 
-    def on_validation_epoch_end(self, trainer, pl_module: LightningModule) -> None:
+    def on_validation_epoch_end(self, trainer, pl_module: LightningModule, outputs: List[Any]) -> None:
         """Called when the val epoch ends."""
         pass
 
     def on_test_epoch_start(self, trainer, pl_module: LightningModule) -> None:
         """Called when the test epoch begins."""
         pass
 
-    def on_test_epoch_end(self, trainer, pl_module: LightningModule) -> None:
+    def on_test_epoch_end(self, trainer, pl_module: LightningModule, outputs: List[Any]) -> None:
         """Called when the test epoch ends."""
         pass
 

pytorch_lightning/core/hooks.py

@@ -240,7 +240,7 @@ def on_train_epoch_start(self) -> None:
         """
         # do something when the epoch starts
 
-    def on_train_epoch_end(self, outputs) -> None:
+    def on_train_epoch_end(self, outputs: List[Any]) -> None:
         """
         Called in the training loop at the very end of the epoch.
         """
@@ -252,7 +252,7 @@ def on_validation_epoch_start(self) -> None:
         """
         # do something when the epoch starts
 
-    def on_validation_epoch_end(self) -> None:
+    def on_validation_epoch_end(self, outputs: List[Any]) -> None:
         """
         Called in the validation loop at the very end of the epoch.
         """
@@ -264,7 +264,7 @@ def on_test_epoch_start(self) -> None:
         """
         # do something when the epoch starts
 
-    def on_test_epoch_end(self) -> None:
+    def on_test_epoch_end(self, outputs: List[Any]) -> None:
         """
         Called in the test loop at the very end of the epoch.
         """

pytorch_lightning/core/step_result.py

@@ -20,8 +20,8 @@
 
 import torch
 from torch import Tensor
+from torchmetrics import Metric
 
-from pytorch_lightning.metrics import Metric
 from pytorch_lightning.utilities.distributed import sync_ddp_if_available
 
 

pytorch_lightning/metrics/classification/accuracy.py

@@ -14,9 +14,9 @@
 from typing import Any, Callable, Optional
 
 import torch
+from torchmetrics import Metric
 
 from pytorch_lightning.metrics.functional.accuracy import _accuracy_compute, _accuracy_update
-from pytorch_lightning.metrics.metric import Metric
 
 
 class Accuracy(Metric):

pytorch_lightning/metrics/classification/auc.py

@@ -14,9 +14,9 @@
 from typing import Any, Callable, Optional
 
 import torch
+from torchmetrics import Metric
 
 from pytorch_lightning.metrics.functional.auc import _auc_compute, _auc_update
-from pytorch_lightning.metrics.metric import Metric
 from pytorch_lightning.utilities import rank_zero_warn
 
 

pytorch_lightning/metrics/classification/auroc.py

@@ -15,9 +15,9 @@
 from typing import Any, Callable, Optional
 
 import torch
+from torchmetrics import Metric
 
 from pytorch_lightning.metrics.functional.auroc import _auroc_compute, _auroc_update
-from pytorch_lightning.metrics.metric import Metric
 from pytorch_lightning.utilities import rank_zero_warn
 
 

pytorch_lightning/metrics/classification/average_precision.py

@@ -14,9 +14,9 @@
 from typing import Any, List, Optional, Union
 
 import torch
+from torchmetrics import Metric
 
 from pytorch_lightning.metrics.functional.average_precision import _average_precision_compute, _average_precision_update
-from pytorch_lightning.metrics.metric import Metric
 from pytorch_lightning.utilities import rank_zero_warn
 
 

pytorch_lightning/metrics/classification/confusion_matrix.py

@@ -14,9 +14,9 @@
 from typing import Any, Optional
 
 import torch
+from torchmetrics import Metric
 
 from pytorch_lightning.metrics.functional.confusion_matrix import _confusion_matrix_compute, _confusion_matrix_update
-from pytorch_lightning.metrics.metric import Metric
 
 
 class ConfusionMatrix(Metric):