Repository: dmlc/dgl

Issue:
DDP Tensorized Dataset fails when passing a dictionary of indices.
## 🐛 Bug

Specifying a dictionary of indices to load from (mapping node types to tensor arrays) causes DDP datasets to fail

## To Reproduce

Steps to reproduce the behavior:

1. Create a node dataloader with use_ddp=True and an index set specified via dictionary (mapping node type to tensors). The mapping can be only from a single type to a single tensor array, and this will still fail.
2. Attempt to load data.


## Expected behavior

DDP Dataloading should function as before 0.8, where specifying a dictionary of indices causes no issue.

## Environment

DGL 0.8
Pytorch 
Linux
Source

## Additional context

The functionality to create a dataloader with a dictionary of indices existed prior to 0.8, and worked correctly.



Discussion:
This is the exact error we're seeing:
```
raceback (most recent call last):
  File "/opt/conda/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 69, in _wrap
    fn(i, *args)
  File "/workspace/gp/workflow/trainers/utils.py", line 4, in run
    trainer.fit_process(rank)
  File "/workspace/gp/workflow/trainers/base_trainer.py", line 99, in fit_process
    self.setup(rank)
  File "/workspace/gp/workflow/trainers/base_trainer.py", line 144, in setup
    self.data_object.build_train_dataloader_post_dist(device, True)
  File "/workspace/gp/data_loaders/data_object.py", line 149, in build_train_dataloader_post_dist
    self.train_dataloader = dgl.dataloading.NodeDataLoader(
  File "/opt/conda/lib/python3.8/site-packages/dgl-0.8-py3.8-linux-x86_64.egg/dgl/dataloading/dataloader.py", line 778, in __init__
    self.dataset = create_tensorized_dataset(
  File "/opt/conda/lib/python3.8/site-packages/dgl-0.8-py3.8-linux-x86_64.egg/dgl/dataloading/dataloader.py", line 554, in create_tensorized_dataset
    return DDPTensorizedDataset(indices, batch_size, drop_last, ddp_seed)
  File "/opt/conda/lib/python3.8/site-packages/dgl-0.8-py3.8-linux-x86_64.egg/dgl/dataloading/dataloader.py", line 208, in __init__
    self._indices[:self._id_tensor.shape[0]] = torch.arange(self._id_tensor.shape[0])
RuntimeError: The expanded size of the tensor (2) must match the existing size (629571) at non-singleton dimension 0.  Target sizes: [2].  Tensor sizes: [629571]
```
It looks like only part of DDPTensorizedDataset's constructor handles the indices being mapping.

Resolving patch:
diff --git a/examples/pytorch/graphsage/multi_gpu_node_classification.py b/examples/pytorch/graphsage/multi_gpu_node_classification.py
--- a/examples/pytorch/graphsage/multi_gpu_node_classification.py
+++ b/examples/pytorch/graphsage/multi_gpu_node_classification.py
@@ -36,7 +36,7 @@ def inference(self, g, device, batch_size, num_workers, buffer_device=None):
         # example is that the intermediate results can also benefit from prefetching.
         g.ndata['h'] = g.ndata['feat']
         sampler = dgl.dataloading.MultiLayerFullNeighborSampler(1, prefetch_node_feats=['h'])
-        dataloader = dgl.dataloading.NodeDataLoader(
+        dataloader = dgl.dataloading.DataLoader(
                 g, torch.arange(g.num_nodes()).to(g.device), sampler, device=device,
                 batch_size=1000, shuffle=False, drop_last=False, num_workers=num_workers,
                 persistent_workers=(num_workers > 0))
@@ -77,7 +77,7 @@ def train(rank, world_size, graph, num_classes, split_idx):
             graph, train_idx, sampler,
             device='cuda', batch_size=1000, shuffle=True, drop_last=False,
             num_workers=0, use_ddp=True, use_uva=True)
-    valid_dataloader = dgl.dataloading.NodeDataLoader(
+    valid_dataloader = dgl.dataloading.DataLoader(
             graph, valid_idx, sampler, device='cuda', batch_size=1024, shuffle=True,
             drop_last=False, num_workers=0, use_uva=True)
 
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
@@ -169,8 +169,10 @@ class DDPTensorizedDataset(torch.utils.data.IterableDataset):
     def __init__(self, indices, batch_size, drop_last, ddp_seed):
         if isinstance(indices, Mapping):
             self._mapping_keys = list(indices.keys())
+            len_indices = sum(len(v) for v in indices.values())
         else:
             self._mapping_keys = None
+            len_indices = len(indices)
 
         self.rank = dist.get_rank()
         self.num_replicas = dist.get_world_size()
@@ -179,17 +181,17 @@ def __init__(self, indices, batch_size, drop_last, ddp_seed):
         self.batch_size = batch_size
         self.drop_last = drop_last
 
-        if self.drop_last and len(indices) % self.num_replicas != 0:
-            self.num_samples = math.ceil((len(indices) - self.num_replicas) / self.num_replicas)
+        if self.drop_last and len_indices % self.num_replicas != 0:
+            self.num_samples = math.ceil((len_indices - self.num_replicas) / self.num_replicas)
         else:
-            self.num_samples = math.ceil(len(indices) / self.num_replicas)
+            self.num_samples = math.ceil(len_indices / self.num_replicas)
         self.total_size = self.num_samples * self.num_replicas
         # If drop_last is True, we create a shared memory array larger than the number
         # of indices since we will need to pad it after shuffling to make it evenly
         # divisible before every epoch.  If drop_last is False, we create an array
         # with the same size as the indices so we can trim it later.
-        self.shared_mem_size = self.total_size if not self.drop_last else len(indices)
-        self.num_indices = len(indices)
+        self.shared_mem_size = self.total_size if not self.drop_last else len_indices
+        self.num_indices = len_indices
 
         if isinstance(indices, Mapping):
             self._device = next(iter(indices.values())).device


Test patch: