Repository: Lightning-AI/pytorch-lightning
Pull request: Add warning if multiple batch_sizes are found from ambiguous batch
## 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 #10093

### 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:
`utilities.data.extract_batch_size` can silently fail -- maybe it needs to be less automatic?
## Motivation

Hey! So from what I understand `extract_batch_size` is used by `ResultCollection` to reduce/re-weight various logging metrics such as validation loss, etc.

We had an issue where our batch collation function aggregated Tensors into a list instead of stacking them (needed because of size mismatches), and `extract_batch_size` would recursively crawl through the batch object until it found a Tensor and returned its first dimention size -- which is wrong (the list length is correct batch size). So we would silently get "random" batch sizes and the val_loss logging becomes incorrect and then the LR scheduling / early stopping incorrectly triggers etc.

### Pitch

1. Maybe there should be some warning / note in the documentation (didn't find any)
2. If the batch object isn't a tensor (like a dict/dataclass), maybe try to `getattr(batch, 'batch_size')` and only then produce a warning and try to infer it automatically.

Hints from issue discussion:
Dear @ronif, if you provide the batch_size to self.log(..., batch_size=batch_size), it should alleviate this problem as it will override the auto-detected value. We are aware of this and we introduced the batch_size in `self.log` for this reason.

Thanks for the quick reply @tchaton. We're using `self.log` in `on_validation_end` -- is that the idiomatic way? Should we provide batch_size there even though it's supposedly irrelevant at that point? (since the logged val loss is already reduced from the values collected at the different batches). In this situation batch_size should be 1 by default, and not extracted from the last batch, no?

In any case, just wanted to note it was a bit confusing for us -- if you're aware of the issue you can close.
since logging inside `on_validation_end` is supposed to happen at epoch level, there is no need for batch size since it's only used to accumulate metrics logged with `on_step`. Are you seeing any issue there?
Okay, it seems a bit more nuanced than that. And also, I meant `validation_epoch_end` and not `on_validation_end`. Look at this code
```
import torch
import torch.nn as nn
import pytorch_lightning as pl
from torch.utils.data import Dataset, DataLoader
from collections import OrderedDict
from torch.optim import SGD

from pytorch_lightning import loggers

FAKE_BATCH_SIZE = 3

class MyDataset(Dataset):
    def __len__(self):
        return 32

    def __getitem__(self, index):
        return torch.rand((3,32,32))

def collate_fn(batch):
    out = OrderedDict()
    # this will be the first item extract_batch_size finds, and will infer wrong batch_size
    out['confuse_extract_batch_size'] = torch.empty(size=(FAKE_BATCH_SIZE, 1))
    out['batch'] = torch.stack(batch)
    return out

train_loader = DataLoader(MyDataset(), batch_size=4, collate_fn=collate_fn)
val_loader = DataLoader(MyDataset(), batch_size=4, collate_fn=collate_fn)

class Model(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 1, 3)

    def training_step(self, batch, batch_idx):
        x = self.conv(batch['batch']).mean() # just something with a grad
        return {'loss': x}

    def validation_step(self, batch, batch_idx):
        return {'val_loss': torch.tensor(.5)}

    def validation_epoch_end(self, outputs):
        avg_loss = torch.stack([o['val_loss'] for o in outputs]).mean()
        self.log('val_loss', avg_loss)
        print(str(self.trainer._results))

    def configure_optimizers(self):
        return SGD(self.parameters(), lr=0.5)


trainer = pl.Trainer(max_epochs=2, logger=loggers.CSVLogger('/tmp/', name='test_extract_batch_size', version='ver'))
m = Model()
trainer.fit(m, train_dataloaders=train_loader, val_dataloaders=val_loader)
```

When `FAKE_BATCH_SIZE` is 0, this messes up the logging -- you will get `val_loss = nan`  in `metrics.csv`. If it's not 0, then the logged value will be correct (`0.5`), but `self.trainer._results` will have `cumulated_batch_size = FAKE_BATCH_SIZE` and  `value = val_loss * FAKE_BATCH_SIZE`, which is a bit weird. Are the values in `trainer._results` used by early stopping or LR schedulers?
Dear @ronif,

Unfortunately, Lightning `extract_batch_size` relies on to infer the batch size.

```py
    if isinstance(batch, dict):
        sample = next(iter(batch.values()), 1)
```

As you are using an OrderedDict, I would advise you to provide the batch as the first key:

```py
def collate_fn(batch):
    out = OrderedDict()
    # this will be the first item extract_batch_size finds, and will infer wrong batch_size
    out['batch'] = torch.stack(batch)
    out['confuse_extract_batch_size'] = torch.empty(size=(FAKE_BATCH_SIZE, 1))
    return out
```

Best,
T.C

Dear @ronif,

But I agree the batch_size should be reset to `1` on epoch end as the responsability is left on the user to properly compute the metric.

Would you be interested in making a fix ?

Best,
T.C
@tchaton a bit busy right now, but will try to take a look next week if no one else wants it by then. I think this fix will require some additional test coverage because it modifies critical paths (valiation metric calculations), so I don't want to rush it.
Dear @ronif,

I thought the simplest solution would be to provide a warning in case the batch_size is ambiguous to resolve.

Best,
T.C

Commit message: Update pytorch_lightning/utilities/data.py

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 60702a2866c83..85385072aad57 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -125,6 +125,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added support for `devices="auto"` ([#10264](https://github.com/PyTorchLightning/pytorch-lightning/pull/10264))
 - Added a `filename` argument in `ModelCheckpoint.format_checkpoint_name` ([#9818](https://github.com/PyTorchLightning/pytorch-lightning/pull/9818))
 - Added `configure_columns` method to `RichProgressBar` ([#10288](https://github.com/PyTorchLightning/pytorch-lightning/pull/10288))
+- Added a warning if multiple batch sizes are found from ambiguous batch ([#10247](https://github.com/PyTorchLightning/pytorch-lightning/pull/10247))
 
 
 ### Changed
diff --git a/pytorch_lightning/utilities/data.py b/pytorch_lightning/utilities/data.py
index 4669cc2020b16..13f39d9a482c5 100644
--- a/pytorch_lightning/utilities/data.py
+++ b/pytorch_lightning/utilities/data.py
@@ -18,28 +18,46 @@
 from torch.utils.data import DataLoader, IterableDataset
 
 from pytorch_lightning.utilities import rank_zero_warn
+from pytorch_lightning.utilities.warnings import WarningCache
 
 BType = Union[torch.Tensor, str, Mapping[Any, "BType"], Iterable["BType"]]
 
+warning_cache = WarningCache()
+
+
+def _extract_batch_size(batch: BType) -> int:
+    if isinstance(batch, torch.Tensor):
+        yield batch.size(0)
+    elif isinstance(batch, str):
+        yield len(batch)
+    elif isinstance(batch, (Iterable, Mapping)):
+        if isinstance(batch, Mapping):
+            batch = batch.values()
+
+        for sample in batch:
+            yield from _extract_batch_size(sample)
+    else:
+        yield 1
+
 
 def extract_batch_size(batch: BType) -> int:
-    """Recursively unpack a batch to find a torch.Tensor.
+    """Unpack a batch to find a ``torch.Tensor``.
 
     Returns:
         ``len(tensor)`` when found, or ``1`` when it hits an empty or non iterable.
     """
-    if isinstance(batch, torch.Tensor):
-        return batch.size(0)
-    if isinstance(batch, str):
-        return len(batch)
-    if isinstance(batch, dict):
-        sample = next(iter(batch.values()), 1)
-        return extract_batch_size(sample)
-    if isinstance(batch, Iterable):
-        sample = next(iter(batch), 1)
-        return extract_batch_size(sample)
-
-    return 1
+    batch_size = None
+    for bs in _extract_batch_size(batch):
+        if batch_size is None:
+            batch_size = bs
+        elif batch_size != bs:
+            warning_cache.warn(
+                "Trying to infer the `batch_size` from an ambiguous collection. The batch size we"
+                f" found is {batch_size}. To avoid any miscalculations, use `self.log(..., batch_size=batch_size)`."
+            )
+            break
+
+    return batch_size
 
 
 def has_iterable_dataset(dataloader: DataLoader) -> bool:
diff --git a/tests/utilities/test_data.py b/tests/utilities/test_data.py
index 6c5b5a8e33ffb..098012d128167 100644
--- a/tests/utilities/test_data.py
+++ b/tests/utilities/test_data.py
@@ -2,7 +2,7 @@
 import torch
 from torch.utils.data.dataloader import DataLoader
 
-from pytorch_lightning.utilities.data import extract_batch_size, get_len, has_iterable_dataset, has_len
+from pytorch_lightning.utilities.data import extract_batch_size, get_len, has_iterable_dataset, has_len, warning_cache
 from tests.helpers.boring_model import RandomDataset, RandomIterableDataset
 
 
@@ -23,6 +23,22 @@ def test_extract_batch_size():
     batch = {"test": [{"test": [torch.zeros(11, 10)]}]}
     assert extract_batch_size(batch) == 11
 
+    batch = {"test": [{"test": [torch.zeros(11, 10), torch.zeros(10, 10)]}]}
+    with pytest.warns(UserWarning, match="Trying to infer the `batch_size` .* we found is 11."):
+        extract_batch_size(batch)
+
+    warning_cache.clear()
+
+    batch = {"test": [{"test": [torch.zeros(10, 10), torch.zeros(11, 10)]}]}
+    with pytest.warns(UserWarning, match="Trying to infer the `batch_size` .* we found is 10."):
+        extract_batch_size(batch)
+
+    warning_cache.clear()
+
+    batch = [{"test": torch.zeros(10, 10), "test_1": torch.zeros(11, 10)}]
+    with pytest.warns(UserWarning, match="Trying to infer the `batch_size` .* we found is 10."):
+        extract_batch_size(batch)
+
 
 def test_has_iterable_dataset():
     assert has_iterable_dataset(DataLoader(RandomIterableDataset(1, 1)))


Code review comments:
diff_hunk: @@ -23,6 +23,22 @@ def test_extract_batch_size():
     batch = {"test": [{"test": [torch.zeros(11, 10)]}]}
     assert extract_batch_size(batch) == 11
 
+    batch = {"test": [{"test": [torch.zeros(11, 10), torch.zeros(10, 10)]}]}
+    with pytest.warns(UserWarning, match="Trying to infer the `batch_size` .* we found is 11."):
+        extract_batch_size(batch)
+
+    warning_cache.clear()
+
+    batch = {"test": [{"test": [torch.zeros(10, 10), torch.zeros(11, 10)]}]}
+    with pytest.warns(UserWarning, match="Trying to infer the `batch_size` .* we found is 10."):
+        extract_batch_size(batch)
+
+    warning_cache.clear()
+
+    batch = [{"test": torch.zeros(10, 10), "test_1": torch.zeros(11, 10)}]
+    with pytest.warns(UserWarning, match="Trying to infer the `batch_size` .* we found is 10."):
path: tests/utilities/test_data.py
text: @user1:
I think this test needs a check that no warning is raised if the collection has a uniform batch size