[Dataset] Add transform argument to built-in datasets (#3733)

* Update

* Fix

* Update

python/dgl/data/bitcoinotc.py

@@ -39,6 +39,10 @@ class BitcoinOTCDataset(DGLBuiltinDataset):
     verbose: bool
         Whether to print out progress information.
         Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -67,12 +71,13 @@ class BitcoinOTCDataset(DGLBuiltinDataset):
     _url = 'https://snap.stanford.edu/data/soc-sign-bitcoinotc.csv.gz'
     _sha1_str = 'c14281f9e252de0bd0b5f1c6e2bae03123938641'
 
-    def __init__(self, raw_dir=None, force_reload=False, verbose=False):
+    def __init__(self, raw_dir=None, force_reload=False, verbose=False, transform=None):
         super(BitcoinOTCDataset, self).__init__(name='bitcoinotc',
                                                 url=self._url,
                                                 raw_dir=raw_dir,
                                                 force_reload=force_reload,
-                                                verbose=verbose)
+                                                verbose=verbose,
+                                                transform=transform)
 
     def download(self):
         gz_file_path = os.path.join(self.raw_dir, self.name + '.csv.gz')
@@ -143,7 +148,10 @@ class BitcoinOTCDataset(DGLBuiltinDataset):
 
             - ``edata['h']`` : edge weights
         """
+        if self._transform is None:
             return self.graphs[item]
+        else:
+            return self._transform(self.graphs[item])
 
     @property
     def is_temporal(self):

python/dgl/data/citation_graph.py

@@ -43,10 +43,14 @@ class CitationGraphDataset(DGLBuiltinDataset):
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
-    reverse_edge: bool
+    reverse_edge : bool
         Whether to add reverse edges in graph. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
     """
     _urls = {
         'cora_v2' : 'dataset/cora_v2.zip',
@@ -54,7 +58,8 @@ class CitationGraphDataset(DGLBuiltinDataset):
         'pubmed' : 'dataset/pubmed.zip',
     }
 
-    def __init__(self, name, raw_dir=None, force_reload=False, verbose=True, reverse_edge=True):
+    def __init__(self, name, raw_dir=None, force_reload=False,
+                 verbose=True, reverse_edge=True, transform=None):
         assert name.lower() in ['cora', 'citeseer', 'pubmed']
 
         # Previously we use the pre-processing in pygcn (https://github.com/tkipf/pygcn)
@@ -69,7 +74,8 @@ class CitationGraphDataset(DGLBuiltinDataset):
                                                    url=url,
                                                    raw_dir=raw_dir,
                                                    force_reload=force_reload,
-                                                   verbose=verbose)
+                                                   verbose=verbose,
+                                                   transform=transform)
 
     def process(self):
         """Loads input data from data directory and reorder graph for better locality
@@ -213,7 +219,10 @@ class CitationGraphDataset(DGLBuiltinDataset):
 
     def __getitem__(self, idx):
         assert idx == 0, "This dataset has only one graph"
+        if self._transform is None:
             return self._g
+        else:
+            return self._transform(self._g)
 
     def __len__(self):
         return 1
@@ -356,10 +365,14 @@ class CoraGraphDataset(CitationGraphDataset):
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
-    reverse_edge: bool
+    reverse_edge : bool
         Whether to add reverse edges in graph. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -400,10 +413,12 @@ class CoraGraphDataset(CitationGraphDataset):
     >>> label = g.ndata['label']
 
     """
-    def __init__(self, raw_dir=None, force_reload=False, verbose=True, reverse_edge=True):
+    def __init__(self, raw_dir=None, force_reload=False, verbose=True,
+                 reverse_edge=True, transform=None):
         name = 'cora'
 
-        super(CoraGraphDataset, self).__init__(name, raw_dir, force_reload, verbose, reverse_edge)
+        super(CoraGraphDataset, self).__init__(name, raw_dir, force_reload,
+                                               verbose, reverse_edge, transform)
 
     def __getitem__(self, idx):
         r"""Gets the graph object
@@ -496,10 +511,14 @@ class CiteseerGraphDataset(CitationGraphDataset):
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
-    reverse_edge: bool
+    reverse_edge : bool
         Whether to add reverse edges in graph. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -543,10 +562,12 @@ class CiteseerGraphDataset(CitationGraphDataset):
     >>> label = g.ndata['label']
 
     """
-    def __init__(self, raw_dir=None, force_reload=False, verbose=True, reverse_edge=True):
+    def __init__(self, raw_dir=None, force_reload=False,
+                 verbose=True, reverse_edge=True, transform=None):
         name = 'citeseer'
 
-        super(CiteseerGraphDataset, self).__init__(name, raw_dir, force_reload, verbose, reverse_edge)
+        super(CiteseerGraphDataset, self).__init__(name, raw_dir, force_reload,
+                                                   verbose, reverse_edge, transform)
 
     def __getitem__(self, idx):
         r"""Gets the graph object
@@ -639,10 +660,14 @@ class PubmedGraphDataset(CitationGraphDataset):
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
-    reverse_edge: bool
+    reverse_edge : bool
         Whether to add reverse edges in graph. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -683,10 +708,12 @@ class PubmedGraphDataset(CitationGraphDataset):
     >>> label = g.ndata['label']
 
     """
-    def __init__(self, raw_dir=None, force_reload=False, verbose=True, reverse_edge=True):
+    def __init__(self, raw_dir=None, force_reload=False, verbose=True,
+                 reverse_edge=True, transform=None):
         name = 'pubmed'
 
-        super(PubmedGraphDataset, self).__init__(name, raw_dir, force_reload, verbose, reverse_edge)
+        super(PubmedGraphDataset, self).__init__(name, raw_dir, force_reload,
+                                                 verbose, reverse_edge, transform)
 
     def __getitem__(self, idx):
         r"""Gets the graph object
@@ -714,7 +741,7 @@ class PubmedGraphDataset(CitationGraphDataset):
         r"""The number of graphs in the dataset."""
         return super(PubmedGraphDataset, self).__len__()
 
-def load_cora(raw_dir=None, force_reload=False, verbose=True, reverse_edge=True):
+def load_cora(raw_dir=None, force_reload=False, verbose=True, reverse_edge=True, transform=None):
     """Get CoraGraphDataset
 
     Parameters
@@ -724,19 +751,24 @@ def load_cora(raw_dir=None, force_reload=False, verbose=True, reverse_edge=True)
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
-    reverse_edge: bool
+    reverse_edge : bool
         Whether to add reverse edges in graph. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Return
     -------
     CoraGraphDataset
     """
-    data = CoraGraphDataset(raw_dir, force_reload, verbose, reverse_edge)
+    data = CoraGraphDataset(raw_dir, force_reload, verbose, reverse_edge, transform)
     return data
 
-def load_citeseer(raw_dir=None, force_reload=False, verbose=True, reverse_edge=True):
+def load_citeseer(raw_dir=None, force_reload=False, verbose=True,
+                  reverse_edge=True, transform=None):
     """Get CiteseerGraphDataset
 
     Parameters
@@ -746,19 +778,24 @@ def load_citeseer(raw_dir=None, force_reload=False, verbose=True, reverse_edge=T
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
-    reverse_edge: bool
+    reverse_edge : bool
         Whether to add reverse edges in graph. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Return
     -------
     CiteseerGraphDataset
     """
-    data = CiteseerGraphDataset(raw_dir, force_reload, verbose, reverse_edge)
+    data = CiteseerGraphDataset(raw_dir, force_reload, verbose, reverse_edge, transform)
     return data
 
-def load_pubmed(raw_dir=None, force_reload=False, verbose=True, reverse_edge=True):
+def load_pubmed(raw_dir=None, force_reload=False, verbose=True,
+                reverse_edge=True, transform=None):
     """Get PubmedGraphDataset
 
     Parameters
@@ -768,16 +805,20 @@ def load_pubmed(raw_dir=None, force_reload=False, verbose=True, reverse_edge=Tru
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-        verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
-    reverse_edge: bool
+    reverse_edge : bool
         Whether to add reverse edges in graph. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Return
     -------
     PubmedGraphDataset
     """
-    data = PubmedGraphDataset(raw_dir, force_reload, verbose, reverse_edge)
+    data = PubmedGraphDataset(raw_dir, force_reload, verbose, reverse_edge, transform)
     return data
 
 class CoraBinary(DGLBuiltinDataset):
@@ -798,15 +839,20 @@ class CoraBinary(DGLBuiltinDataset):
         Whether to reload the dataset. Default: False
     verbose: bool
         Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
     """
-    def __init__(self, raw_dir=None, force_reload=False, verbose=True):
+    def __init__(self, raw_dir=None, force_reload=False, verbose=True, transform=None):
         name = 'cora_binary'
         url = _get_dgl_url('dataset/cora_binary.zip')
         super(CoraBinary, self).__init__(name,
                                          url=url,
                                          raw_dir=raw_dir,
                                          force_reload=force_reload,
-                                         verbose=verbose)
+                                         verbose=verbose,
+                                         transform=transform)
 
     def process(self):
         root = self.raw_path
@@ -894,7 +940,11 @@ class CoraBinary(DGLBuiltinDataset):
         (dgl.DGLGraph, scipy.sparse.coo_matrix, int)
             The graph, scipy sparse coo_matrix and its label.
         """
-        return (self.graphs[i], self.pmpds[i], self.labels[i])
+        if self._transform is None:
+            g = self.graphs[i]
+        else:
+            g = self._transform(self.graphs[i])
+        return (g, self.pmpds[i], self.labels[i])
 
     @property
     def save_name(self):

python/dgl/data/csv_dataset.py

@@ -33,6 +33,10 @@ class DGLCSVDataset(DGLDataset):
         A callable object which is used to parse corresponding column graph
         data. Default: None. If None, a default data parser is applied
         which load data directly and tries to convert list into array.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -46,7 +50,8 @@ class DGLCSVDataset(DGLDataset):
     """
     META_YAML_NAME = 'meta.yaml'
 
-    def __init__(self, data_path, force_reload=False, verbose=True, node_data_parser=None, edge_data_parser=None, graph_data_parser=None):
+    def __init__(self, data_path, force_reload=False, verbose=True, node_data_parser=None,
+                 edge_data_parser=None, graph_data_parser=None, transform=None):
         from .csv_dataset_base import load_yaml_with_sanity_check, DefaultDataParser
         self.graphs = None
         self.data = None
@@ -61,7 +66,7 @@ class DGLCSVDataset(DGLDataset):
         self.meta_yaml = load_yaml_with_sanity_check(meta_yaml_path)
         ds_name = self.meta_yaml.dataset_name
         super().__init__(ds_name, raw_dir=os.path.dirname(
-            meta_yaml_path), force_reload=force_reload, verbose=verbose)
+            meta_yaml_path), force_reload=force_reload, verbose=verbose, transform=transform)
 
 
     def process(self):
@@ -122,10 +127,15 @@ class DGLCSVDataset(DGLDataset):
         self.graphs, self.data = load_graphs(graph_path)
 
     def __getitem__(self, i):
+        if self._transform is None:
+            g = self.graphs[i]
+        else:
+            g = self._transform(self.graphs[i])
+
         if 'label' in self.data:
-            return self.graphs[i], self.data['label'][i]
+            return g, self.data['label'][i]
         else:
-            return self.graphs[i]
+            return g
 
     def __len__(self):
         return len(self.graphs)

python/dgl/data/dgl_dataset.py

@@ -49,6 +49,10 @@ class DGLDataset(object):
         Whether to reload the dataset. Default: False
     verbose : bool
         Whether to print out progress information
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -72,13 +76,14 @@ class DGLDataset(object):
         Hash value for the dataset and the setting.
     """
     def __init__(self, name, url=None, raw_dir=None, save_dir=None,
-                 hash_key=(), force_reload=False, verbose=False):
+                 hash_key=(), force_reload=False, verbose=False, transform=None):
         self._name = name
         self._url = url
         self._force_reload = force_reload
         self._verbose = verbose
         self._hash_key = hash_key
         self._hash = self._get_hash()
+        self._transform = transform
 
         # if no dir is provided, the default dgl download dir is used.
         if raw_dir is None:
@@ -287,17 +292,23 @@ class DGLBuiltinDataset(DGLDataset):
         from the same dataset class by comparing the hash values.
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: False
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
     """
-    def __init__(self, name, url, raw_dir=None, hash_key=(), force_reload=False, verbose=False):
+    def __init__(self, name, url, raw_dir=None, hash_key=(),
+                 force_reload=False, verbose=False, transform=None):
         super(DGLBuiltinDataset, self).__init__(name,
                                                 url=url,
                                                 raw_dir=raw_dir,
                                                 save_dir=None,
                                                 hash_key=hash_key,
                                                 force_reload=force_reload,
-                                                verbose=verbose)
+                                                verbose=verbose,
+                                                transform=transform)
 
     def download(self):
         r""" Automatically download data and extract it.

python/dgl/data/fakenews.py

@@ -76,6 +76,10 @@ class FakeNewsDataset(DGLBuiltinDataset):
         downloaded data or the directory that
         already stores the input data.
         Default: ~/.dgl/
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -113,7 +117,7 @@ class FakeNewsDataset(DGLBuiltinDataset):
         'politifact': 'dataset/FakeNewsPOL.zip'
     }
 
-    def __init__(self, name, feature_name, raw_dir=None):
+    def __init__(self, name, feature_name, raw_dir=None, transform=None):
         assert name in ['gossipcop', 'politifact'], \
             "Only supports 'gossipcop' or 'politifact'."
         url = _get_dgl_url(self.file_urls[name])
@@ -123,7 +127,8 @@ class FakeNewsDataset(DGLBuiltinDataset):
         self.feature_name = feature_name
         super(FakeNewsDataset, self).__init__(name=name,
                                               url=url,
-                                              raw_dir=raw_dir)
+                                              raw_dir=raw_dir,
+                                              transform=transform)
 
     def process(self):
         """process raw data to graph, labels and masks"""
@@ -213,7 +218,11 @@ class FakeNewsDataset(DGLBuiltinDataset):
         -------
         (:class:`dgl.DGLGraph`, Tensor)
         """
-        return self.graphs[i], self.labels[i]
+        if self._transform is None:
+            g = self.graphs[i]
+        else:
+            g = self._transform(self.graphs[i])
+        return g, self.labels[i]
 
     def __len__(self):
         r"""Number of graphs in the dataset.

python/dgl/data/fraud.py

@@ -48,8 +48,12 @@ class FraudDataset(DGLBuiltinDataset):
         Default: 0.1
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -90,7 +94,7 @@ class FraudDataset(DGLBuiltinDataset):
     }
 
     def __init__(self, name, raw_dir=None, random_seed=717, train_size=0.7,
-                 val_size=0.1, force_reload=False, verbose=True):
+                 val_size=0.1, force_reload=False, verbose=True, transform=None):
         assert name in ['yelp', 'amazon'], "only supports 'yelp', or 'amazon'"
         url = _get_dgl_url(self.file_urls[name])
         self.seed = random_seed
@@ -101,7 +105,8 @@ class FraudDataset(DGLBuiltinDataset):
                                            raw_dir=raw_dir,
                                            hash_key=(random_seed, train_size, val_size),
                                            force_reload=force_reload,
-                                           verbose=verbose)
+                                           verbose=verbose,
+                                           transform=transform)
 
     def process(self):
         """process raw data to graph, labels, splitting masks"""
@@ -145,7 +150,10 @@ class FraudDataset(DGLBuiltinDataset):
             - ``ndata['test_mask']``: mask of testing set
         """
         assert idx == 0, "This dataset has only one graph"
+        if self._transform is None:
             return self.graph
+        else:
+            return self._transform(self.graph)
 
     def __len__(self):
         """number of data examples"""
@@ -254,8 +262,12 @@ class FraudYelpDataset(FraudDataset):
         Default: 0.1
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Examples
     --------
@@ -267,14 +279,15 @@ class FraudYelpDataset(FraudDataset):
     """
 
     def __init__(self, raw_dir=None, random_seed=717, train_size=0.7,
-                 val_size=0.1, force_reload=False, verbose=True):
+                 val_size=0.1, force_reload=False, verbose=True, transform=None):
         super(FraudYelpDataset, self).__init__(name='yelp',
                                                raw_dir=raw_dir,
                                                random_seed=random_seed,
                                                train_size=train_size,
                                                val_size=val_size,
                                                force_reload=force_reload,
-                                               verbose=verbose)
+                                               verbose=verbose,
+                                               transform=transform)
 
 
 class FraudAmazonDataset(FraudDataset):
@@ -330,8 +343,12 @@ class FraudAmazonDataset(FraudDataset):
         Default: 0.1
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Examples
     --------
@@ -343,11 +360,12 @@ class FraudAmazonDataset(FraudDataset):
     """
 
     def __init__(self, raw_dir=None, random_seed=717, train_size=0.7,
-                 val_size=0.1, force_reload=False, verbose=True):
+                 val_size=0.1, force_reload=False, verbose=True, transform=None):
         super(FraudAmazonDataset, self).__init__(name='amazon',
                                                  raw_dir=raw_dir,
                                                  random_seed=random_seed,
                                                  train_size=train_size,
                                                  val_size=val_size,
                                                  force_reload=force_reload,
-                                                 verbose=verbose)
+                                                 verbose=verbose,
+                                                 transform=transform)

python/dgl/data/gdelt.py

