Repository: dmlc/dgl

Issue:
pytorch/tgn example fails: cannot import name 'EdgeCollator' from 'dgl.dataloading.dataloader'
## 🐛 Bug

I got an error saying: cannot import name 'EdgeCollator' from 'dgl.dataloading.dataloader' in the TGN example (wikipedia).
https://github.com/dmlc/dgl/tree/master/examples/pytorch/tgn

I was trying to adapt the code according to the errors I got (several iterations), but it seems like this example hasn't been updated to correspond newer release of dgl.

## To Reproduce

Steps to reproduce the behavior:

1. Clone the repository and follow the instructions in the tutorial:
```
python train.py --dataset wikipedia
```

leading to the following stack trace:
```
Traceback (most recent call last):
  File "train.py", line 12, in <module>
    from dataloading import (FastTemporalEdgeCollator, FastTemporalSampler,
  File "/mnt/DS_SHARED/users/dvirb/dev/code/graph_learning/dgl/examples/pytorch/tgn/dataloading.py", line 4, in <module>
    from dgl.dataloading.dataloader import EdgeCollator
ImportError: cannot import name 'EdgeCollator' from 'dgl.dataloading.dataloader' (/opt/conda/lib/python3.8/site-packages/dgl/dataloading/dataloader.py)
```

## Expected behavior

The training script should run without failures.

## Environment

 - DGL Version: 0.8.1
 - Backend Library & Version: pytorch 1.10
 - OS: Linux
 - How you installed DGL: conda
 - Build command you used (if compiling from source): --
 - Python version: 3.8.10
 - CUDA/cuDNN version: 11.4
 - GPU models and configuration: NVIDIA A100-SXM 40GB
 - Any other relevant information: -- 

## Additional context
### Attempt 1
I tried modifying the import commands on `examples/pytorch/tgn/dataloading.py` as follows:
instead of:

```python
from dgl.dataloading.dataloader import EdgeCollator
from dgl.dataloading import BlockSampler
from dgl.dataloading.pytorch import _pop_subgraph_storage, _pop_storages
from dgl.base import DGLError
```

I tried  importing some of the classes from `_dataloading.dataloader` using:

```python
from dgl._dataloading.dataloader import EdgeCollator
from dgl.dataloading import BlockSampler
from dgl._dataloading.pytorch import _pop_subgraph_storage, _pop_storages
from dgl.base import DGLError
```

yielding the following error:
```
opt/conda/lib/python3.8/site-packages/dgl/dataloading/dataloader.py:909: DGLWarning: EdgeDataLoader directly taking a BlockSampler will be deprecated and it will not support feature prefetching. Please use dgl.dataloading.as_edge_prediction_sampler to wrap it.
  dgl_warning(
Traceback (most recent call last):
  File "train.py", line 200, in <module>
    train_dataloader = TemporalEdgeDataLoader(graph_no_new_node,
  File "/mnt/DS_SHARED/users/dvirb/dev/code/graph_learning/dgl/examples/pytorch/tgn/dataloading.py", line 269, in __init__
    super().__init__(g, eids, graph_sampler, device, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/dgl/dataloading/dataloader.py", line 925, in __init__
    super().__init__(
  File "/opt/conda/lib/python3.8/site-packages/dgl/dataloading/dataloader.py", line 828, in __init__
    super().__init__(
TypeError: __init__() got an unexpected keyword argument 'g_sampling'
```

### Attempt 2
Then, tried adding the following line to the `__init__` method of `dataloading.TemporalEdgeDataLoader` class:
```python
def __init__(self, g, eids, graph_sampler, device='cpu', collator=TemporalEdgeCollator, **kwargs):
        g_sampling = kwargs.pop('g_sampling')
```

followed by this stack trace:
```
Traceback (most recent call last):
  File "train.py", line 200, in <module>
    train_dataloader = TemporalEdgeDataLoader(graph_no_new_node,
  File "/mnt/DS_SHARED/users/dvirb/dev/code/graph_learning/dgl/examples/pytorch/tgn/dataloading.py", line 276, in __init__
    if k in self.collator_arglist:
AttributeError: 'TemporalEdgeDataLoader' object has no attribute 'collator_arglist'
```

### Attempt 3
tried populating `collator_arglist` member of `TemporalEdgeDataLoader` by adding the following line:
```python
self.collator_arglist = inspect.getfullargspec(collator).args
```
to the constructor of `TemporalEdgeDataLoader`.
But, it yields the following stack trace:

