Repository: dmlc/dgl

Issue:
[Example][Bug] Running error on the example case: example/pytorch/caregnn
## 🐛 Bug

Example for CARE-GNN model is crashed when training with mini-batch. Full graph can pass though.
From the core-dump stack, it seems one function, sample_blocks() (/dgl/dataloading/base.py), has not been implemented?

## To Reproduce
`python main_sampling.py --gpu 0`
```
Traceback (most recent call last):
  File "main_sampling.py", line 194, in <module>
    main(args)
  File "main_sampling.py", line 115, in main
    for input_nodes, output_nodes, blocks in train_dataloader:
  File "/opt/conda/lib/python3.8/site-packages/dgl-0.9-py3.8-linux-x86_64.egg/dgl/dataloading/dataloader.py", line 498, in __next__
    self._next_non_threaded() if not self.use_thread else self._next_threaded()
  File "/opt/conda/lib/python3.8/site-packages/dgl-0.9-py3.8-linux-x86_64.egg/dgl/dataloading/dataloader.py", line 478, in _next_non_threaded
    batch = next(self.dataloader_it)
  File "/opt/conda/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 578, in __next__
    data = self._next_data()
  File "/opt/conda/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 618, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/opt/conda/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 40, in fetch
    return self.collate_fn(data)
  File "/opt/conda/lib/python3.8/site-packages/dgl-0.9-py3.8-linux-x86_64.egg/dgl/dataloading/dataloader.py", line 522, in __call__
    batch = self.sample_func(self.g, items)
  File "/opt/conda/lib/python3.8/site-packages/dgl-0.9-py3.8-linux-x86_64.egg/dgl/dataloading/base.py", line 243, in sample
    result = self.sample_blocks(g, seed_nodes, exclude_eids=exclude_eids)
  File "/opt/conda/lib/python3.8/site-packages/dgl-0.9-py3.8-linux-x86_64.egg/dgl/dataloading/base.py", line 230, in sample_blocks
    raise NotImplementedError
NotImplementedError

```

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

It should run through.

## Environment

 - DGL Version (e.g., 1.0): 0.9
 - Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3): 1.12
 - OS (e.g., Linux): ubuntu
 - How you installed DGL (`conda`, `pip`, source): source
 - Build command you used (if compiling from source):
 - Python version: 3.8
 - CUDA/cuDNN version (if applicable): 11.7
 - GPU models and configuration (e.g. V100): A100
 - Any other relevant information:



Discussion:
It uses a custom sampler https://github.com/dmlc/dgl/blob/0f0e7c7fad6ec1399a67881ce9a7e60d14cfe1a0/examples/pytorch/caregnn/model_sampling.py#L14 that needs to be updated.

Resolving patch:
diff --git a/examples/pytorch/caregnn/main_sampling.py b/examples/pytorch/caregnn/main_sampling.py
--- a/examples/pytorch/caregnn/main_sampling.py
+++ b/examples/pytorch/caregnn/main_sampling.py
@@ -92,9 +92,10 @@ def main(args):
             graph.ndata['nd'] = th.tanh(model.layers[i].MLP(layers_feat[i]))
             for etype in graph.canonical_etypes:
                 graph.apply_edges(_l1_dist, etype=etype)
-                dist[etype] = graph.edges[etype].data['ed']
+                dist[etype] = graph.edges[etype].data.pop('ed').detach().cpu()
             dists.append(dist)
             p.append(model.layers[i].p)
+        graph.ndata.pop('nd')
         sampler = CARESampler(p, dists, args.num_layers)
 
         # train
@@ -103,14 +104,9 @@ def main(args):
         tr_recall = 0
         tr_auc = 0
         tr_blk = 0
-        train_dataloader = dgl.dataloading.DataLoader(graph,
-                                                          train_idx,
-                                                          sampler,
-                                                          batch_size=args.batch_size,
-                                                          shuffle=True,
-                                                          drop_last=False,
-                                                          num_workers=args.num_workers
-                                                          )
+        train_dataloader = dgl.dataloading.DataLoader(
+            graph, train_idx, sampler, batch_size=args.batch_size,
+            shuffle=True,  drop_last=False, num_workers=args.num_workers)
 
         for input_nodes, output_nodes, blocks in train_dataloader:
             blocks = [b.to(device) for b in blocks]