@@ -37,8 +37,12 @@ class GDELTDataset(DGLBuiltinDataset):
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -65,7 +69,8 @@ class GDELTDataset(DGLBuiltinDataset):
     ....
     >>>
     """
-    def __init__(self, mode='train', raw_dir=None, force_reload=False, verbose=False):
+    def __init__(self, mode='train', raw_dir=None,
+                 force_reload=False, verbose=False, transform=None):
         mode = mode.lower()
         assert mode in ['train', 'valid', 'test'], "Mode not valid."
         self.mode = mode
@@ -75,7 +80,8 @@ class GDELTDataset(DGLBuiltinDataset):
                                            url=_url,
                                            raw_dir=raw_dir,
                                            force_reload=force_reload,
-                                           verbose=verbose)
+                                           verbose=verbose,
+                                           transform=transform)
 
     def process(self):
         file_path = os.path.join(self.raw_path, self.mode + '.txt')
@@ -148,6 +154,8 @@ class GDELTDataset(DGLBuiltinDataset):
         rate = self.data[row_mask][:, 1]
         g = dgl_graph((edges[:, 0], edges[:, 1]))
         g.edata['rel_type'] = F.tensor(rate.reshape(-1, 1), dtype=F.data_type_dict['int64'])
+        if self._transform is not None:
+            g = self._transform(g)
         return g
 
     def __len__(self):

python/dgl/data/gindt.py

@@ -44,6 +44,10 @@ class GINDataset(DGLBuiltinDataset):
         add self to self edge if true
     degree_as_nlabel: bool
         take node degree as label and feature if true
+    transform: callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Examples
     --------
@@ -73,7 +77,7 @@ class GINDataset(DGLBuiltinDataset):
     """
 
     def __init__(self, name, self_loop, degree_as_nlabel=False,
-                 raw_dir=None, force_reload=False, verbose=False):
+                 raw_dir=None, force_reload=False, verbose=False, transform=None):
 
         self._name = name  # MUTAG
         gin_url = 'https://raw.githubusercontent.com/weihua916/powerful-gnns/master/dataset.zip'
@@ -106,7 +110,8 @@ class GINDataset(DGLBuiltinDataset):
         self.nlabels_flag = False
 
         super(GINDataset, self).__init__(name=name, url=gin_url, hash_key=(name, self_loop, degree_as_nlabel),
-                                         raw_dir=raw_dir, force_reload=force_reload, verbose=verbose)
+                                         raw_dir=raw_dir, force_reload=force_reload,
+                                         verbose=verbose, transform=transform)
 
     @property
     def raw_path(self):
@@ -136,7 +141,11 @@ class GINDataset(DGLBuiltinDataset):
         (:class:`dgl.Graph`, Tensor)
             The graph and its label.
         """
-        return self.graphs[idx], self.labels[idx]
+        if self._transform is None:
+            g = self.graphs[idx]
+        else:
+            g = self._transform(self.graphs[idx])
+        return g, self.labels[idx]
 
     def _file_path(self):
         return os.path.join(self.raw_dir, "GINDataset", 'dataset', self.name, "{}.txt".format(self.name))

python/dgl/data/gnn_benchmark.py

@@ -27,13 +27,14 @@ class GNNBenchmarkDataset(DGLBuiltinDataset):
 
     Reference: https://github.com/shchur/gnn-benchmark#datasets
     """
-    def __init__(self, name, raw_dir=None, force_reload=False, verbose=False):
+    def __init__(self, name, raw_dir=None, force_reload=False, verbose=False, transform=None):
         _url = _get_dgl_url('dataset/' + name + '.zip')
         super(GNNBenchmarkDataset, self).__init__(name=name,
                                                   url=_url,
                                                   raw_dir=raw_dir,
                                                   force_reload=force_reload,
-                                                  verbose=verbose)
+                                                  verbose=verbose,
+                                                  transform=transform)
 
     def process(self):
         npz_path = os.path.join(self.raw_path, self.name + '.npz')
@@ -128,7 +129,10 @@ class GNNBenchmarkDataset(DGLBuiltinDataset):
             - ``ndata['label']``: node labels
         """
         assert idx == 0, "This dataset has only one graph"
+        if self._transform is None:
             return self._graph
+        else:
+            return self._transform(self._graph)
 
     def __len__(self):
         r"""Number of graphs in the dataset"""
@@ -164,8 +168,12 @@ class CoraFullDataset(GNNBenchmarkDataset):
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -182,11 +190,12 @@ class CoraFullDataset(GNNBenchmarkDataset):
     >>> feat = g.ndata['feat']  # get node feature
     >>> label = g.ndata['label']  # get node labels
     """
-    def __init__(self, raw_dir=None, force_reload=False, verbose=False):
+    def __init__(self, raw_dir=None, force_reload=False, verbose=False, transform=None):
         super(CoraFullDataset, self).__init__(name="cora_full",
                                               raw_dir=raw_dir,
                                               force_reload=force_reload,
-                                              verbose=verbose)
+                                              verbose=verbose,
+                                              transform=transform)
 
     @property
     def num_classes(self):
@@ -231,8 +240,12 @@ class CoauthorCSDataset(GNNBenchmarkDataset):
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -249,11 +262,12 @@ class CoauthorCSDataset(GNNBenchmarkDataset):
     >>> feat = g.ndata['feat']  # get node feature
     >>> label = g.ndata['label']  # get node labels
     """
-    def __init__(self, raw_dir=None, force_reload=False, verbose=False):
+    def __init__(self, raw_dir=None, force_reload=False, verbose=False, transform=None):
         super(CoauthorCSDataset, self).__init__(name='coauthor_cs',
                                                 raw_dir=raw_dir,
                                                 force_reload=force_reload,
-                                                verbose=verbose)
+                                                verbose=verbose,
+                                                transform=transform)
 
     @property
     def num_classes(self):
@@ -298,8 +312,12 @@ class CoauthorPhysicsDataset(GNNBenchmarkDataset):
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -316,11 +334,12 @@ class CoauthorPhysicsDataset(GNNBenchmarkDataset):
     >>> feat = g.ndata['feat']  # get node feature
     >>> label = g.ndata['label']  # get node labels
     """
-    def __init__(self, raw_dir=None, force_reload=False, verbose=False):
+    def __init__(self, raw_dir=None, force_reload=False, verbose=False, transform=None):
         super(CoauthorPhysicsDataset, self).__init__(name='coauthor_physics',
                                                      raw_dir=raw_dir,
                                                      force_reload=force_reload,
-                                                     verbose=verbose)
+                                                     verbose=verbose,
+                                                     transform=transform)
 
     @property
     def num_classes(self):
