Repository: dmlc/dgl
Pull request: [DataLoader] Allow batch_size=None for GraphDataLoader
Fixes #3358.

Problem statement:
Support `batch_size = None` for GraphDataLoader
## 🚀 Feature
Supporting `None` batch size for unbatched parallel data loading, as the PyTorch DataLoader supports.

## Motivation
Currently, when passing None as `batch_size` for GraphDataLoader, it will give an error
```
  ...
  File "dgl/dataloading/dataloader.py", line 823, in collate
    elem = items[0]
  File "dgl/heterograph.py", line 2231, in __getitem__
    raise DGLError('Invalid key "{}". Must be one of the edge types.'.format(orig_key))
dgl._ffi.base.DGLError: Invalid key "0". Must be one of the edge types.
```

In some cases, for example as we want to preprocess data in parallel, it is simpler to use None as batch size so we do not need to implement the parallel data pipeline ourselves.

## Alternatives
A batch size of 1 with manual unbatching may be used, but it is less convenient.

Hints from issue discussion:
Actually in DGL a single graph always has a batch size of 1.  So you don't have to do manual unbatching.  Setting batch size to 1 and then directly do stuff on the sampled minibatch should work.
We have graph level data, such as graph classification labels.

Commit message: Update dataloader.py

Patch to review:
diff --git a/python/dgl/dataloading/dataloader.py b/python/dgl/dataloading/dataloader.py
index 4eed99587c6f..478c135f20d5 100644
--- a/python/dgl/dataloading/dataloader.py
+++ b/python/dgl/dataloading/dataloader.py
@@ -1197,17 +1197,18 @@ def __init__(self, dataset, collate_fn=None, use_ddp=False, ddp_seed=0, **kwargs
             else:
                 dataloader_kwargs[k] = v
 
-        if collate_fn is None:
-            self.collate = GraphCollator(**collator_kwargs).collate
-        else:
-            self.collate = collate_fn
-
         self.use_ddp = use_ddp
         if use_ddp:
             self.dist_sampler = _create_dist_sampler(dataset, dataloader_kwargs, ddp_seed)
             dataloader_kwargs['sampler'] = self.dist_sampler
 
-        super().__init__(dataset=dataset, collate_fn=self.collate, **dataloader_kwargs)
+        super().__init__(dataset=dataset, collate_fn=collate_fn, **dataloader_kwargs)
+
+        # Overwrite PyTorch's default collator function with ours.
+        # After calling super().__init__() PyTorch should have set the collate_fn attribute.
+        # pylint: disable=access-member-before-definition
+        if self.collate_fn is torch.utils.data._utils.collate.default_collate:
+            self.collate_fn = GraphCollator(**collator_kwargs).collate
 
     def set_epoch(self, epoch):
         """Sets the epoch number for the underlying sampler which ensures all replicas
diff --git a/tests/pytorch/test_dataloader.py b/tests/pytorch/test_dataloader.py
index 2f15ec6c18bf..9832180584a3 100644
--- a/tests/pytorch/test_dataloader.py
+++ b/tests/pytorch/test_dataloader.py
@@ -12,15 +12,19 @@
 import pytest
 
 
-def test_graph_dataloader():
-    batch_size = 16
+@pytest.mark.parametrize('batch_size', [None, 16])
+def test_graph_dataloader(batch_size):
     num_batches = 2
-    minigc_dataset = dgl.data.MiniGCDataset(batch_size * num_batches, 10, 20)
+    num_samples = num_batches * (batch_size if batch_size is not None else 1)
+    minigc_dataset = dgl.data.MiniGCDataset(num_samples, 10, 20)
     data_loader = dgl.dataloading.GraphDataLoader(minigc_dataset, batch_size=batch_size, shuffle=True)
     assert isinstance(iter(data_loader), Iterator)
     for graph, label in data_loader:
         assert isinstance(graph, dgl.DGLGraph)
-        assert F.asnumpy(label).shape[0] == batch_size
+        if batch_size is not None:
+            assert F.asnumpy(label).shape[0] == batch_size
+        else:
+            assert F.asnumpy(label).ndim == 0
 
 @unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
 @pytest.mark.parametrize('num_workers', [0, 4])


Code review comments:
diff_hunk: @@ -1197,17 +1197,18 @@ def __init__(self, dataset, collate_fn=None, use_ddp=False, ddp_seed=0, **kwargs
             else:
                 dataloader_kwargs[k] = v
 
-        if collate_fn is None:
-            self.collate = GraphCollator(**collator_kwargs).collate
-        else:
-            self.collate = collate_fn
-
         self.use_ddp = use_ddp
         if use_ddp:
             self.dist_sampler = _create_dist_sampler(dataset, dataloader_kwargs, ddp_seed)
             dataloader_kwargs['sampler'] = self.dist_sampler
 
-        super().__init__(dataset=dataset, collate_fn=self.collate, **dataloader_kwargs)
+        super().__init__(dataset=dataset, collate_fn=collate_fn, **dataloader_kwargs)
+
+        # Overwrite PyTorch's default collator function with ours.
+        # After calling super().__init__() PyTorch should have set the collate_fn attribute.
+        # pylint: disable=access-member-before-definition
+        if self.collate_fn is torch.utils.data._utils.collate.default_collate:
+            self.collate_fn = GraphCollator(**collator_kwargs).collate
path: python/dgl/dataloading/dataloader.py
text: @user1:
I'm concerned we're relying on implementation details of torch's DataLoader, both in how it chooses to set its `self.collate_fn`, and that it make no other decisions in the constructor based on this. Could we instead explicitly check if batch_size is None, and set our collate_fn input to Torch's DataLoader based on that?
```python3
if batch_size is not None and collate_fn is None:
    collate_fn = GraphCollator(**collator_kwargs).collate
```
Or is there another reason we need to check what collator the parent DataLoader chooses?

@author:
Makes sense.