add exclude self option for EdgeDataLoader (dmlc#3122)

Co-authored-by: Minjie Wang <[email protected]>

python/dgl/dataloading/dataloader.py

@@ -77,6 +77,9 @@ def _find_exclude_eids(g, exclude_mode, eids, **kwargs):
         None (default)
             Does not exclude any edge.
 
+        'self'
+            Exclude the given edges themselves but nothing else.
+
         'reverse_id'
             Exclude all edges specified in ``eids``, as well as their reverse edges
             of the same edge type.
@@ -105,6 +108,8 @@ def _find_exclude_eids(g, exclude_mode, eids, **kwargs):
     """
     if exclude_mode is None:
         return None
+    elif exclude_mode == 'self':
+        return eids
     elif exclude_mode == 'reverse_id':
         return _find_exclude_eids_with_reverse_id(g, eids, kwargs['reverse_eid_map'])
     elif exclude_mode == 'reverse_types':
@@ -493,6 +498,8 @@ class EdgeCollator(Collator):
 
         * None, which excludes nothing.
 
+        * ``'self'``, which excludes the sampled edges themselves but nothing else.
+
         * ``'reverse_id'``, which excludes the reverse edges of the sampled edges.  The said
           reverse edges have the same edge type as the sampled edges.  Only works
           on edge types whose source node type is the same as its destination node type.

python/dgl/dataloading/pytorch/dataloader.py

@@ -607,6 +607,7 @@ class EdgeDataLoader:
         minibatch.  Possible values are
 
         * None,
+        * ``self``,
         * ``reverse_id``,
         * ``reverse_types``
 

tests/pytorch/test_dataloader.py

@@ -121,6 +121,12 @@ def test_neighbor_sampler_dataloader():
         nids.append({'follow': seeds})
         modes.append('edge')
 
+        collators.append(dgl.dataloading.EdgeCollator(
+            g, seeds, sampler, exclude='self'))
+        graphs.append(g)
+        nids.append({'follow': seeds})
+        modes.append('edge')
+
         collators.append(dgl.dataloading.EdgeCollator(
             g, seeds, sampler, exclude='reverse_id', reverse_eids=reverse_eids))
         graphs.append(g)
@@ -133,6 +139,12 @@ def test_neighbor_sampler_dataloader():
         nids.append({'follow': seeds})
         modes.append('link')
 
+        collators.append(dgl.dataloading.EdgeCollator(
+            g, seeds, sampler, exclude='self', negative_sampler=dgl.dataloading.negative_sampler.Uniform(2)))
+        graphs.append(g)
+        nids.append({'follow': seeds})
+        modes.append('link')
+
         collators.append(dgl.dataloading.EdgeCollator(
             g, seeds, sampler, exclude='reverse_id', reverse_eids=reverse_eids,
             negative_sampler=dgl.dataloading.negative_sampler.Uniform(2)))