@@ -135,14 +131,9 @@ def main(args):
 
         # validation
         model.eval()
-        val_dataloader = dgl.dataloading.DataLoader(graph,
-                                                        val_idx,
-                                                        sampler,
-                                                        batch_size=args.batch_size,
-                                                        shuffle=True,
-                                                        drop_last=False,
-                                                        num_workers=args.num_workers
-                                                        )
+        val_dataloader = dgl.dataloading.DataLoader(
+            graph, val_idx, sampler, batch_size=args.batch_size,
+            shuffle=True, drop_last=False, num_workers=args.num_workers)
 
         val_recall, val_auc, val_loss = evaluate(model, loss_fn, val_dataloader, device)
 
@@ -159,14 +150,9 @@ def main(args):
     model.eval()
     if args.early_stop:
         model.load_state_dict(th.load('es_checkpoint.pt'))
-    test_dataloader = dgl.dataloading.DataLoader(graph,
-                                                     test_idx,
-                                                     sampler,
-                                                     batch_size=args.batch_size,
-                                                     shuffle=True,
-                                                     drop_last=False,
-                                                     num_workers=args.num_workers
-                                                     )
+    test_dataloader = dgl.dataloading.DataLoader(
+        graph, test_idx, sampler, batch_size=args.batch_size,
+        shuffle=True, drop_last=False, num_workers=args.num_workers)
 
     test_recall, test_auc, test_loss = evaluate(model, loss_fn, test_dataloader, device)
 
diff --git a/examples/pytorch/caregnn/model_sampling.py b/examples/pytorch/caregnn/model_sampling.py
--- a/examples/pytorch/caregnn/model_sampling.py
+++ b/examples/pytorch/caregnn/model_sampling.py
@@ -13,9 +13,10 @@ def _l1_dist(edges):
 
 class CARESampler(dgl.dataloading.BlockSampler):
     def __init__(self, p, dists, num_layers):
-        super().__init__(num_layers)
+        super().__init__()
         self.p = p
         self.dists = dists
+        self.num_layers = num_layers
 
     def sample_frontier(self, block_id, g, seed_nodes, *args, **kwargs):
         with g.local_scope():
@@ -28,7 +29,7 @@ def sample_frontier(self, block_id, g, seed_nodes, *args, **kwargs):
                     num_neigh = th.ceil(g.in_degrees(node, etype=etype) * self.p[block_id][etype]).int().item()
                     neigh_dist = self.dists[block_id][etype][edges]
                     if neigh_dist.shape[0] > num_neigh:
-                        neigh_index = np.argpartition(neigh_dist.cpu().detach(), num_neigh)[:num_neigh]
+                        neigh_index = np.argpartition(neigh_dist, num_neigh)[:num_neigh]
                     else:
                         neigh_index = np.arange(num_neigh)
                     edge_mask[edges[neigh_index]] = 1
@@ -36,6 +37,19 @@ def sample_frontier(self, block_id, g, seed_nodes, *args, **kwargs):
 
             return dgl.edge_subgraph(g, new_edges_masks, relabel_nodes=False)
 
+    def sample_blocks(self, g, seed_nodes, exclude_eids=None):
+        output_nodes = seed_nodes
+        blocks = []
+        for block_id in reversed(range(self.num_layers)):
+            frontier = self.sample_frontier(block_id, g, seed_nodes)
+            eid = frontier.edata[dgl.EID]
+            block = dgl.to_block(frontier, seed_nodes)
+            block.edata[dgl.EID] = eid
+            seed_nodes = block.srcdata[dgl.NID]
+            blocks.insert(0, block)
+
+        return seed_nodes, output_nodes, blocks
+
     def __len__(self):
         return self.num_layers
 


Test patch: