[feat] 1/2 Add trainer.predict (#5579)

* start adding predict

* add predict

* resolve test

* add predict

* remove limit_predict

* update

* add test for predict

* typo

* update on comments

* remove predict_step

* update ddp_shareded

* check ddp_sharded

* resolve on comments

* resolve isort

* update dp

* add test dp 1 gpu

* made default forward

* resolve path

* resolve bug

* update on comments

* resolve doc

* resolve bug

* update

* resolve bug

* update on comments

* resolve pep8

* update test doc

* update on comments

* solve special tests

* resolve bug

* resolve flake8

* Update pytorch_lightning/callbacks/progress.py

Co-authored-by: Carlos Mocholí <carlossmocholi@gmail.com>

* Update pytorch_lightning/trainer/trainer.py

Co-authored-by: Carlos Mocholí <carlossmocholi@gmail.com>

* add predict to LightningModule

* missing predict

* typo

* rename is_prediction to _predicting

* add

* update

* update

* update doc

Co-authored-by: Carlos Mocholí <carlossmocholi@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Author: 3 people

CHANGELOG.md

@@ -11,6 +11,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Add support for summarized model total params size in megabytes ([#5590](https://github.com/PyTorchLightning/pytorch-lightning/pull/5590))
 
+
 - Add Support for multiple train loaders ([#1959](https://github.com/PyTorchLightning/pytorch-lightning/pull/1959))
 
 
@@ -68,6 +69,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added compositional metrics ([#5464](https://github.com/PyTorchLightning/pytorch-lightning/pull/5464))
 
 
+- Added Trainer method `predict(...)` for high performence predictions ([#5579](https://github.com/PyTorchLightning/pytorch-lightning/pull/5579))
+
+
 - Added AUC/AUROC class interface ([#5479](https://github.com/PyTorchLightning/pytorch-lightning/pull/5479))
 
 
@@ -120,7 +124,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 
 - Removed deprecated `TrainResult` ([#5323](https://github.com/PyTorchLightning/pytorch-lightning/pull/5323))
-  
+
 
 - Removed deprecated `EvalResult` ([#5633](https://github.com/PyTorchLightning/pytorch-lightning/pull/5633))
 
@@ -155,7 +159,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Fixed `transfer_batch_to_device` for DDP with `len(devices_ids) == 1` ([#5195](https://github.com/PyTorchLightning/pytorch-lightning/pull/5195))
 - Logging only on `not should_accumulate()` during training ([#5417](https://github.com/PyTorchLightning/pytorch-lightning/pull/5417))
-- Resolve interpolation bug with Hydra ([#5406](https://github.com/PyTorchLightning/pytorch-lightning/pull/5406)) 
+- Resolve interpolation bug with Hydra ([#5406](https://github.com/PyTorchLightning/pytorch-lightning/pull/5406))
 - Check environ before selecting a seed to prevent warning message ([#4743](https://github.com/PyTorchLightning/pytorch-lightning/pull/4743))
 
 

docs/source/starter/introduction_guide.rst

@@ -881,8 +881,30 @@ Or maybe we have a model that we use to do generation
     z = sample_noise()
     generated_imgs = model(z)
 
-How you split up what goes in ``forward`` vs ``training_step`` depends on how you want to use this model for
+
+To perform inference at scale, it is possible to use ``trainer.predict`` with LightningModule ``predict`` function
+By default, LightningModule ``predict`` calls forward, but it can be overriden to add any processing logic.
+
+.. code-block:: python
+
+    class LitMNISTDreamer(LightningModule):
+
+        def forward(self, z):
+            imgs = self.decoder(z)
+            return imgs
+
+        def predict(self, batch, batch_idx: int , dataloader_idx: int = None):
+            return self(batch)
+
+
+    model = LitMNISTDreamer()
+    trainer.predict(model, datamodule)
+
+
+How you split up what goes in ``forward`` vs ``training_step`` vs ``predict`` depends on how you want to use this model for
 prediction.
+However, we recommend ``forward`` to contain only tensor operation with your model, ``training_step`` to encapsulate ``forward`` logic with logging,
+metrics and loss computation and ``predict`` to encapsulate ``forward`` with preprocess, postprocess functions.
 
 ----------------
 

pytorch_lightning/accelerators/legacy/cpu_accelerator.py

@@ -77,6 +77,9 @@ def validation_step(self, args):
     def test_step(self, args):
         return self._step(self.trainer.model.test_step, args)
 
+    def predict(self, args):
+        return self._step(self.trainer.model.predict, args)
+
     def sync_tensor(self,
                     tensor: Union[torch.Tensor],
                     group: Optional[Any] = None,

pytorch_lightning/accelerators/legacy/ddp2_accelerator.py

@@ -66,6 +66,9 @@ def validation_step(self, args):
     def test_step(self, args):
         return self._step(args)
 
+    def predict(self, args):
+        return self._step(args)
+
     def _step(self, args):
         args = self.ddp_plugin.on_before_forward(self.trainer.get_model(), *args)
         if self.trainer.amp_backend == AMPType.NATIVE:

pytorch_lightning/accelerators/legacy/ddp_accelerator.py

@@ -164,6 +164,9 @@ def validation_step(self, args):
     def test_step(self, args):
         return self._step(args)
 
+    def predict(self, args):
+        return self._step(args)
+
     def _step(self, args):
         args = self.ddp_plugin.on_before_forward(self.trainer.get_model(), *args)
         if self.trainer.amp_backend == AMPType.NATIVE:

pytorch_lightning/accelerators/legacy/ddp_cpu_spawn_accelerator.py

@@ -178,6 +178,9 @@ def validation_step(self, args):
     def test_step(self, args):
         return self._step(args)
 
+    def predict(self, args):
+        return self._step(args)
+
     def _step(self, args):
         args = self.ddp_plugin.on_before_forward(self.trainer.get_model(), *args)
         if self.trainer.amp_backend == AMPType.NATIVE:

pytorch_lightning/accelerators/legacy/ddp_hpc_accelerator.py

@@ -83,6 +83,9 @@ def validation_step(self, args):
     def test_step(self, args):
         return self._step(args)
 
+    def predict(self, args):
+        return self._step(args)
+
     def _step(self, args):
         args = self.ddp_plugin.on_before_forward(self.trainer.get_model(), *args)
         if self.trainer.amp_backend == AMPType.NATIVE:

pytorch_lightning/accelerators/legacy/ddp_spawn_accelerator.py

@@ -212,6 +212,9 @@ def validation_step(self, args):
     def test_step(self, args):
         return self._step(args)
 
+    def predict(self, args):
+        return self._step(args)
+
     def _step(self, args):
         args = self.ddp_plugin.on_before_forward(self.trainer.get_model(), *args)
         if self.trainer.amp_backend == AMPType.NATIVE:

pytorch_lightning/accelerators/legacy/dp_accelerator.py

@@ -132,6 +132,9 @@ def validation_step(self, args):
     def test_step(self, args):
         return self._step(args)
 
+    def predict(self, args):
+        return self._step(args)
+
     def training_step_end(self, output):
         if isinstance(output, Result):
             output.dp_reduce()

pytorch_lightning/accelerators/legacy/gpu_accelerator.py

@@ -85,6 +85,9 @@ def validation_step(self, args):
     def test_step(self, args):
         return self._step(self.trainer.model.test_step, args)
 
+    def predict(self, args):
+        return self._step(self.trainer.model.predict, args)
+
     def to_device(self, batch):
         gpu_id = 0
         if isinstance(self.trainer.data_parallel_device_ids, list):