Repository: dmlc/dgl

Issue:
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.



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.

Resolving patch:
diff --git a/python/dgl/dataloading/dataloader.py b/python/dgl/dataloading/dataloader.py
--- a/python/dgl/dataloading/dataloader.py
+++ b/python/dgl/dataloading/dataloader.py
@@ -1124,17 +1124,15 @@ 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)
+        if collate_fn is None and kwargs.get('batch_size', 1) is not None:
+            collate_fn = GraphCollator(**collator_kwargs).collate
+
+        super().__init__(dataset=dataset, collate_fn=collate_fn, **dataloader_kwargs)
 
     def set_epoch(self, epoch):
         """Sets the epoch number for the underlying sampler which ensures all replicas


Test patch:
diff --git a/tests/pytorch/test_dataloader.py b/tests/pytorch/test_dataloader.py
--- a/tests/pytorch/test_dataloader.py
+++ b/tests/pytorch/test_dataloader.py
@@ -12,15 +12,21 @@
 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:
+            # If batch size is None, the label element will be a single scalar following
+            # PyTorch's practice.
+            assert F.asnumpy(label).ndim == 0
 
 @unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
 @pytest.mark.parametrize('num_workers', [0, 4])
diff --git a/tests/pytorch/test_nn.py b/tests/pytorch/test_nn.py
--- a/tests/pytorch/test_nn.py
+++ b/tests/pytorch/test_nn.py
@@ -1516,7 +1516,7 @@ def test_hgt(idtype, in_size, num_heads):
     sorted_y = m(sorted_g, sorted_x, sorted_ntype, sorted_etype, presorted=False)
     assert sorted_y.shape == (g.num_nodes(), head_size * num_heads)
     # mini-batch
-    train_idx = th.randint(0, 100, (10, ), dtype = idtype)
+    train_idx = th.randperm(100, dtype = idtype)[:10]
     sampler = dgl.dataloading.NeighborSampler([-1])
     train_loader = dgl.dataloading.DataLoader(g, train_idx.to(dev), sampler,
                                             batch_size=8, device=dev,