Forward extra keyword arguments in LightningDataModule.from_datasets (#14185)

Co-authored-by: otaj <[email protected]>
Co-authored-by: Carlos Mocholí <[email protected]>

Author: 3 people

src/pytorch_lightning/CHANGELOG.md

@@ -12,7 +12,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added prefix to log message in `seed_everything` with rank info ([#13290](https://github.com/Lightning-AI/lightning/issues/13290))
 
 
--
+- Added support for passing extra init-parameters to the `LightningDataModule.from_datasets` ([#14185](https://github.com/Lightning-AI/lightning/issues/14185))
+
 
 
 ### Changed

src/pytorch_lightning/core/datamodule.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """LightningDataModule for loading DataLoaders with ease."""
+import inspect
 from argparse import ArgumentParser, Namespace
 from typing import Any, Dict, IO, List, Mapping, Optional, Sequence, Tuple, Union
 
@@ -109,19 +110,22 @@ def from_datasets(
         predict_dataset: Optional[Union[Dataset, Sequence[Dataset]]] = None,
         batch_size: int = 1,
         num_workers: int = 0,
+        **datamodule_kwargs: Any,
     ):
         r"""
         Create an instance from torch.utils.data.Dataset.
 
         Args:
-            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.
+            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. This parameter gets forwarded to the
+                ``__init__`` if the datamodule has such a name defined in its signature.
             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.
-
+                data will be loaded in the main process. Number of CPUs available. This parameter gets forwarded to the
+                ``__init__`` if the datamodule has such a name defined in its signature.
+            **datamodule_kwargs: Additional parameters that get passed down to the datamodule's ``__init__``.
         """
 
         def dataloader(ds: Dataset, shuffle: bool = False) -> DataLoader:
@@ -150,7 +154,17 @@ def predict_dataloader():
                 return [dataloader(ds) for ds in predict_dataset]
             return dataloader(predict_dataset)
 
-        datamodule = cls()
+        candidate_kwargs = dict(batch_size=batch_size, num_workers=num_workers)
+        accepted_params = inspect.signature(cls.__init__).parameters
+        accepts_kwargs = any(param.kind == param.VAR_KEYWORD for param in accepted_params.values())
+        if accepts_kwargs:
+            special_kwargs = candidate_kwargs
+        else:
+            accepted_params = set(accepted_params)
+            accepted_params.discard("self")
+            special_kwargs = {k: v for k, v in candidate_kwargs.items() if k in accepted_params}
+
+        datamodule = cls(**datamodule_kwargs, **special_kwargs)
         if train_dataset is not None:
             datamodule.train_dataloader = train_dataloader
         if val_dataset is not None:

tests/tests_pytorch/core/test_datamodules.py

@@ -366,6 +366,54 @@ def test_dm_init_from_datasets_dataloaders(iterable):
         )
 
 
+def test_dm_init_from_datasets_with_init_params():
+    """Test that extra kwargs can be passed down to the init via the ``LightningDataModule.from_datasets`` method.
+
+    The two special arguments batch_size and num_workers get passed down depending on whether the __init__ accepts them.
+    """
+    # No additional parameters
+    LightningDataModule.from_datasets(DummyDS(), batch_size=4, num_workers=2)
+
+    class KnownExtraParametersDataModule(LightningDataModule):
+        def __init__(self, batch_size=1, num_workers=0):
+            super().__init__()
+            self.batch_size = batch_size
+            self.num_workers = num_workers
+
+    # batch_size and num_workers get special treatment - they are part of the `from_datasets` signature
+    dm = KnownExtraParametersDataModule.from_datasets(DummyDS(), batch_size=4, num_workers=2)
+    assert dm.batch_size == 4
+    assert dm.num_workers == 2
+
+    class UnknownExtraParametersDataModule(LightningDataModule):
+        def __init__(self, other, batch_size=1):
+            super().__init__()
+            self.other = other
+            self.batch_size = batch_size
+
+    # additional parameter `other` gets forwarded, alongside the special `batch_size` parameter
+    dm = UnknownExtraParametersDataModule.from_datasets(DummyDS(), batch_size=4, num_workers=2, other=5)
+    assert dm.batch_size == 4
+    assert dm.other == 5
+
+    # positional arguments raise an error as they would when instantiating the datamodule normally
+    with pytest.raises(TypeError, match="missing 1 required positional argument: 'other'"):
+        UnknownExtraParametersDataModule.from_datasets(DummyDS(), batch_size=4, num_workers=2)
+
+    class KwargsParametersDataModule(LightningDataModule):
+        def __init__(self, num_workers, **kwargs):
+            super().__init__()
+            self.num_workers = num_workers
+            for key, value in kwargs.items():
+                setattr(self, key, value)
+
+    # everything gets forwarded, because there is `**kwargs` present
+    dm = KwargsParametersDataModule.from_datasets(DummyDS(), batch_size=10, num_workers=100, another=None)
+    assert dm.batch_size == 10
+    assert dm.num_workers == 100
+    assert dm.another is None
+
+
 # all args
 class DataModuleWithHparams_0(LightningDataModule):
     def __init__(self, arg0, arg1, kwarg0=None):