@@ -364,8 +383,12 @@ class AmazonCoBuyComputerDataset(GNNBenchmarkDataset):
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -382,11 +405,12 @@ class AmazonCoBuyComputerDataset(GNNBenchmarkDataset):
     >>> feat = g.ndata['feat']  # get node feature
     >>> label = g.ndata['label']  # get node labels
     """
-    def __init__(self, raw_dir=None, force_reload=False, verbose=False):
+    def __init__(self, raw_dir=None, force_reload=False, verbose=False, transform=None):
         super(AmazonCoBuyComputerDataset, self).__init__(name='amazon_co_buy_computer',
                                                          raw_dir=raw_dir,
                                                          force_reload=force_reload,
-                                                         verbose=verbose)
+                                                         verbose=verbose,
+                                                         transform=transform)
 
     @property
     def num_classes(self):
@@ -430,8 +454,12 @@ class AmazonCoBuyPhotoDataset(GNNBenchmarkDataset):
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -448,11 +476,12 @@ class AmazonCoBuyPhotoDataset(GNNBenchmarkDataset):
     >>> feat = g.ndata['feat']  # get node feature
     >>> label = g.ndata['label']  # get node labels
     """
-    def __init__(self, raw_dir=None, force_reload=False, verbose=False):
+    def __init__(self, raw_dir=None, force_reload=False, verbose=False, transform=None):
         super(AmazonCoBuyPhotoDataset, self).__init__(name='amazon_co_buy_photo',
                                                       raw_dir=raw_dir,
                                                       force_reload=force_reload,
-                                                      verbose=verbose)
+                                                      verbose=verbose,
+                                                      transform=transform)
 
     @property
     def num_classes(self):

python/dgl/data/icews18.py

@@ -39,8 +39,12 @@ class ICEWS18Dataset(DGLBuiltinDataset):
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     -------
@@ -61,7 +65,7 @@ class ICEWS18Dataset(DGLBuiltinDataset):
     ....
     >>>
     """
-    def __init__(self, mode='train', raw_dir=None, force_reload=False, verbose=False):
+    def __init__(self, mode='train', raw_dir=None, force_reload=False, verbose=False, transform=None):
         mode = mode.lower()
         assert mode in ['train', 'valid', 'test'], "Mode not valid"
         self.mode = mode
@@ -70,7 +74,8 @@ class ICEWS18Dataset(DGLBuiltinDataset):
                                              url=_url,
                                              raw_dir=raw_dir,
                                              force_reload=force_reload,
-                                             verbose=verbose)
+                                             verbose=verbose,
+                                             transform=transform)
 
     def process(self):
         data = loadtxt(os.path.join(self.save_path, '{}.txt'.format(self.mode)),
@@ -118,7 +123,10 @@ class ICEWS18Dataset(DGLBuiltinDataset):
 
             - ``edata['rel_type']``: edge type
         """
+        if self._transform is None:
             return self._graphs[idx]
+        else:
+            return self._transform(self._graphs[idx])
 
     def __len__(self):
         r"""Number of graphs in the dataset.

python/dgl/data/karate.py

@@ -34,6 +34,13 @@ class KarateClubDataset(DGLDataset):
     - Edges: 156
     - Number of Classes: 2
 
+    Parameters
+    ----------
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
     Attributes
     ----------
     num_classes : int
@@ -48,8 +55,8 @@ class KarateClubDataset(DGLDataset):
     >>> g = dataset[0]
     >>> labels = g.ndata['label']
     """
-    def __init__(self):
-        super(KarateClubDataset, self).__init__(name='karate_club')
+    def __init__(self, transform=None):
+        super(KarateClubDataset, self).__init__(name='karate_club', transform=transform)
 
     def process(self):
         kc_graph = nx.karate_club_graph()
@@ -88,7 +95,10 @@ class KarateClubDataset(DGLDataset):
             - ``ndata['label']``: ground truth labels
         """
         assert idx == 0, "This dataset has only one graph"
+        if self._transform is None:
             return self._graph
+        else:
+            return self._transform(self._graph)
 
     def __len__(self):
         r"""The number of graphs in the dataset."""

python/dgl/data/knowledge_graph.py

@@ -25,19 +25,24 @@ class KnowledgeGraphDataset(DGLBuiltinDataset):
 
     Parameters
     -----------
-    name: str
+    name : str
         Name can be 'FB15k-237', 'FB15k' or 'wn18'.
-    reverse: bool
+    reverse : bool
         Whether add reverse edges. Default: True.
     raw_dir : str
         Raw file directory to download/contains the input data directory.
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
     """
-    def __init__(self, name, reverse=True, raw_dir=None, force_reload=False, verbose=True):
+    def __init__(self, name, reverse=True, raw_dir=None, force_reload=False,
+                 verbose=True, transform=None):
         self._name = name
         self.reverse = reverse
         url = _get_dgl_url('dataset/') + '{}.tgz'.format(name)
@@ -45,7 +50,8 @@ class KnowledgeGraphDataset(DGLBuiltinDataset):
                                                     url=url,
                                                     raw_dir=raw_dir,
                                                     force_reload=force_reload,
-                                                    verbose=verbose)
+                                                    verbose=verbose,
+                                                    transform=transform)
 
     def download(self):
         r""" Automatically download data and extract it.
@@ -112,7 +118,10 @@ class KnowledgeGraphDataset(DGLBuiltinDataset):
 
     def __getitem__(self, idx):
         assert idx == 0, "This dataset has only one graph"
+        if self._transform is None:
             return self._g
+        else:
+            return self._transform(self._g)
 
     def __len__(self):
         return 1
@@ -389,8 +398,12 @@ class FB15k237Dataset(KnowledgeGraphDataset):
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -433,9 +446,11 @@ class FB15k237Dataset(KnowledgeGraphDataset):
     >>>
     >>> # Train, Validation and Test
     """
-    def __init__(self, reverse=True, raw_dir=None, force_reload=False, verbose=True):
+    def __init__(self, reverse=True, raw_dir=None, force_reload=False,
+                 verbose=True, transform=None):
         name = 'FB15k-237'
-        super(FB15k237Dataset, self).__init__(name, reverse, raw_dir, force_reload, verbose)
+        super(FB15k237Dataset, self).__init__(name, reverse, raw_dir,
+                                              force_reload, verbose, transform)
 
     def __getitem__(self, idx):
         r"""Gets the graph object
@@ -526,8 +541,12 @@ class FB15kDataset(KnowledgeGraphDataset):
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -570,9 +589,11 @@ class FB15kDataset(KnowledgeGraphDataset):
     >>> # Train, Validation and Test
     >>>
     """
