Repository: Lightning-AI/pytorch-lightning
Pull request: Add support for `len(datamodule)`
## What does this PR do?

<!--
Please include a summary of the change and which issue is fixed.
Please also include relevant motivation and context.
List any dependencies that are required for this change.

If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

The following links the related issue to the PR (https://docs.github.com/en/free-pro-team@latest/github/managing-your-work-on-github/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword)
-->

Fixes #5965

### Does your PR introduce any breaking changes? If yes, please list them.

<!-- FILL IN or None -->

## Before submitting

- [X] Was this **discussed/approved** via a GitHub issue? (not for typos and docs)
- [X] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [X] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- [X] Did you make sure to **update the documentation** with your changes? (if necessary)
- [X] Did you write any **new necessary tests**? (not for typos and docs)
- [X] Did you verify new and **existing tests pass** locally with your changes?
- [X] Did you list all the **breaking changes** introduced by this pull request?
- [X] Did you **update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)**? (not for typos, docs, test updates, or internal minor changes/refactorings)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing make sure you have read [Review guidelines](https://github.com/PyTorchLightning/pytorch-lightning/wiki/Review-guidelines). In short, see the following bullet-list:

- [X] Is this pull request ready for review? (if not, please submit in draft mode)
- [X] Check that all items from **Before submitting** are resolved
- [X] Make sure the title is self-explanatory and the description concisely explains the PR
- [X] Add labels and milestones (and optionally projects) to the PR so it can be classified

## Did you have fun?

Make sure you had fun coding 🙃


Problem statement:
support len(datamodule)
Let's add support for len(datamodule) so we can get the following:

```python
len(datamodule)

# prints:  
# train_dataloder_1:   200 samples
# train_dataloader_2:  500 samples
# tval_dataloder_1:    200 samples
# val_dataloader_2:    500 samples
# test_dataloder_1:    200 samples
# test_dataloader_2:   500 samples
```
cc @edenafek

Hints from issue discussion:
I think this is better suited for `__str__`, `__len__` should aways return an int
This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!

@carmocca I'd like to work on this. Do I need to write tests for this method?
@kingyiusuen Go for it! And yes, it needs to be tested.

The idea would be to iterate over all the dataloaders defined in the datamodule and sum all their individual lengths. Note that not all `*_dataloader` methods might be defined and that the dataloaders might not have `__len__` defined. In the latter case, I would suggest you count it as 0 and print a warning.
> @kingyiusuen Go for it! And yes, it needs to be tested.
> 
> The idea would be to iterate over all the dataloaders defined in the datamodule and sum all their individual lengths. Note that not all `*_dataloader` methods might be defined and that the dataloaders might not have `__len__` defined. In the latter case, I would suggest you count it as 0 and print a warning.

Just wanted to make sure, are we printing the number of batches (`len(dataloader)`) or number of samples (`len(dataloader.dataset)`)?
Dear @kingyiusuen,

I would argue batch_size seems more intuitive for users.

We could print in the current structure of the dataloaders implemented by the users.

```py
MyDataModuleClass(
    train_dataloader: {"a": DataLoaderClass(batch_size=8, num_batches=16, num_workers=2), "b":  DataLoaderClass(batch_size=2, num_batches=16, num_workers=2)]
    val_dataloader: [ DataLoaderClass(batch_size=3, num_batches=14, num_workers=0), DataLoaderClass(batch_size=8, num_batches=4, num_workers=0)]
    train_dataloader:DataLoaderClass(batch_size=4, num_batches=7, num_workers=2)
)
```


Note that it should not print anything but instead return an `int`. 

We can explore overridding `__str__` for the actual structure etc
> Note that it should not print anything but instead return an `int`.
> 
> We can explore overridding `__str__` for the actual structure etc

Got it. I will make one PR for `__len__` and one PR for `__str__` later.

Commit message: Fix flake8 issue

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index fa206f3784c49..7515e9a40d980 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -177,6 +177,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added support for `torch.autograd.set_detect_anomaly` through `Trainer` constructor argument `detect_anomaly` ([#9848](https://github.com/PyTorchLightning/pytorch-lightning/pull/9848))
 
+- Added a `len` method to `LightningDataModule` ([#9895](https://github.com/PyTorchLightning/pytorch-lightning/pull/9895))
 
 ### Changed
 
diff --git a/pytorch_lightning/core/datamodule.py b/pytorch_lightning/core/datamodule.py
index f3a5c855fe07a..46936b352ebf6 100644
--- a/pytorch_lightning/core/datamodule.py
+++ b/pytorch_lightning/core/datamodule.py
@@ -22,7 +22,10 @@
 from pytorch_lightning.core.hooks import CheckpointHooks, DataHooks
 from pytorch_lightning.core.mixins import HyperparametersMixin
 from pytorch_lightning.utilities import rank_zero_deprecation
+from pytorch_lightning.utilities.apply_func import apply_to_collection
 from pytorch_lightning.utilities.argparse import add_argparse_args, from_argparse_args, get_init_arguments_and_types
+from pytorch_lightning.utilities.data import has_len
+from pytorch_lightning.utilities.warnings import rank_zero_warn
 
 
 class LightningDataModule(CheckpointHooks, DataHooks, HyperparametersMixin):
@@ -481,3 +484,35 @@ def __getstate__(self) -> dict:
         for fn in ("prepare_data", "setup", "teardown"):
             del d[fn]
         return d
+
+    def __len__(self) -> int:
+        """Returns the total number of batches in all dataloaders defined in the datamodule."""
+
+        from pytorch_lightning.trainer.supporters import CombinedLoader
+
+        num_batches = 0
+        not_implemented_count = 0
+
+        def get_num_batches(dataloader: DataLoader) -> None:
+            nonlocal num_batches
+            if not has_len(dataloader):
+                rank_zero_warn(
+                    "The number of batches for a dataloader is counted as 0 because it does not have `__len__` defined."
+                )
+            else:
+                num_batches += len(dataloader)
+
+        for method_name in ("train_dataloader", "val_dataloader", "test_dataloader", "predict_dataloader"):
+            dataloader_method = getattr(self, method_name)
+            try:
+                dataloader = dataloader_method()
+                if isinstance(dataloader, CombinedLoader):
+                    dataloader = dataloader.loaders
+                apply_to_collection(dataloader, DataLoader, get_num_batches)
+            except NotImplementedError:
+                not_implemented_count += 1
+
+        if not_implemented_count == 4:
+            rank_zero_warn("You datamodule does not have any valid dataloader so `__len__` will be returned as 0.")
+
+        return num_batches
diff --git a/tests/core/test_datamodules.py b/tests/core/test_datamodules.py
index 2f84032593472..5890a9c49fd96 100644
--- a/tests/core/test_datamodules.py
+++ b/tests/core/test_datamodules.py
@@ -21,13 +21,15 @@
 import pytest
 import torch
 from omegaconf import OmegaConf
+from torch.utils.data import DataLoader
 
 from pytorch_lightning import LightningDataModule, Trainer
 from pytorch_lightning.callbacks import ModelCheckpoint
+from pytorch_lightning.trainer.supporters import CombinedLoader
 from pytorch_lightning.utilities import AttributeDict
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.model_helpers import is_overridden
-from tests.helpers import BoringDataModule, BoringModel
+from tests.helpers import BoringDataModule, BoringModel, RandomDataset
 from tests.helpers.datamodules import ClassifDataModule
 from tests.helpers.runif import RunIf
 from tests.helpers.simple_models import ClassificationModel
@@ -564,13 +566,14 @@ class BoringDataModule1(LightningDataModule):
         batch_size: int
         dims: int = 2
 
-        def __post_init__(self):
-            super().__init__(dims=self.dims)
+        def train_dataloader(self):
+            return DataLoader(torch.randn(self.batch_size * 2, 10), batch_size=self.batch_size)
 
     # asserts for the different dunder methods added by dataclass, when __init__ is implemented, i.e.
     # __repr__, __eq__, __lt__, __le__, etc.
     assert BoringDataModule1(batch_size=64).dims == 2
     assert BoringDataModule1(batch_size=32)
+    assert len(BoringDataModule1(batch_size=32)) == 2
     assert hasattr(BoringDataModule1, "__repr__")
     assert BoringDataModule1(batch_size=32) == BoringDataModule1(batch_size=32)
 
@@ -581,7 +584,9 @@ class BoringDataModule2(LightningDataModule):
 
     # asserts for the different dunder methods added by dataclass, when super class is inherently initialized, i.e.
     # __init__, __repr__, __eq__, __lt__, __le__, etc.
-    assert BoringDataModule2(batch_size=32)
+    assert BoringDataModule2(batch_size=32) is not None
+    assert BoringDataModule2(batch_size=32).batch_size == 32
+    assert len(BoringDataModule2(batch_size=32)) == 0
     assert hasattr(BoringDataModule2, "__repr__")
     assert BoringDataModule2(batch_size=32).prepare_data() is None
     assert BoringDataModule2(batch_size=32) == BoringDataModule2(batch_size=32)
@@ -625,3 +630,69 @@ def test_inconsistent_prepare_data_per_node(tmpdir):
         trainer.model = model
         trainer.datamodule = dm
         trainer.data_connector.prepare_data()
+
+
+DATALOADER = DataLoader(RandomDataset(1, 32))
+
+
+@pytest.mark.parametrize("method_name", ["train_dataloader", "val_dataloader", "test_dataloader", "predict_dataloader"])
+@pytest.mark.parametrize(
+    ["dataloader", "expected"],
+    [
+        [DATALOADER, 32],
+        [[DATALOADER, DATALOADER], 64],
+        [[[DATALOADER], [DATALOADER, DATALOADER]], 96],
+        [[{"foo": DATALOADER}, {"foo": DATALOADER, "bar": DATALOADER}], 96],
+        [{"foo": DATALOADER, "bar": DATALOADER}, 64],
+        [{"foo": {"foo": DATALOADER}, "bar": {"foo": DATALOADER, "bar": DATALOADER}}, 96],
+        [{"foo": [DATALOADER], "bar": [DATALOADER, DATALOADER]}, 96],
+        [CombinedLoader({"foo": DATALOADER, "bar": DATALOADER}), 64],
+    ],
+)
+def test_len_different_types(method_name, dataloader, expected):
+    dm = LightningDataModule()
+    setattr(dm, method_name, lambda: dataloader)
+    assert len(dm) == expected
+
+
+@pytest.mark.parametrize("method_name", ["train_dataloader", "val_dataloader", "test_dataloader", "predict_dataloader"])
+def test_len_dataloader_no_len(method_name):
+    class CustomNotImplementedErrorDataloader(DataLoader):
+        def __len__(self):
+            raise NotImplementedError
+
+    dataloader = CustomNotImplementedErrorDataloader(RandomDataset(1, 32))
+    dm = LightningDataModule()
+    setattr(dm, method_name, lambda: dataloader)
+    with pytest.warns(UserWarning, match="The number of batches for a dataloader is counted as 0"):
+        assert len(dm) == 0
+
+
+def test_len_all_dataloader_methods_implemented():
+    class BoringDataModule(LightningDataModule):
+        def __init__(self, dataloader):
+            super().__init__()
+            self.dataloader = dataloader
+
+        def train_dataloader(self):
+            return {"foo": self.dataloader, "bar": self.dataloader}
+
+        def val_dataloader(self):
+            return self.dataloader
+
+        def test_dataloader(self):
+            return [self.dataloader]
+
+        def predict_dataloader(self):
+            return [self.dataloader, self.dataloader]
+
+    dm = BoringDataModule(DATALOADER)
+
+    # 6 dataloaders each producing 32 batches: 6 * 32 = 192
+    assert len(dm) == 192
+
+
+def test_len_no_dataloader_methods_implemented():
+    dm = LightningDataModule()
+    with pytest.warns(UserWarning, match="You datamodule does not have any valid dataloader"):
+        assert len(dm) == 0


Code review comments:
diff_hunk: @@ -481,3 +484,35 @@ def __getstate__(self) -> dict:
         for fn in ("prepare_data", "setup", "teardown"):
             del d[fn]
         return d
+
+    def __len__(self) -> int:
+        """Returns the total number of batches in all dataloaders defined in the datamodule."""
+
+        from pytorch_lightning.trainer.supporters import CombinedLoader
+
+        num_batches = 0
+        not_implemented_count = 0
+
+        def get_num_batches(dataloader: DataLoader) -> None:
+            nonlocal num_batches
+            if not has_len(dataloader):
+                rank_zero_warn(
+                    "The number of batches for a dataloader is counted as 0 because it does not have `__len__` defined."
path: pytorch_lightning/core/datamodule.py
text: @user1:
```suggestion
                    f"The number of batches for {name} is counted as 0 because it does not have `__len__` defined."
```

You'll need to set `def get_num_batches(dataloader: DataLoader, name: str)`

and pass the name below with `apply_to_collection(dataloader, DataLoader, get_num_batches, method_name)`

diff_hunk: @@ -481,3 +484,35 @@ def __getstate__(self) -> dict:
         for fn in ("prepare_data", "setup", "teardown"):
             del d[fn]
         return d
+
+    def __len__(self) -> int:
+        """Returns the total number of batches in all dataloaders defined in the datamodule."""
+
+        from pytorch_lightning.trainer.supporters import CombinedLoader
+
+        num_batches = 0
+        not_implemented_count = 0
+
+        def get_num_batches(dataloader: DataLoader) -> None:
+            nonlocal num_batches
+            if not has_len(dataloader):
+                rank_zero_warn(
+                    "The number of batches for a dataloader is counted as 0 because it does not have `__len__` defined."
+                )
+            else:
+                num_batches += len(dataloader)
+
+        for method_name in ("train_dataloader", "val_dataloader", "test_dataloader", "predict_dataloader"):
+            dataloader_method = getattr(self, method_name)
+            try:
+                dataloader = dataloader_method()
+                if isinstance(dataloader, CombinedLoader):
+                    dataloader = dataloader.loaders
+                apply_to_collection(dataloader, DataLoader, get_num_batches)
+            except NotImplementedError:
+                not_implemented_count += 1
+
+        if not_implemented_count == 4:
path: pytorch_lightning/core/datamodule.py
text: @user1:
Is the additional variable nessesary? To me it seem to be the same as:
```suggestion
        if num_batches == 0:
```

@author:
It is possible that a dataloader is implemented but it does not have `__len__` defined. `num_batches` will also be 0 in this case.