Repository: pyg-team/pytorch_geometric
Pull request #8137 (2023-10-06 04:29:27)

Raise error if `len(DynamicBatchSampler())` is ambiguous
Part of a fix for #7724.

Assuming `len(DynamicBatchSampler(dataset)) == len(dataset)` is wrong in AFAIK all cases `max_num>1`, which can lead to undesired behaviour in 3rd party libraries like PyTorch Lightning as reported in the linked issue.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8137?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#8137](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8137?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (a63d7a2) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/aae8dfd0d73341e9214200808c92e5ed1a7a2f1c?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (aae8dfd) will **decrease** coverage by `0.01%`.
> The diff coverage is `n/a`.

```diff
@@            Coverage Diff             @@
##           master    #8137      +/-   ##
==========================================
- Coverage   88.14%   88.14%   -0.01%     
==========================================
  Files         471      471              
  Lines       28274    28272       -2     
==========================================
- Hits        24922    24920       -2     
  Misses       3352     3352              
```


| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8137?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/loader/dynamic\_batch\_sampler.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8137?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2xvYWRlci9keW5hbWljX2JhdGNoX3NhbXBsZXIucHk=) | `95.34% <ø> (+2.01%)` | :arrow_up: |

... and [1 file with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8137/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 09c60e2b1..114a8dafd 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -117,6 +117,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Fix `DynamicBatchSampler.__len__` to raise an error in case `num_steps` is undefined ([#8137](https://github.com/pyg-team/pytorch_geometric/pull/8137))
 - Enabled pickling of `DimeNet` models ([#8019](https://github.com/pyg-team/pytorch_geometric/pull/8019))
 - Changed the `trim_to_layer` function to filter out non-reachable node and edge types when operating on heterogeneous graphs ([#7942](https://github.com/pyg-team/pytorch_geometric/pull/7942))
 - Accelerated and simplified `top_k` computation in `TopKPooling` ([#7737](https://github.com/pyg-team/pytorch_geometric/pull/7737))
diff --git a/torch_geometric/loader/dynamic_batch_sampler.py b/torch_geometric/loader/dynamic_batch_sampler.py
index 0622706a9..6df8f1432 100644
--- a/torch_geometric/loader/dynamic_batch_sampler.py
+++ b/torch_geometric/loader/dynamic_batch_sampler.py
@@ -1,4 +1,3 @@
-import warnings
 from typing import Iterator, List, Optional
 
 import torch
@@ -38,20 +37,23 @@ class DynamicBatchSampler(torch.utils.data.sampler.Sampler):
         num_steps (int, optional): The number of mini-batches to draw for a
             single epoch. If set to :obj:`None`, will iterate through all the
             underlying examples, but :meth:`__len__` will be :obj:`None` since
-            it is be ambiguous. (default: :obj:`None`)
+            it is ambiguous. (default: :obj:`None`)
     """
-    def __init__(self, dataset: Dataset, max_num: int, mode: str = 'node',
-                 shuffle: bool = False, skip_too_big: bool = False,
-                 num_steps: Optional[int] = None):
-        if not isinstance(max_num, int) or max_num <= 0:
-            raise ValueError("`max_num` should be a positive integer value "
-                             "(got {max_num}).")
+    def __init__(
+        self,
+        dataset: Dataset,
+        max_num: int,
+        mode: str = 'node',
+        shuffle: bool = False,
+        skip_too_big: bool = False,
+        num_steps: Optional[int] = None,
+    ):
+        if max_num <= 0:
+            raise ValueError(f"`max_num` should be a positive integer value "
+                             f"(got {max_num})")
         if mode not in ['node', 'edge']:
-            raise ValueError("`mode` choice should be either "
-                             f"'node' or 'edge' (got '{mode}').")
-
-        if num_steps is None:
-            num_steps = len(dataset)
+            raise ValueError(f"`mode` choice should be either "
+                             f"'node' or 'edge' (got '{mode}')")
 
         self.dataset = dataset
         self.max_num = max_num
@@ -59,48 +61,47 @@ class DynamicBatchSampler(torch.utils.data.sampler.Sampler):
         self.shuffle = shuffle
         self.skip_too_big = skip_too_big
         self.num_steps = num_steps
+        self.max_steps = num_steps or len(dataset)
 
     def __iter__(self) -> Iterator[List[int]]:
-        batch = []
-        batch_n = 0
-        num_steps = 0
-        num_processed = 0
-
         if self.shuffle:
-            indices = torch.randperm(len(self.dataset), dtype=torch.long)
+            indices = torch.randperm(len(self.dataset)).tolist()
         else:
-            indices = torch.arange(len(self.dataset), dtype=torch.long)
+            indices = range(len(self.dataset))
+
+        samples: List[int] = []
+        current_num: int = 0
+        num_steps: int = 0
+        num_processed: int = 0
 
         while (num_processed < len(self.dataset)
-               and num_steps < self.num_steps):
-            # Fill batch
-            for idx in indices[num_processed:]:
-                # Size of sample
-                data = self.dataset[idx]
-                n = data.num_nodes if self.mode == 'node' else data.num_edges
-
-                if batch_n + n > self.max_num:
-                    if batch_n == 0:
+               and num_steps < self.max_steps):
+
+            for i in indices[num_processed:]:
+                data = self.dataset[i]
+                num = data.num_nodes if self.mode == 'node' else data.num_edges
+
+                if current_num + num > self.max_num:
+                    if current_num == 0:
                         if self.skip_too_big:
                             continue
-                        else:
-                            warnings.warn("Size of data sample at index "
-                                          f"{idx} is larger than "
-                                          f"{self.max_num} {self.mode}s "
-                                          f"(got {n} {self.mode}s.")
-                    else:
-                        # Mini-batch filled
+                    else:  # Mini-batch filled:
                         break
 
-                # Add sample to current batch
-                batch.append(idx.item())
+                samples.append(i)
                 num_processed += 1
-                batch_n += n
+                current_num += num
 
-            yield batch
-            batch = []
-            batch_n = 0
+            yield samples
+            samples: List[int] = []
+            current_num = 0
             num_steps += 1
 
     def __len__(self) -> int:
+        if self.num_steps is None:
+            raise ValueError(f"The length of '{self.__class__.__name__}' is "
+                             f"undefined since the number of steps per epoch "
+                             f"is ambiguous. Either specify `num_steps` or "
+                             f"use a static batch sampler.")
+
         return self.num_steps


Test patch:
diff --git a/test/loader/test_dynamic_batch_sampler.py b/test/loader/test_dynamic_batch_sampler.py
index 57108335d..6c2a4285b 100644
--- a/test/loader/test_dynamic_batch_sampler.py
+++ b/test/loader/test_dynamic_batch_sampler.py
@@ -33,13 +33,6 @@ def test_dataloader_with_dynamic_batches():
         num_nodes_total += data.num_nodes
     assert num_nodes_total == 404
 
-    # Test warning
-    batch_sampler = DynamicBatchSampler(data_list, 300, skip_too_big=False,
-                                        num_steps=2)
-    loader = DataLoader(data_list, batch_sampler=batch_sampler)
-
-    with pytest.warns(UserWarning, match="is larger than 300 nodes"):
-        num_nodes_total = 0
-        for data in loader:
-            num_nodes_total += data.num_nodes
-        assert num_nodes_total == 601
+    with pytest.raises(ValueError, match="length of 'DynamicBatchSampler'"):
+        len(DynamicBatchSampler(data_list, max_num=300))
+    assert len(DynamicBatchSampler(data_list, max_num=300, num_steps=2)) == 2