-    def __init__(self, reverse=True, raw_dir=None, force_reload=False, verbose=True):
+    def __init__(self, reverse=True, raw_dir=None, force_reload=False,
+                 verbose=True, transform=None):
         name = 'FB15k'
-        super(FB15kDataset, self).__init__(name, reverse, raw_dir, force_reload, verbose)
+        super(FB15kDataset, self).__init__(name, reverse, raw_dir,
+                                           force_reload, verbose, transform)
 
     def __getitem__(self, idx):
         r"""Gets the graph object
@@ -662,8 +683,12 @@ class WN18Dataset(KnowledgeGraphDataset):
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -706,9 +731,11 @@ class WN18Dataset(KnowledgeGraphDataset):
     >>> # Train, Validation and Test
     >>>
     """
-    def __init__(self, reverse=True, raw_dir=None, force_reload=False, verbose=True):
+    def __init__(self, reverse=True, raw_dir=None, force_reload=False,
+                 verbose=True, transform=None):
         name = 'wn18'
-        super(WN18Dataset, self).__init__(name, reverse, raw_dir, force_reload, verbose)
+        super(WN18Dataset, self).__init__(name, reverse, raw_dir,
+                                          force_reload, verbose, transform)
 
     def __getitem__(self, idx):
         r"""Gets the graph object

python/dgl/data/minigc.py

@@ -33,8 +33,12 @@ class MiniGCDataset(DGLDataset):
         Minimum number of nodes for graphs
     max_num_v: int
         Maximum number of nodes for graphs
-    seed : int, default is 0
+    seed: int, default is 0
         Random seed for data generation
+    transform: callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -75,7 +79,7 @@ class MiniGCDataset(DGLDataset):
     """
 
     def __init__(self, num_graphs, min_num_v, max_num_v, seed=0,
-                 save_graph=True, force_reload=False, verbose=False):
+                 save_graph=True, force_reload=False, verbose=False, transform=None):
         self.num_graphs = num_graphs
         self.min_num_v = min_num_v
         self.max_num_v = max_num_v
@@ -84,7 +88,7 @@ class MiniGCDataset(DGLDataset):
 
         super(MiniGCDataset, self).__init__(name="minigc", hash_key=(num_graphs, min_num_v, max_num_v, seed),
                                             force_reload=force_reload,
-                                            verbose=verbose)
+                                            verbose=verbose, transform=transform)
 
     def process(self):
         self.graphs = []
@@ -108,7 +112,11 @@ class MiniGCDataset(DGLDataset):
         (:class:`dgl.Graph`, Tensor)
             The graph and its label.
         """
-        return self.graphs[idx], self.labels[idx]
+        if self._transform is None:
+            g = self.graphs[idx]
+        else:
+            g = self._transform(self.graphs[idx])
+        return g, self.labels[idx]
 
     def has_cache(self):
         graph_path = os.path.join(self.save_path, 'dgl_graph_{}.bin'.format(self.hash))

python/dgl/data/ppi.py

@@ -56,9 +56,13 @@ class PPIDataset(DGLBuiltinDataset):
     force_reload : bool
         Whether to reload the dataset.
         Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information.
         Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -79,7 +83,8 @@ class PPIDataset(DGLBuiltinDataset):
     ....    # your code here
     >>>
     """
-    def __init__(self, mode='train', raw_dir=None, force_reload=False, verbose=False):
+    def __init__(self, mode='train', raw_dir=None, force_reload=False,
+                 verbose=False, transform=None):
         assert mode in ['train', 'valid', 'test']
         self.mode = mode
         _url = _get_dgl_url('dataset/ppi.zip')
@@ -87,7 +92,8 @@ class PPIDataset(DGLBuiltinDataset):
                                          url=_url,
                                          raw_dir=raw_dir,
                                          force_reload=force_reload,
-                                         verbose=verbose)
+                                         verbose=verbose,
+                                         transform=transform)
 
     def process(self):
         graph_file = os.path.join(self.save_path, '{}_graph.json'.format(self.mode))
@@ -178,7 +184,10 @@ class PPIDataset(DGLBuiltinDataset):
             - ``ndata['feat']``: node features
             - ``ndata['label']``: node labels
         """
+        if self._transform is None:
             return self.graphs[item]
+        else:
+            return self._transform(self.graphs[item])
 
 
 class LegacyPPIDataset(PPIDataset):
@@ -198,5 +207,8 @@ class LegacyPPIDataset(PPIDataset):
         (dgl.DGLGraph, Tensor, Tensor)
             The graph, features and its label.
         """
-
-        return self.graphs[item], self.graphs[item].ndata['feat'], self.graphs[item].ndata['label']
+        if self._transform is None:
+            g = self.graphs[item]
+        else:
+            g = self._transform(self.graphs[item])
+        return g, g.ndata['feat'], g.ndata['label']

python/dgl/data/qm7b.py

@@ -34,8 +34,12 @@ class QM7bDataset(DGLDataset):
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -65,12 +69,13 @@ class QM7bDataset(DGLDataset):
            'datasets/qm7b.mat'
     _sha1_str = '4102c744bb9d6fd7b40ac67a300e49cd87e28392'
 
-    def __init__(self, raw_dir=None, force_reload=False, verbose=False):
+    def __init__(self, raw_dir=None, force_reload=False, verbose=False, transform=None):
         super(QM7bDataset, self).__init__(name='qm7b',
                                           url=self._url,
                                           raw_dir=raw_dir,
                                           force_reload=force_reload,
-                                          verbose=verbose)
+                                          verbose=verbose,
+                                          transform=transform)
 
     def process(self):
         mat_path = self.raw_path + '.mat'
@@ -129,7 +134,11 @@ class QM7bDataset(DGLDataset):
         -------
         (:class:`dgl.DGLGraph`, Tensor)
         """
-        return self.graphs[idx], self.label[idx]
+        if self._transform is None:
+            g = self.graphs[idx]
+        else:
+            g = self._transform(self.graphs[idx])
+        return g, self.label[idx]
 
     def __len__(self):
         r"""Number of graphs in the dataset.

