[feat] 1/2 Add trainer.predict

pytorch_lightning/accelerators/cpu_accelerator.py

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

pytorch_lightning/accelerators/ddp2_accelerator.py

@@ -66,6 +66,9 @@ def validation_step(self, args):
     def test_step(self, args):
         return self._step(args)
 
+    def forward(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/ddp_accelerator.py

@@ -164,6 +164,9 @@ def validation_step(self, args):
     def test_step(self, args):
         return self._step(args)
 
+    def forward(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/ddp_cpu_spawn_accelerator.py

@@ -180,6 +180,9 @@ def validation_step(self, args):
     def test_step(self, args):
         return self._step(args)
 
+    def forward(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/ddp_hpc_accelerator.py

@@ -83,6 +83,9 @@ def validation_step(self, args):
     def test_step(self, args):
         return self._step(args)
 
+    def forward(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/ddp_spawn_accelerator.py

@@ -214,6 +214,9 @@ def validation_step(self, args):
     def test_step(self, args):
         return self._step(args)
 
+    def forward(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/dp_accelerator.py

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

pytorch_lightning/accelerators/gpu_accelerator.py

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

pytorch_lightning/accelerators/horovod_accelerator.py

@@ -136,6 +136,9 @@ def validation_step(self, args):
     def test_step(self, args):
         return self._step(self.trainer.model.test_step, args)
 
+    def forward(self, args):
+        return self._step(self.trainer.model.forward, args)
+
     def backward(self, closure_loss, optimizer, opt_idx, *args, **kwargs):
         super().backward(closure_loss, optimizer, opt_idx, *args, **kwargs)
         optimizer.synchronize()

pytorch_lightning/accelerators/tpu_accelerator.py

@@ -159,6 +159,9 @@ def validation_step(self, args):
     def test_step(self, args):
         return self._step(self.trainer.model.test_step, args)
 
+    def forward(self, args):
+        return self._step(self.trainer.model.forward, args)
+
     def process_dataloader(self, dataloader):
         device = xm.xla_device(self.trainer.tpu_id)
         dataloader = xla_pl.ParallelLoader(dataloader, [device])

pytorch_lightning/callbacks/progress.py

@@ -291,10 +291,13 @@ def init_validation_tqdm(self) -> tqdm:
         )
         return bar
 
-    def init_test_tqdm(self) -> tqdm:
+    def init_test_tqdm(self, trainer=None) -> tqdm:
         """ Override this to customize the tqdm bar for testing. """
+        desc = "Testing"
+        if trainer is not None and getattr(trainer, "is_predicting", False):
+            desc = "Predicting"
         bar = tqdm(
-            desc='Testing',
+            desc=desc,
             position=(2 * self.process_position),
             disable=self.is_disabled,
             leave=True,
@@ -361,7 +364,7 @@ def on_train_end(self, trainer, pl_module):
 
     def on_test_start(self, trainer, pl_module):
         super().on_test_start(trainer, pl_module)
-        self.test_progress_bar = self.init_test_tqdm()
+        self.test_progress_bar = self.init_test_tqdm(trainer=trainer)
         self.test_progress_bar.total = convert_inf(self.total_test_batches)
 
     def on_test_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):

pytorch_lightning/core/lightning.py

@@ -14,16 +14,16 @@
 
 """nn.Module with additional great features."""
 
-from abc import ABC
-from argparse import Namespace
 import collections
 import copy
-from functools import partial
 import inspect
 import os
-from pathlib import Path
 import re
 import tempfile
+from abc import ABC
+from argparse import Namespace
+from functools import partial
+from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import torch
@@ -66,6 +66,7 @@ class LightningModule(
         "on_gpu",
         "current_epoch",
         "global_step",
+        "running_stage",
     ] + DeviceDtypeModuleMixin.__jit_unused_properties__
 
     def __init__(self, *args, **kwargs):
@@ -102,6 +103,7 @@ def __init__(self, *args, **kwargs):
         self._running_manual_backward = False
         self._current_hook_fx_name = None
         self._current_dataloader_idx = None
+        self.running_stage = None
 
     def optimizers(self):
         opts = self.trainer.optimizers

pytorch_lightning/overrides/data_parallel.py

@@ -28,6 +28,7 @@
 
 from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.core.step_result import Result
+from pytorch_lightning.trainer.states import RunningStage
 from pytorch_lightning.utilities.warnings import WarningCache
 
 
@@ -78,14 +79,22 @@ def forward(self, *inputs, **kwargs):
                                    "them on device: {}".format(self.src_device_obj, t.device))
 
         inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
+
         if len(self.device_ids) == 1:
-            # lightning
-            if self.module.training:
+
+            running_stage = self.module.running_stage
+
+            if running_stage == RunningStage.TRAINING:
                 return self.module.training_step(*inputs[0], **kwargs[0])
-            if self.module.testing:
+
+            elif running_stage == RunningStage.TESTING:
                 return self.module.test_step(*inputs[0], **kwargs[0])
 
-            return self.module.validation_step(*inputs[0], **kwargs[0])
+            elif running_stage == RunningStage.EVALUATING:
+                return self.module.validation_step(*inputs[0], **kwargs[0])
+
+            else:
+                return self.module(*inputs[0], **kwargs[0])
 
         replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
         outputs = self.parallel_apply(replicas, inputs, kwargs)
@@ -187,15 +196,24 @@ def __init__(self, pl_module: LightningModule):
         self.module = pl_module
 
     def forward(self, *inputs, **kwargs):
-        if self.module.training:
+
+        running_stage = self.module.running_stage
+
+        if running_stage == RunningStage.TRAINING:
             output = self.module.training_step(*inputs, **kwargs)
             warn_if_output_is_none(output, "training_step")
-        elif self.module.testing:
+
+        elif running_stage == RunningStage.TESTING:
             output = self.module.test_step(*inputs, **kwargs)
             warn_if_output_is_none(output, "test_step")
-        else:
+
+        elif running_stage == RunningStage.EVALUATING:
             output = self.module.validation_step(*inputs, **kwargs)
             warn_if_output_is_none(output, "validation_step")
+
+        else:
+            output = self.module(*inputs, **kwargs)
+
         return output
 
 
@@ -276,16 +294,22 @@ def _worker(i, module, input, kwargs, device=None):
 
                 # ---------------
                 # CHANGE
-                if module.training:
+                if module.running_stage == RunningStage.TRAINING:
                     output = module.training_step(*input, **kwargs)
                     fx_called = 'training_step'
-                elif module.testing:
+
+                elif module.running_stage == RunningStage.TESTING:
                     output = module.test_step(*input, **kwargs)
                     fx_called = 'test_step'
-                else:
+
+                elif module.running_stage == RunningStage.EVALUATING:
                     output = module.validation_step(*input, **kwargs)
                     fx_called = 'validation_step'
 
+                else:
+                    output = module(*input, **kwargs)
+                    fx_called = 'forward'
+
                 if output is None:
                     warn_missing_output(fx_called)
 

pytorch_lightning/overrides/fairscale.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from pytorch_lightning.trainer.states import RunningStage
 from pytorch_lightning.utilities import _FAIRSCALE_AVAILABLE
 
 LightningShardedDataParallel = None
@@ -23,10 +24,18 @@ def forward(self, *inputs, **kwargs):
             if self.enable_broadcast_buffers:
                 self.sync_buffers()
 
-            if self.module.training:
+            running_stage = self.module.running_stage
+
+            if running_stage == RunningStage.TRAINING:
                 outputs = self.module.training_step(*inputs, **kwargs)
-            elif self.module.testing:
+
+            elif running_stage == RunningStage.TESTING:
                 outputs = self.module.test_step(*inputs, **kwargs)
-            else:
+
+            elif running_stage == RunningStage.EVALUATING:
                 outputs = self.module.validation_step(*inputs, **kwargs)
+
+            else:
+                outputs = self.module(*inputs, **kwargs)
+
             return outputs

pytorch_lightning/trainer/configuration_validator.py

@@ -52,7 +52,7 @@ def __verify_train_loop_configuration(self, model):
         # verify model has a train dataloader
         # -----------------------------------
         has_train_dataloader = is_overridden('train_dataloader', model)
-        if not has_train_dataloader:
+        if not has_train_dataloader and not self.trainer.is_predicting:
             raise MisconfigurationException(
                 'No `train_dataloader()` method defined. Lightning `Trainer` expects as minimum a'
                 ' `training_step()`, `train_dataloader()` and `configure_optimizers()` to be defined.'
@@ -62,7 +62,7 @@ def __verify_train_loop_configuration(self, model):
         # verify model has optimizer
         # -----------------------------------
         has_optimizers = is_overridden('configure_optimizers', model)
-        if not has_optimizers:
+        if not has_optimizers and not self.trainer.is_predicting:
             raise MisconfigurationException(
                 'No `configure_optimizers()` method defined. Lightning `Trainer` expects as minimum a'
                 ' `training_step()`, `train_dataloader()` and `configure_optimizers()` to be defined.'

pytorch_lightning/trainer/connectors/logger_connector/epoch_result_store.py

@@ -12,17 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from collections import defaultdict
+from enum import Enum
 from typing import Any, Dict, List, Optional, Union
 
 import torch
 
 from pytorch_lightning.core.step_result import Result
+from pytorch_lightning.trainer.states import RunningStage
 from pytorch_lightning.utilities import DistributedType, LightningEnum
 
 
-class LoggerStages(LightningEnum):
+class LoggerStages(str, Enum):
     """ Train/validation/test phase in each training step.
-
     >>> # you can math the type with string
     >>> LoggerStages.TRAIN == 'train'
     True
@@ -371,7 +372,7 @@ def update_logger_connector(self) -> None:
         callback_metrics = {}
         batch_pbar_metrics = {}
         batch_log_metrics = {}
-        is_train = self._stage in LoggerStages.TRAIN.value
+        is_train = self._stage in RunningStage.TRAINING
 
         if not self._has_batch_loop_finished:
             # get pbar

pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py

@@ -22,8 +22,9 @@
 from pytorch_lightning.core.step_result import Result
 from pytorch_lightning.loggers import LoggerCollection, TensorBoardLogger
 from pytorch_lightning.trainer.connectors.logger_connector.callback_hook_validator import CallbackHookNameValidator
-from pytorch_lightning.trainer.connectors.logger_connector.epoch_result_store import EpochResultStore, LoggerStages
+from pytorch_lightning.trainer.connectors.logger_connector.epoch_result_store import EpochResultStore
 from pytorch_lightning.trainer.connectors.logger_connector.metrics_holder import MetricsHolder
+from pytorch_lightning.trainer.states import RunningStage
 from pytorch_lightning.utilities import DeviceType, flatten_dict
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.model_helpers import is_overridden
@@ -37,9 +38,9 @@ def __init__(self, trainer):
         self._logged_metrics = MetricsHolder()
         self._progress_bar_metrics = MetricsHolder()
         self.eval_loop_results = []
-        self._cached_results = {stage: EpochResultStore(trainer, stage) for stage in LoggerStages}
+        self._cached_results = {stage: EpochResultStore(trainer, stage) for stage in RunningStage}
+        self._cached_results[None] = EpochResultStore(trainer, None)
         self._callback_hook_validator = CallbackHookNameValidator()
-        self._current_stage = None
 
     @property
     def callback_metrics(self) -> Dict:
@@ -75,7 +76,7 @@ def progress_bar_metrics(self, progress_bar_metrics: Dict) -> None:
 
     @property
     def cached_results(self) -> Union[EpochResultStore, None]:
-        return self._cached_results.get(self._current_stage)    # type: ignore
+        return self._cached_results.get(self.trainer._running_stage)    # type: ignore
 
     def get_metrics(self, key: str) -> Dict:
         metrics_holder = getattr(self, f"_{key}", None)
@@ -90,10 +91,8 @@ def set_metrics(self, key: str, val: Any) -> None:
         metrics_holder = getattr(self, f"_{key}", None)
         metrics_holder.reset(val)
 
-    def set_stage(self, stage_or_testing: Union[str, bool], reset: bool = False) -> None:
-        self._current_stage = LoggerStages.determine_stage(stage_or_testing)
-        if reset:
-            self.cached_results.reset()
+    def reset(self) -> None:
+        self.cached_results.reset()
 
     def check_logging_in_callbacks(self, hook_fx_name, on_step: bool = None, on_epoch: bool = None) -> None:
         self._callback_hook_validator.check_logging_in_callbacks(
@@ -119,8 +118,7 @@ def on_train_batch_end(self) -> None:
         self.cached_results._batch_size = None
 
     def cache_logged_metrics(self):
-        if self._current_stage is not None:
-            self._cached_results[self._current_stage].cache_result()
+        self._cached_results[self.trainer._running_stage].cache_result()
 
     def on_trainer_init(self, logger, flush_logs_every_n_steps: int, log_every_n_steps: int, move_metrics_to_cpu: bool):
         # logging