Simplify several profile calls (#11031)

Author: daniellepintz

CHANGELOG.md

@@ -102,6 +102,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
     * Some configuration errors that were previously raised as `MisconfigurationException`s will now be raised as `ProcessRaisedException` (torch>=1.8) or as `Exception` (torch<1.8)
 
 
+- Changed `batch_to_device` entry in profiling from stage-specific to generic, to match profiling of other hooks ([#11031](https://github.com/PyTorchLightning/pytorch-lightning/pull/11031))
+
+
 - Changed the info message for finalizing ddp-spawn worker processes to a debug-level message ([#10864](https://github.com/PyTorchLightning/pytorch-lightning/pull/10864))
 
 

pytorch_lightning/loops/batch/training_batch_loop.py

@@ -116,9 +116,7 @@ def _tbptt_split_batch(self, batch: Any) -> List[Any]:
         if tbptt_steps == 0:
             return [batch]
 
-        model_ref = self.trainer.lightning_module
-        with self.trainer.profiler.profile("tbptt_split_batch"):
-            splits = model_ref.tbptt_split_batch(batch, tbptt_steps)
+        splits = self.trainer._call_lightning_module_hook("tbptt_split_batch", batch, tbptt_steps)
         return splits
 
     def _update_running_loss(self, current_loss: Tensor) -> None:

pytorch_lightning/loops/epoch/evaluation_epoch_loop.py

@@ -107,8 +107,7 @@ def advance(  # type: ignore[override]
             raise StopIteration
 
         if not data_fetcher.store_on_device:
-            with self.trainer.profiler.profile("evaluation_batch_to_device"):
-                batch = self.trainer.training_type_plugin.batch_to_device(batch, dataloader_idx=(dataloader_idx or 0))
+            batch = self.trainer._call_ttp_hook("batch_to_device", batch, dataloader_idx=(dataloader_idx or 0))
 
         self.batch_progress.increment_ready()
 
@@ -121,9 +120,9 @@ def advance(  # type: ignore[override]
         self.batch_progress.increment_started()
 
         # lightning module methods
-        with self.trainer.profiler.profile("evaluation_step_and_end"):
-            output = self._evaluation_step(**kwargs)
-            output = self._evaluation_step_end(output)
+
+        output = self._evaluation_step(**kwargs)
+        output = self._evaluation_step_end(output)
 
         self.batch_progress.increment_processed()
 

pytorch_lightning/loops/epoch/prediction_epoch_loop.py

@@ -96,13 +96,11 @@ def advance(  # type: ignore[override]
         if batch is None:
             raise StopIteration
 
-        with self.trainer.profiler.profile("predict_batch_to_device"):
-            batch = self.trainer.training_type_plugin.batch_to_device(batch, dataloader_idx=dataloader_idx)
+        batch = self.trainer._call_ttp_hook("batch_to_device", batch, dataloader_idx=dataloader_idx)
 
         self.batch_progress.increment_ready()
 
-        with self.trainer.profiler.profile("predict_step"):
-            self._predict_step(batch, batch_idx, dataloader_idx)
+        self._predict_step(batch, batch_idx, dataloader_idx)
 
     def on_run_end(self) -> Tuple[List[Any], List[List[int]]]:
         """Returns the predictions and the corresponding batch indices."""

pytorch_lightning/loops/epoch/training_epoch_loop.py

@@ -156,8 +156,7 @@ def advance(self, data_fetcher: AbstractDataFetcher) -> None:  # type: ignore[ov
         batch_idx, (batch, self.batch_progress.is_last_batch) = next(self._dataloader_iter)
 
         if not data_fetcher.store_on_device:
-            with self.trainer.profiler.profile("training_batch_to_device"):
-                batch = self.trainer.training_type_plugin.batch_to_device(batch)
+            batch = self.trainer._call_ttp_hook("batch_to_device", batch)
 
         self.batch_progress.increment_ready()