python/dgl/data/qm9.py

@@ -60,9 +60,9 @@ class QM9Dataset(DGLDataset):
 
     Parameters
     ----------
-    label_keys: list
+    label_keys : list
         Names of the regression property, which should be a subset of the keys in the table above.
-    cutoff: float
+    cutoff : float
         Cutoff distance for interatomic interactions, i.e. two atoms are connected in the corresponding graph if the distance between them is no larger than this.
         Default: 5.0 Angstrom
     raw_dir : str
@@ -70,8 +70,12 @@ class QM9Dataset(DGLDataset):
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -102,7 +106,8 @@ class QM9Dataset(DGLDataset):
                  cutoff=5.0,
                  raw_dir=None,
                  force_reload=False,
-                 verbose=False):
+                 verbose=False,
+                 transform=None):
 
         self.cutoff = cutoff
         self.label_keys = label_keys
@@ -112,7 +117,8 @@ class QM9Dataset(DGLDataset):
                                          url=self._url,
                                          raw_dir=raw_dir,
                                          force_reload=force_reload,
-                                         verbose=verbose)
+                                         verbose=verbose,
+                                         transform=transform)
 
     def process(self):
         npz_path = f'{self.raw_dir}/qm9_eV.npz'
@@ -172,6 +178,10 @@ class QM9Dataset(DGLDataset):
         g.ndata['R'] = F.tensor(R, dtype=F.data_type_dict['float32'])
         g.ndata['Z'] = F.tensor(self.Z[self.N_cumsum[idx]:self.N_cumsum[idx + 1]],
                                 dtype=F.data_type_dict['int64'])
+
+        if self._transform is not None:
+            g = self._transform(g)
+
         return g, label
 
     def __len__(self):

python/dgl/data/qm9_edge.py

@@ -88,7 +88,7 @@ class QM9EdgeDataset(DGLDataset):
 
     Parameters
     ----------
-    label_keys: list
+    label_keys : list
         Names of the regression property, which should be a subset of the keys in the table above.
         If not provided, it will load all the labels.
     raw_dir : str
@@ -96,8 +96,12 @@ class QM9EdgeDataset(DGLDataset):
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False.
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -134,7 +138,8 @@ class QM9EdgeDataset(DGLDataset):
                  label_keys=None,
                  raw_dir=None,
                  force_reload=False,
-                 verbose=True):
+                 verbose=True,
+                 transform=None):
         if label_keys is None:
             self.label_keys = None
             self.num_labels = 19
@@ -149,7 +154,8 @@ class QM9EdgeDataset(DGLDataset):
                                              raw_dir=raw_dir,
                                              url=self._url,
                                              force_reload=force_reload,
-                                             verbose=verbose)
+                                             verbose=verbose,
+                                             transform=transform)
 
     def download(self):
         file_path = f'{self.raw_dir}/qm9_edge.npz'
@@ -225,6 +231,9 @@ class QM9EdgeDataset(DGLDataset):
 
         label = F.tensor(self.targets[idx][self.label_keys], dtype=F.data_type_dict['float32'])
 
+        if self._transform is not None:
+            g = self._transform(g)
+
         return g, label
 
     def __len__(self):

python/dgl/data/rdf.py

@@ -94,15 +94,20 @@ class RDFGraphDataset(DGLBuiltinDataset):
         Default: ~/.dgl/
     force_reload : bool, optional
         If true, force load and process from raw data. Ignore cached pre-processed data.
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
     """
     def __init__(self, name, url, predict_category,
                  print_every=10000,
                  insert_reverse=True,
                  raw_dir=None,
                  force_reload=False,
-                 verbose=True):
+                 verbose=True,
+                 transform=None):
         self._insert_reverse = insert_reverse
         self._print_every = print_every
         self._predict_category = predict_category
@@ -110,7 +115,8 @@ class RDFGraphDataset(DGLBuiltinDataset):
         super(RDFGraphDataset, self).__init__(name, url,
                                               raw_dir=raw_dir,
                                               force_reload=force_reload,
-                                              verbose=verbose)
+                                              verbose=verbose,
+                                              transform=transform)
 
     def process(self):
         raw_tuples = self.load_raw_tuples(self.raw_path)
@@ -409,6 +415,8 @@ class RDFGraphDataset(DGLBuiltinDataset):
         r"""Gets the graph object
         """
         g = self._hg
+        if self._transform is not None:
+            g = self._transform(g)
         return g
 
     def __len__(self):
@@ -523,17 +531,21 @@ class AIFBDataset(RDFGraphDataset):
 
     Parameters
     -----------
-    print_every: int
+    print_every : int
         Preprocessing log for every X tuples. Default: 10000.
-    insert_reverse: bool
+    insert_reverse : bool
         If true, add reverse edge and reverse relations to the final graph. Default: True.
     raw_dir : str
         Raw file directory to download/contains the input data directory.
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -562,7 +574,8 @@ class AIFBDataset(RDFGraphDataset):
                  insert_reverse=True,
                  raw_dir=None,
                  force_reload=False,
-                 verbose=True):
+                 verbose=True,
+                 transform=None):
         import rdflib as rdf
         self.employs = rdf.term.URIRef("http://swrc.ontoware.org/ontology#employs")
         self.affiliation = rdf.term.URIRef("http://swrc.ontoware.org/ontology#affiliation")
@@ -574,7 +587,8 @@ class AIFBDataset(RDFGraphDataset):
                                           insert_reverse=insert_reverse,
                                           raw_dir=raw_dir,
                                           force_reload=force_reload,
-                                          verbose=verbose)
+                                          verbose=verbose,
+                                          transform=transform)
 
     def __getitem__(self, idx):
         r"""Gets the graph object
@@ -653,17 +667,21 @@ class MUTAGDataset(RDFGraphDataset):
 
     Parameters
     -----------
-    print_every: int
+    print_every : int
         Preprocessing log for every X tuples. Default: 10000.
