Fix trainer.predict(return_predictions=False) does not track batch_indices (#13629)

* Pull request for fixing issue #13580
* chlog and test
* disable track for epoch

Co-authored-by: rohitgr7 <[email protected]>

Author: georgestein

src/pytorch_lightning/CHANGELOG.md

@@ -342,6 +342,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed the input validation for the accelerator Trainer argument when passed as a string ([#13417](https://github.com/PyTorchLightning/pytorch-lightning/pull/13417))
 
 
+- Fixed `Trainer.predict(return_predictions=False)` to track prediction's batch_indices ([#13629](https://github.com/Lightning-AI/lightning/pull/13629))
+
+
 ## [1.6.5] - 2022-07-13
 
 ### Fixed

src/pytorch_lightning/loops/epoch/prediction_epoch_loop.py

@@ -67,7 +67,7 @@ def on_run_start(  # type: ignore[override]
         self._dl_max_batches = dl_max_batches
         self._num_dataloaders = num_dataloaders
         # this call requires that `self.return_predictions` is set
-        self._seen_batch_indices = self._get_batch_indices(dataloader_idx)
+        self._seen_batch_indices = self._get_batch_indices(dataloader_idx) if self.should_store_predictions else []
 
     def advance(  # type: ignore[override]
         self,
@@ -87,7 +87,7 @@ def advance(  # type: ignore[override]
         action_name = f"[{self.__class__.__name__}].predict_dataloader_idx_{dataloader_idx}_next"
         with self.trainer.profiler.profile(action_name):
             batch_idx, batch = next(dataloader_iter)
-        self._seen_batch_indices = self._get_batch_indices(dataloader_idx)
+        self._seen_batch_indices = self._get_batch_indices(dataloader_idx) if self.should_store_predictions else []
         # we need to truncate the list of batch indices due to prefetching in the dataloader and Lightning
         self._seen_batch_indices = self._seen_batch_indices[: (self.batch_progress.current.completed + 1)]
 
@@ -119,7 +119,8 @@ def _predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int) -> None
         step_kwargs = self._build_kwargs(batch, batch_idx, dataloader_idx)
 
         # extract batch_indices and store them
-        self.current_batch_indices = self._seen_batch_indices[batch_idx] if self._seen_batch_indices else []
+        batch_indices = self._get_batch_indices(dataloader_idx)
+        self.current_batch_indices = batch_indices[batch_idx] if batch_indices else []
 
         self.trainer._call_callback_hooks("on_predict_batch_start", batch, batch_idx, dataloader_idx)
         self.trainer._call_lightning_module_hook("on_predict_batch_start", batch, batch_idx, dataloader_idx)
@@ -166,7 +167,7 @@ def _get_batch_indices(self, dataloader_idx: int) -> List[List[int]]:
             "batch_sampler",
             None,
         )
-        if isinstance(batch_sampler, IndexBatchSamplerWrapper) and self.should_store_predictions:
+        if isinstance(batch_sampler, IndexBatchSamplerWrapper):
             return batch_sampler.seen_batch_indices
 
         warning_cache.warn("Lightning couldn't infer the indices fetched for your dataloader.")

tests/tests_pytorch/callbacks/test_prediction_writer.py

@@ -26,10 +26,10 @@
 
 
 class DummyPredictionWriter(BasePredictionWriter):
-    def write_on_batch_end(self, *args, **kwargs):
+    def write_on_batch_end(self, *_, **__):
         pass
 
-    def write_on_epoch_end(self, *args, **kwargs):
+    def write_on_epoch_end(self, *_, **__):
         pass
 
 
@@ -39,7 +39,7 @@ def test_prediction_writer_invalid_write_interval():
         DummyPredictionWriter("something")
 
 
-def test_prediction_writer_hook_call_intervals(tmpdir):
+def test_prediction_writer_hook_call_intervals():
     """Test that the `write_on_batch_end` and `write_on_epoch_end` hooks get invoked based on the defined
     interval."""
     DummyPredictionWriter.write_on_batch_end = Mock()
@@ -84,7 +84,7 @@ def test_prediction_writer_hook_call_intervals(tmpdir):
 
 
 @pytest.mark.parametrize("num_workers", [0, pytest.param(2, marks=RunIf(slow=True))])
-def test_prediction_writer_batch_indices(tmpdir, num_workers):
+def test_prediction_writer_batch_indices(num_workers):
     DummyPredictionWriter.write_on_batch_end = Mock()
     DummyPredictionWriter.write_on_epoch_end = Mock()
 
@@ -110,7 +110,7 @@ def test_prediction_writer_batch_indices(tmpdir, num_workers):
     )
 
 
-def test_prediction_writer_partial_support_for_combined_loader(tmpdir):
+def test_prediction_writer_partial_support_for_combined_loader():
     """Test partial support for CombinedLoader: prediction works but sample indices don't get tracked."""
     pl.loops.epoch.prediction_epoch_loop.warning_cache.clear()
 
@@ -140,3 +140,27 @@ def predict_step(self, batch, *args, **kwargs):
     )
 
     writer.write_on_epoch_end.assert_has_calls([call(trainer, model, ANY, [[]])])
+
+
+def test_batch_level_batch_indices():
+    """Test that batch_indices are returned when `return_predictions=False`."""
+    DummyPredictionWriter.write_on_batch_end = Mock()
+
+    class CustomBoringModel(BoringModel):
+        def on_predict_epoch_end(self, *args, **kwargs):
+            assert self.trainer.predict_loop.epoch_batch_indices == [[]]
+
+    writer = DummyPredictionWriter("batch")
+    model = CustomBoringModel()
+    dataloader = DataLoader(RandomDataset(32, 64), batch_size=4)
+    trainer = Trainer(limit_predict_batches=4, callbacks=writer)
+    trainer.predict(model, dataloaders=dataloader, return_predictions=False)
+
+    writer.write_on_batch_end.assert_has_calls(
+        [
+            call(trainer, model, ANY, [0, 1, 2, 3], ANY, 0, 0),
+            call(trainer, model, ANY, [4, 5, 6, 7], ANY, 1, 0),
+            call(trainer, model, ANY, [8, 9, 10, 11], ANY, 2, 0),
+            call(trainer, model, ANY, [12, 13, 14, 15], ANY, 3, 0),
+        ]
+    )