Support predict_dataset in LightningDataModule.from_datasets (#12942)

Co-authored-by: Rohit Gupta <rohitgr1998@gmail.com>

Author: digital-idiot

CHANGELOG.md

@@ -33,8 +33,13 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Include a version suffix for new "last" checkpoints of later runs in the same directory ([#12902](https://github.com/PyTorchLightning/pytorch-lightning/pull/12902))
 
 
+-  Added missing `predict_dataset` argument in `LightningDataModule.from_datasets` to create predict dataloaders ([#12942](https://github.com/PyTorchLightning/pytorch-lightning/pull/12942))
+
+
 - Added class name prefix to metrics logged by `DeviceStatsMonitor` ([#12228](https://github.com/PyTorchLightning/pytorch-lightning/pull/12228))
 
+
+
 ### Changed
 
 - Enable validation during overfitting ([#12527](https://github.com/PyTorchLightning/pytorch-lightning/pull/12527))

pytorch_lightning/core/datamodule.py

@@ -102,6 +102,7 @@ def from_datasets(
         train_dataset: Optional[Union[Dataset, Sequence[Dataset], Mapping[str, Dataset]]] = None,
         val_dataset: Optional[Union[Dataset, Sequence[Dataset]]] = None,
         test_dataset: Optional[Union[Dataset, Sequence[Dataset]]] = None,
+        predict_dataset: Optional[Union[Dataset, Sequence[Dataset]]] = None,
         batch_size: int = 1,
         num_workers: int = 0,
     ):
@@ -112,6 +113,7 @@ def from_datasets(
             train_dataset: (optional) Dataset to be used for train_dataloader()
             val_dataset: (optional) Dataset or list of Dataset to be used for val_dataloader()
             test_dataset: (optional) Dataset or list of Dataset to be used for test_dataloader()
+            predict_dataset: (optional) Dataset or list of Dataset to be used for predict_dataloader()
             batch_size: Batch size to use for each dataloader. Default is 1.
             num_workers: Number of subprocesses to use for data loading. 0 means that the
                 data will be loaded in the main process. Number of CPUs available.
@@ -139,13 +141,20 @@ def test_dataloader():
                 return [dataloader(ds) for ds in test_dataset]
             return dataloader(test_dataset)
 
+        def predict_dataloader():
+            if isinstance(predict_dataset, Sequence):
+                return [dataloader(ds) for ds in predict_dataset]
+            return dataloader(predict_dataset)
+
         datamodule = cls()
         if train_dataset is not None:
             datamodule.train_dataloader = train_dataloader
         if val_dataset is not None:
             datamodule.val_dataloader = val_dataloader
         if test_dataset is not None:
             datamodule.test_dataloader = test_dataloader
+        if predict_dataset is not None:
+            datamodule.predict_dataloader = predict_dataloader
         return datamodule
 
     def state_dict(self) -> Dict[str, Any]:
@@ -154,7 +163,7 @@ def state_dict(self) -> Dict[str, Any]:
         Returns:
             A dictionary containing datamodule state.
         """
-        return {}
+        return dict()
 
     def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
         """Called when loading a checkpoint, implement to reload datamodule state given datamodule state_dict.

tests/core/test_datamodules.py

@@ -377,9 +377,9 @@ def test_dm_init_from_datasets_dataloaders(iterable):
     with mock.patch("pytorch_lightning.core.datamodule.DataLoader") as dl_mock:
         dm.train_dataloader()
         dl_mock.assert_called_once_with(train_ds, batch_size=4, shuffle=not iterable, num_workers=0, pin_memory=True)
-    with pytest.raises(MisconfigurationException):
+    with pytest.raises(MisconfigurationException, match="`val_dataloader` must be implemented"):
         _ = dm.val_dataloader()
-    with pytest.raises(MisconfigurationException):
+    with pytest.raises(MisconfigurationException, match="`test_dataloader` must be implemented"):
         _ = dm.test_dataloader()
 
     train_ds_sequence = [ds(), ds()]
@@ -392,9 +392,9 @@ def test_dm_init_from_datasets_dataloaders(iterable):
                 call(train_ds_sequence[1], batch_size=4, shuffle=not iterable, num_workers=0, pin_memory=True),
             ]
         )
-    with pytest.raises(MisconfigurationException):
+    with pytest.raises(MisconfigurationException, match="`val_dataloader` must be implemented"):
         _ = dm.val_dataloader()
-    with pytest.raises(MisconfigurationException):
+    with pytest.raises(MisconfigurationException, match="`test_dataloader` must be implemented"):
         _ = dm.test_dataloader()
 
     valid_ds = ds()
@@ -405,21 +405,25 @@ def test_dm_init_from_datasets_dataloaders(iterable):
         dl_mock.assert_called_with(valid_ds, batch_size=2, shuffle=False, num_workers=0, pin_memory=True)
         dm.test_dataloader()
         dl_mock.assert_called_with(test_ds, batch_size=2, shuffle=False, num_workers=0, pin_memory=True)
-    with pytest.raises(MisconfigurationException):
+    with pytest.raises(MisconfigurationException, match="`train_dataloader` must be implemented"):
         _ = dm.train_dataloader()
 
     valid_dss = [ds(), ds()]
     test_dss = [ds(), ds()]
-    dm = LightningDataModule.from_datasets(train_ds, valid_dss, test_dss, batch_size=4, num_workers=0)
+    predict_dss = [ds(), ds()]
+    dm = LightningDataModule.from_datasets(train_ds, valid_dss, test_dss, predict_dss, batch_size=4, num_workers=0)
     with mock.patch("pytorch_lightning.core.datamodule.DataLoader") as dl_mock:
         dm.val_dataloader()
         dm.test_dataloader()
+        dm.predict_dataloader()
         dl_mock.assert_has_calls(
             [
                 call(valid_dss[0], batch_size=4, shuffle=False, num_workers=0, pin_memory=True),
                 call(valid_dss[1], batch_size=4, shuffle=False, num_workers=0, pin_memory=True),
                 call(test_dss[0], batch_size=4, shuffle=False, num_workers=0, pin_memory=True),
                 call(test_dss[1], batch_size=4, shuffle=False, num_workers=0, pin_memory=True),
+                call(predict_dss[0], batch_size=4, shuffle=False, num_workers=0, pin_memory=True),
+                call(predict_dss[1], batch_size=4, shuffle=False, num_workers=0, pin_memory=True),
             ]
         )