-    insert_reverse: bool
+    insert_reverse : bool
         If true, add reverse edge and reverse relations to the final graph. Default: True.
     raw_dir : str
         Raw file directory to download/contains the input data directory.
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -697,7 +715,8 @@ class MUTAGDataset(RDFGraphDataset):
                  insert_reverse=True,
                  raw_dir=None,
                  force_reload=False,
-                 verbose=True):
+                 verbose=True,
+                 transform=None):
         import rdflib as rdf
         self.is_mutagenic = rdf.term.URIRef("http://dl-learner.org/carcinogenesis#isMutagenic")
         self.rdf_type = rdf.term.URIRef("http://www.w3.org/1999/02/22-rdf-syntax-ns#type")
@@ -712,7 +731,8 @@ class MUTAGDataset(RDFGraphDataset):
                                            insert_reverse=insert_reverse,
                                            raw_dir=raw_dir,
                                            force_reload=force_reload,
-                                           verbose=verbose)
+                                           verbose=verbose,
+                                           transform=transform)
 
     def __getitem__(self, idx):
         r"""Gets the graph object
@@ -814,17 +834,21 @@ class BGSDataset(RDFGraphDataset):
 
     Parameters
     -----------
-    print_every: int
+    print_every : int
         Preprocessing log for every X tuples. Default: 10000.
-    insert_reverse: bool
+    insert_reverse : bool
         If true, add reverse edge and reverse relations to the final graph. Default: True.
     raw_dir : str
         Raw file directory to download/contains the input data directory.
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -854,7 +878,8 @@ class BGSDataset(RDFGraphDataset):
                  insert_reverse=True,
                  raw_dir=None,
                  force_reload=False,
-                 verbose=True):
+                 verbose=True,
+                 transform=None):
         import rdflib as rdf
         url = _get_dgl_url('dataset/rdf/bgs-hetero.zip')
         name = 'bgs-hetero'
@@ -865,7 +890,8 @@ class BGSDataset(RDFGraphDataset):
                                          insert_reverse=insert_reverse,
                                          raw_dir=raw_dir,
                                          force_reload=force_reload,
-                                         verbose=verbose)
+                                         verbose=verbose,
+                                         transform=transform)
 
     def __getitem__(self, idx):
         r"""Gets the graph object
@@ -964,17 +990,21 @@ class AMDataset(RDFGraphDataset):
 
     Parameters
     -----------
-    print_every: int
+    print_every : int
         Preprocessing log for every X tuples. Default: 10000.
-    insert_reverse: bool
+    insert_reverse : bool
         If true, add reverse edge and reverse relations to the final graph. Default: True.
     raw_dir : str
         Raw file directory to download/contains the input data directory.
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -1003,7 +1033,8 @@ class AMDataset(RDFGraphDataset):
                  insert_reverse=True,
                  raw_dir=None,
                  force_reload=False,
-                 verbose=True):
+                 verbose=True,
+                 transform=None):
         import rdflib as rdf
         self.objectCategory = rdf.term.URIRef("http://purl.org/collections/nl/am/objectCategory")
         self.material = rdf.term.URIRef("http://purl.org/collections/nl/am/material")
@@ -1015,7 +1046,8 @@ class AMDataset(RDFGraphDataset):
                                         insert_reverse=insert_reverse,
                                         raw_dir=raw_dir,
                                         force_reload=force_reload,
-                                        verbose=verbose)
+                                        verbose=verbose,
+                                        transform=transform)
 
     def __getitem__(self, idx):
         r"""Gets the graph object

python/dgl/data/reddit.py

@@ -84,8 +84,12 @@ class RedditDataset(DGLBuiltinDataset):
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -125,7 +129,8 @@ class RedditDataset(DGLBuiltinDataset):
     >>>
     >>> # Train, Validation and Test
     """
-    def __init__(self, self_loop=False, raw_dir=None, force_reload=False, verbose=False):
+    def __init__(self, self_loop=False, raw_dir=None, force_reload=False,
+                 verbose=False, transform=None):
         self_loop_str = ""
         if self_loop:
             self_loop_str = "_self_loop"
@@ -135,7 +140,8 @@ class RedditDataset(DGLBuiltinDataset):
                                             url=_url,
                                             raw_dir=raw_dir,
                                             force_reload=force_reload,
-                                            verbose=verbose)
+                                            verbose=verbose,
+                                            transform=transform)
 
     def process(self):
         # graph
@@ -251,7 +257,10 @@ class RedditDataset(DGLBuiltinDataset):
             - ``ndata['test_mask']:`` mask for test node set
         """
         assert idx == 0, "Reddit Dataset only has one graph"
+        if self._transform is None:
             return self._graph
+        else:
+            return self._transform(self._graph)
 
     def __len__(self):
         r"""Number of graphs in the dataset"""

python/dgl/data/tree.py

@@ -63,8 +63,12 @@ class SSTDataset(DGLBuiltinDataset):
         Default: ~/.dgl/
     force_reload : bool
         Whether to reload the dataset. Default: False
-    verbose: bool
+    verbose : bool
         Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
 
     Attributes
     ----------
@@ -120,7 +124,8 @@ class SSTDataset(DGLBuiltinDataset):
                  vocab_file=None,
                  raw_dir=None,
                  force_reload=False,
-                 verbose=False):
+                 verbose=False,
+                 transform=None):
         assert mode in ['train', 'dev', 'test', 'tiny']
         _url = _get_dgl_url('dataset/sst.zip')
         self._glove_embed_file = glove_embed_file if mode == 'train' else None
@@ -130,7 +135,8 @@ class SSTDataset(DGLBuiltinDataset):
                                          url=_url,
                                          raw_dir=raw_dir,
                                          force_reload=force_reload,
-                                         verbose=verbose)
+                                         verbose=verbose,
+                                         transform=transform)
 
     def process(self):
         from nltk.corpus.reader import BracketParseCorpusReader
@@ -255,7 +261,10 @@ class SSTDataset(DGLBuiltinDataset):
             - ``ndata['y']:`` label of the node
             - ``ndata['mask']``: 1 if the node is a leaf, otherwise 0
         """
+        if self._transform is None:
             return self._trees[idx]
+        else:
+            return self._transform(self._trees[idx])
 
     def __len__(self):
         r"""Number of graphs in the dataset."""