```
Traceback (most recent call last):
  File "train.py", line 265, in <module>
    train_loss = train(model, train_dataloader, sampler,
  File "train.py", line 32, in train
    for _, positive_pair_g, negative_pair_g, blocks in dataloader:
  File "/opt/conda/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 521, in __next__
    data = self._next_data()
  File "/opt/conda/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 561, 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 52, in fetch
    return self.collate_fn(data)
  File "/opt/conda/lib/python3.8/site-packages/dgl/_dataloading/dataloader.py", line 897, in collate
    return self._collate_with_negative_sampling(items)
  File "/mnt/DS_SHARED/users/dvirb/dev/code/graph_learning/dgl/examples/pytorch/tgn/dataloading.py", line 208, in _collate_with_negative_sampling
    subg = self.graph_sampler.sample_blocks(self.g_sampling,
  File "/mnt/DS_SHARED/users/dvirb/dev/code/graph_learning/dgl/examples/pytorch/tgn/dataloading.py", line 92, in sample_blocks
    frontier = self.sampler_frontier(0, g, seed_nodes, timestamp)
  File "/mnt/DS_SHARED/users/dvirb/dev/code/graph_learning/dgl/examples/pytorch/tgn/dataloading.py", line 79, in sampler_frontier
    temporal_subgraph.ndata[dgl.NID] = g.ndata[dgl.NID][temp2origin]
  File "/opt/conda/lib/python3.8/site-packages/dgl/view.py", line 67, in __getitem__
    return self._graph._get_n_repr(self._ntid, self._nodes)[key]
  File "/opt/conda/lib/python3.8/site-packages/dgl/frame.py", line 523, in __getitem__
    return self._columns[name].data
KeyError: '_ID'
```



Discussion:
Should be fixed in https://github.com/dmlc/dgl/pull/3972.

Resolving patch:
diff --git a/examples/pytorch/graphsage/node_classification.py b/examples/pytorch/graphsage/node_classification.py
--- a/examples/pytorch/graphsage/node_classification.py
+++ b/examples/pytorch/graphsage/node_classification.py
@@ -127,6 +127,6 @@ def inference(self, g, device, batch_size, num_workers, buffer_device=None):
 with torch.no_grad():
     pred = model.inference(graph, device, 4096, 0, 'cpu')
     pred = pred[test_idx].to(device)
-    label = graph.ndata['label'][test_idx]
+    label = graph.ndata['label'][test_idx].to(device)
     acc = MF.accuracy(pred, label)
     print('Test acc:', acc.item())
diff --git a/examples/pytorch/tgn/dataloading.py b/examples/pytorch/tgn/dataloading.py
--- a/examples/pytorch/tgn/dataloading.py
+++ b/examples/pytorch/tgn/dataloading.py
@@ -3,7 +3,7 @@
 
 from dgl._dataloading.dataloader import EdgeCollator
 from dgl._dataloading import BlockSampler
-from dgl._dataloading.pytorch import _pop_subgraph_storage, _pop_storages
+from dgl._dataloading.pytorch import _pop_subgraph_storage, _pop_storages, EdgeDataLoader
 from dgl.base import DGLError
 
 from functools import partial
@@ -234,7 +234,7 @@ def collator(self, items):
         return result
 
 
-class TemporalEdgeDataLoader(dgl.dataloading.EdgeDataLoader):
+class TemporalEdgeDataLoader(EdgeDataLoader):
     """ TemporalEdgeDataLoader is an iteratable object to generate blocks for temporal embedding
     as well as pos and neg pair graph for memory update.
 
@@ -600,7 +600,7 @@ def collator(self, items):
 # "APAN: Asynchronous Propagation Attention Network for Real-time Temporal Graph Embedding"
 # that will be appeared in SIGMOD 21, code repo https://github.com/WangXuhongCN/APAN
 
-class SimpleTemporalSampler(dgl.dataloading.BlockSampler):
+class SimpleTemporalSampler(BlockSampler):
     '''
     Simple Temporal Sampler just choose the edges that happen before the current timestamp, to build the subgraph of the corresponding nodes.
     And then the sampler uses the simplest static graph neighborhood sampling methods.
@@ -637,7 +637,7 @@ def sample_frontier(self, block_id, g, seed_nodes):
         return frontier
 
 
-class SimpleTemporalEdgeCollator(dgl.dataloading.EdgeCollator):
+class SimpleTemporalEdgeCollator(EdgeCollator):
     '''
     Temporal Edge collator merge the edges specified by eid: items
 


Test patch: