[feat] Add better support for predict + ddp 2/3

CHANGELOG.md

@@ -111,6 +111,11 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added new `EarlyStopping` parameters `stopping_threshold` and `divergence_threshold` ([#6868](https://github.com/PyTorchLightning/pytorch-lightning/pull/6868))
 
 
+- Added new `UnrepeatedDistributedSampler` and `IndexBatchSamplerWrapper` for tracking distributed predictions ([#7215](https://github.com/PyTorchLightning/pytorch-lightning/pull/7215))
+
+
+- Added `trainer.predict(return_predictions=None|False|True)` ([#7215](https://github.com/PyTorchLightning/pytorch-lightning/pull/7215))
+
 
 ### Changed
 

pytorch_lightning/core/datamodule.py

@@ -26,7 +26,7 @@
 
 class _DataModuleWrapper(type):
 
- def __init__(self, *args, **kwargs):
+ def __init__(self, *args: Any, **kwargs: Any) -> None:
  super().__init__(*args, **kwargs)
  self.__has_added_checks = False
 

pytorch_lightning/core/lightning.py

@@ -74,7 +74,7 @@ class LightningModule(
  "model_size",
  ] + DeviceDtypeModuleMixin.__jit_unused_properties__
 
- def __init__(self, *args, **kwargs):
+ def __init__(self, *args: Any, **kwargs: Any) -> None:
  super().__init__(*args, **kwargs)
 
  # see (https://github.com/pytorch/pytorch/blob/3e6bb5233f9ca2c5aa55d9cda22a7ee85439aa6e/

pytorch_lightning/overrides/distributed.py

@@ -12,10 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import itertools
-from typing import Any
+from typing import Any, Iterator, List, Optional
 
 import torch
 from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import BatchSampler, DistributedSampler, Sampler
 
 from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.overrides.base import _LightningModuleWrapperBase
@@ -75,3 +76,70 @@ def prepare_for_backward(model: DistributedDataParallel, output: Any):
  model.reducer.prepare_for_backward([])
  else:
  model.require_forward_param_sync = False
+
+
+class UnrepeatedDistributedSampler(DistributedSampler):
+ """
+ A fork of the pytorch DistributedSampler that doesn't repeat data, instead
+ allowing the number of batches per process to be off-by-one from each other.
+ This makes this sampler usable for predictions (it's deterministic and
+ doesn't require shuffling). It is potentially unsafe to use this sampler for
+ training, because during training the DistributedDataParallel syncs buffers
+ on each forward pass, so it could freeze if one of the processes runs one
+ fewer batch. During prediction, buffers are only synced on the first batch,
+ so this is safe to use as long as each process runs at least one batch. We
+ verify this in an assert.
+
+ Taken from https://github.com/jpuigcerver/PyLaia/blob/v1.0.0/laia/data/unpadded_distributed_sampler.py
+ and https://github.com/pytorch/pytorch/issues/25162#issuecomment-634146002
+ """
+
+ def __init__(self, *args: Any, **kwargs: Any) -> None:
+ super().__init__(*args, **kwargs)
+ self.num_samples = len(range(self.rank, len(self.dataset), self.num_replicas))
+ self.total_size = len(self.dataset)
+ # If any process has at least one batch, every other process needs to
+ # have at least one batch, or the DistributedDataParallel could lock up.
+ assert self.num_samples >= 1 or self.total_size == 0
+
+ def __iter__(self) -> Iterator[List[int]]:
+ if self.shuffle:
+ # deterministically shuffle based on epoch
+ g = torch.Generator()
+ g.manual_seed(self.epoch)
+ indices = torch.randperm(len(self.dataset), generator=g).tolist()
+ else:
+ indices = list(range(len(self.dataset)))
+
+ assert len(indices) == self.total_size
+
+ # subsample
+ indices = indices[self.rank:self.total_size:self.num_replicas]
+ assert len(indices) == self.num_samples
+
+ return iter(indices)
+
+
+class IndexBatchSamplerWrapper:
+ """This class is used to wrap a :class:`torch.utils.data.BatchSampler` and capture its indices."""
+
+ def __init__(self, sampler: BatchSampler) -> None:
+ self._sampler = sampler
+ self.batch_indices: Optional[List[int]] = None
+
+ def __iter__(self) -> Iterator[List[int]]:
+ for batch in self._sampler:
+ self.batch_indices = batch
+ yield batch
+
+ @property
+ def drop_last(self) -> bool:
+ return self._sampler.drop_last
+
+ @property
+ def batch_size(self) -> int:
+ return self._sampler.batch_size
+
+ @property
+ def sampler(self) -> Sampler:
+ return self._sampler.sampler

pytorch_lightning/trainer/connectors/callback_connector.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 import os
 from datetime import timedelta
-from typing import List, Union, Optional, Dict
+from typing import Dict, List, Optional, Union
 
 from pytorch_lightning.callbacks import Callback, ModelCheckpoint, ProgressBar, ProgressBarBase
 from pytorch_lightning.callbacks.timer import Timer
@@ -58,6 +58,8 @@ def on_trainer_init(
  # configure swa callback
  self._configure_swa_callbacks()
 
+ # configure the timer callback.
+ # responsible to stop the training when max_time is reached.
  self._configure_timer_callback(max_time)
 
  # init progress bar
@@ -115,9 +117,7 @@ def _configure_timer_callback(self, max_time: Optional[Union[str, timedelta, Dic
  if max_time is None:
  return
  if any(isinstance(cb, Timer) for cb in self.trainer.callbacks):
- rank_zero_info(
- "Ignoring `Trainer(max_time=...)`, callbacks list already contains a Timer."
- )
+ rank_zero_info("Ignoring `Trainer(max_time=...)`, callbacks list already contains a Timer.")
  return
  timer = Timer(duration=max_time, interval="step")
  self.trainer.callbacks.append(timer)

pytorch_lightning/trainer/data_loading.py

@@ -17,14 +17,16 @@
 from abc import ABC
 from copy import deepcopy
 from functools import partial
-from typing import Iterable, List, Optional, Tuple, Union
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 from torch.utils.data import BatchSampler, DataLoader, RandomSampler, SequentialSampler
 from torch.utils.data.distributed import DistributedSampler
 
 from pytorch_lightning.accelerators import Accelerator
 from pytorch_lightning.core import LightningModule
+from pytorch_lightning.overrides.distributed import IndexBatchSamplerWrapper, UnrepeatedDistributedSampler
 from pytorch_lightning.trainer.connectors.accelerator_connector import AcceleratorConnector
+from pytorch_lightning.trainer.states import RunningStage
 from pytorch_lightning.trainer.supporters import CombinedLoader
 from pytorch_lightning.utilities import _TORCH_GREATER_EQUAL_1_6, rank_zero_warn
 from pytorch_lightning.utilities.apply_func import apply_to_collection
@@ -107,7 +109,9 @@ def auto_add_worker_init_fn(self, dataloader: DataLoader) -> None:
  if int(os.environ.get("PL_SEED_WORKERS", 0)) and dataloader.worker_init_fn is None:
  dataloader.worker_init_fn = partial(pl_worker_init_function, rank=self.global_rank)
 
- def auto_add_sampler(self, dataloader: DataLoader, shuffle: bool) -> DataLoader:
+ def auto_add_sampler(
+ self, dataloader: DataLoader, shuffle: bool, mode: Optional[RunningStage] = None
+ ) -> DataLoader:
  # don't do anything if it's not a dataloader
  is_dataloader = isinstance(dataloader, DataLoader)
  # don't manipulate iterable datasets
@@ -133,20 +137,24 @@ def auto_add_sampler(self, dataloader: DataLoader, shuffle: bool) -> DataLoader:
  )
 
  # replace with distributed sampler
- sampler = self._get_distributed_sampler(dataloader, shuffle)
- dataloader = self.replace_sampler(dataloader, sampler)
+ sampler = self._get_distributed_sampler(dataloader, shuffle, mode=mode)
+ dataloader = self.replace_sampler(dataloader, sampler, mode=mode)
 
  return dataloader
 
  @staticmethod
- def _resolve_batch_sampler(dl_args, dataloader, sampler):
+ def _resolve_batch_sampler(dl_args, dataloader, sampler, mode: Optional[RunningStage] = None) -> Dict[str, Any]:
  batch_sampler = getattr(dataloader, "batch_sampler")
- if batch_sampler is not None and type(batch_sampler) is not BatchSampler:
+ is_predicting = mode == RunningStage.PREDICTING
+ # checking the batch sampler type is different than PyTorch default.
+ if (batch_sampler is not None and type(batch_sampler) is not BatchSampler) or is_predicting:
  batch_sampler = type(batch_sampler)(
  sampler,
  batch_size=batch_sampler.batch_size,
- drop_last=batch_sampler.drop_last,
+ drop_last=(False if is_predicting else batch_sampler.drop_last),
  )
+ if is_predicting:
+ batch_sampler = IndexBatchSamplerWrapper(batch_sampler)
  dl_args['batch_sampler'] = batch_sampler
  dl_args['batch_size'] = 1
  dl_args['shuffle'] = False
@@ -159,7 +167,7 @@ def _resolve_batch_sampler(dl_args, dataloader, sampler):
 
  return dl_args
 
- def replace_sampler(self, dataloader, sampler):
+ def replace_sampler(self, dataloader: DataLoader, sampler, mode: Optional[RunningStage] = None) -> DataLoader:
  skip_keys = ('sampler', 'batch_sampler', 'dataset_kind')
  skip_signature_keys = ('args', 'kwargs', 'self')
 
@@ -174,7 +182,7 @@ def replace_sampler(self, dataloader, sampler):
 
  dl_args = {name: attrs[name] for name in params if name in attrs and name not in skip_keys}
 
- dl_args = self._resolve_batch_sampler(dl_args, dataloader, sampler)
+ dl_args = self._resolve_batch_sampler(dl_args, dataloader, sampler, mode=mode)
 
  multiprocessing_context = dataloader.multiprocessing_context
  dl_args['multiprocessing_context'] = multiprocessing_context
@@ -205,12 +213,15 @@ def __init__(self, num_features, dataset, *args, **kwargs):
  dataloader.multiprocessing_context = multiprocessing_context
  return dataloader
 
- def _get_distributed_sampler(self, dataloader, shuffle):
+ def _get_distributed_sampler(
+ self, dataloader: DataLoader, shuffle: bool, mode: Optional[RunningStage] = None
+ ) -> DistributedSampler:
  kwargs = self.distributed_sampler_kwargs
  kwargs["shuffle"] = shuffle and not self.overfit_batches
  if _TORCH_GREATER_EQUAL_1_6:
  kwargs.setdefault("seed", int(os.getenv("PL_GLOBAL_SEED", 0)))
- sampler = DistributedSampler(dataloader.dataset, **kwargs)
+ cls = UnrepeatedDistributedSampler if mode == RunningStage.PREDICTING else DistributedSampler
+ sampler = cls(dataloader.dataset, **kwargs)
  return sampler
 
  def reset_train_dataloader(self, model: LightningModule) -> None:
@@ -296,7 +307,7 @@ def _reset_eval_dataloader(
 
  Args:
  model: The current `LightningModule`
- mode: Either `'val'` or `'test'`
+ mode: Either `'val'`, `'test'` or `'predict'`
 
  Returns:
  Tuple (num_batches, dataloaders)
@@ -342,7 +353,9 @@ def _reset_eval_dataloader(
  rank_zero_warn("One of given dataloaders is None and it will be skipped.")
 
  # add samplers
- dataloaders = [self.auto_add_sampler(dl, shuffle=False) for dl in dataloaders if dl is not None]
+ dataloaders = [
+ self.auto_add_sampler(dl, shuffle=False, mode=self._running_stage) for dl in dataloaders if dl is not None
+ ]
 
  # add worker_init_fn for correct seeding in worker processes
  apply_to_collection(dataloaders, dtype=DataLoader, function=self.auto_add_worker_init_fn)

pytorch_lightning/trainer/predict_loop.py

@@ -11,9 +11,14 @@
 # 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 typing import Any, List, Optional, Union
+
 import torch
+from torch.utils.data.dataloader import DataLoader
 
-from pytorch_lightning.utilities.apply_func import apply_to_collection
+from pytorch_lightning.overrides.distributed import IndexBatchSamplerWrapper
+from pytorch_lightning.plugins import DDPSpawnPlugin
+from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.warnings import WarningCache
 
 
@@ -24,6 +29,27 @@ def __init__(self, trainer):
  self.max_batches = None
  self.num_dataloaders = None
  self.warning_cache = WarningCache()
+ self.batch_indices: Optional[List[int]] = None
+ self.epoch_batch_indices: Optional[List[List[int]]] = None
+ # `DDPSpawnPlugin` plugins and derivate don't support return predictions.
+ self._return_predictions: Optional[bool] = None
+ self._previous_grad_status: Optional[bool] = None
+
+ @property
+ def return_predictions(self) -> bool:
+ return self._return_predictions
+
+ @return_predictions.setter
+ def return_predictions(self, return_predictions: Optional[bool] = None) -> None:
+ # ``DDPSpawnPlugin`` plugins and derivate don't support return predictions.
+ is_ddp_spawn = isinstance(self.trainer.training_type_plugin, DDPSpawnPlugin)
+ if return_predictions and is_ddp_spawn:
+ raise MisconfigurationException(
+ "`return_predictions` should be set to `False` when using the `DDPSpawnPlugin` or children class. "
+ f"Found {return_predictions} with training_type_plugin {type(self.trainer.training_type_plugin)}."
+ )
+ # For non ``DDPSpawnPlugin`` plugin, the `return_predictions` is True by default unless user decide otherwise.
+ self._return_predictions = not is_ddp_spawn if return_predictions is None else return_predictions
 
  def on_trainer_init(self):
  self.trainer.num_predict_batches = []
@@ -54,22 +80,26 @@ def setup(self, model, max_batches, dataloaders):
 
  self.max_batches = max_batches
  self.num_dataloaders = self._get_num_dataloaders(dataloaders)
- self._predictions = [[] for _ in range(self.num_dataloaders)]
+ self.predictions = [[] for _ in range(self.num_dataloaders)]
+ self.epoch_batch_indices = [[] for _ in range(self.num_dataloaders)]
 
- def _get_num_dataloaders(self, dataloaders):
+ def _get_num_dataloaders(self, dataloaders: List[DataLoader]) -> int:
  # case where user does:
  # return dl1, dl2
  length = len(dataloaders)
  if len(dataloaders) > 0 and isinstance(dataloaders[0], (list, tuple)):
  length = len(dataloaders[0])
  return length
 
- def predict_step(self, batch, batch_idx, dataloader_idx):
+ def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int) -> None:
  # configure args
  args = [batch, batch_idx]
  if self.num_dataloaders:
  args.append(dataloader_idx)
 
+ # extract batch_indices and store them
+ self._store_batch_indices(dataloader_idx)
+
  model_ref = self.trainer.lightning_module
 
  self.trainer.call_hook("on_predict_batch_start", batch, batch_idx, dataloader_idx)
@@ -82,18 +112,44 @@ def predict_step(self, batch, batch_idx, dataloader_idx):
 
  self.trainer.call_hook("on_predict_batch_end", predictions, batch, batch_idx, dataloader_idx)
 
- self._predictions[dataloader_idx].append(predictions)
+ if self.return_predictions:
+ self.predictions[dataloader_idx].append(predictions)
+
+ def _store_batch_indices(self, dataloader_idx: int) -> None:
+ batch_sampler = self.trainer.predict_dataloaders[dataloader_idx].batch_sampler
+ if isinstance(batch_sampler, IndexBatchSamplerWrapper):
+ self.batch_indices = batch_sampler.batch_indices
+ if self.return_predictions:
+ self.epoch_batch_indices[dataloader_idx].append(batch_sampler.batch_indices)
 
- def on_predict_epoch_end(self):
+ def on_predict_start(self) -> None:
+ # enable eval mode + no grads
+ self.on_predict_model_eval()
+ self.trainer.lightning_module.zero_grad()
+ self._previous_grad_status = torch.is_grad_enabled()
+ torch.set_grad_enabled(False)
+
+ # hook
+ self.trainer.call_hook("on_predict_start")
+ self.trainer.call_hook("on_predict_epoch_start")
+
+ def on_predict_epoch_end(self) -> Optional[Union[List[Any], List[List[Any]]]]:
  self.trainer.profiler.describe()
 
- results = self._predictions
+ results: List[List[Any]] = self.predictions
 
  self.trainer.call_hook("on_predict_epoch_end", results)
 
- def _convert_to_numpy(v):
- return v.cpu().numpy()
+ if self.return_predictions:
+ return results[0] if self.num_dataloaders == 1 else results
+
+ def on_predict_end(self):
+ # clear memory. the predictions are extracted in `on_predict_epoch_end`.
+ self.predictions = None
+ self.batch_indices = None
 
- results = apply_to_collection(results, torch.Tensor, _convert_to_numpy)
+ # reset grad to its previous status.
+ torch.set_grad_enabled(self._previous_grad_status)
 
- return results[0] if len(results) == 1 else results
+ # hook
+ self.trainer.call_hook("on_predict_end")