[Dataset] Citation graph (#1902)

* citation graph

* GCN example use new citatoin dataset

* mxnet gat

* triger

* Fix

* Fix gat

* fix

* Fix tensorflow dgi

* Fix appnp, graphsage for mxnet

* fix monet and sgc for mxnet

* Fix tagcn

* update sgc, appnp
Co-authored-by: Ubuntu <ubuntu@ip-172-31-51-214.ec2.internal>

examples/mxnet/appnp/appnp.py

@@ -4,8 +4,9 @@ import dgl
 import mxnet as mx
 from mxnet import nd, gluon
 from mxnet.gluon import nn
-from dgl import DGLGraph
-from dgl.data import register_data_args, load_data
+import dgl
+from dgl.data import register_data_args
+from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
 from dgl.nn.mxnet.conv import APPNPConv
 
 class APPNP(nn.Block):
@@ -57,13 +58,29 @@ def evaluate(model, features, labels, mask):
 
 def main(args):
     # load and preprocess dataset
-    data = load_data(args)
-    features = nd.array(data.features)
-    labels = nd.array(data.labels)
-    train_mask = nd.array(data.train_mask)
-    val_mask = nd.array(data.val_mask)
-    test_mask = nd.array(data.test_mask)
+    if args.dataset == 'cora':
+        data = CoraGraphDataset()
+    elif args.dataset == 'citeseer':
+        data = CiteseerGraphDataset()
+    elif args.dataset == 'pubmed':
+        data = PubmedGraphDataset()
+    else:
+        raise ValueError('Unknown dataset: {}'.format(args.dataset))
+
+    g = data[0]
+    if args.gpu < 0:
+        cuda = False
+        ctx = mx.cpu(0)
+    else:
+        cuda = True
+        ctx = mx.gpu(args.gpu)
+        g = g.to(ctx)
 
+    features = g.ndata['feat']
+    labels = mx.nd.array(g.ndata['label'], dtype="float32", ctx=ctx)
+    train_mask = g.ndata['train_mask']
+    val_mask = g.ndata['val_mask']
+    test_mask = g.ndata['test_mask']
     in_feats = features.shape[1]
     n_classes = data.num_labels
     n_edges = data.graph.number_of_edges()
@@ -78,24 +95,9 @@ def main(args):
            val_mask.sum().asscalar(),
            test_mask.sum().asscalar()))
 
-    if args.gpu < 0:
-        ctx = mx.cpu()
-    else:
-        ctx = mx.gpu(args.gpu)
-
-    features = features.as_in_context(ctx)
-    labels = labels.as_in_context(ctx)
-    train_mask = train_mask.as_in_context(ctx)
-    val_mask = val_mask.as_in_context(ctx)
-    test_mask = test_mask.as_in_context(ctx)
-
-    # graph preprocess and calculate normalization factor
-    g = DGLGraph(data.graph)
-    n_edges = g.number_of_edges()
     # add self loop
-    g.add_edges(g.nodes(), g.nodes())
-    g.set_n_initializer(dgl.init.zero_initializer)
-    g.set_e_initializer(dgl.init.zero_initializer)
+    g = dgl.remove_self_loop(g)
+    g = dgl.add_self_loop(g)
 
     # create APPNP model
     model = APPNP(g,

examples/mxnet/gat/train.py

@@ -14,8 +14,10 @@ import time
 import mxnet as mx
 from mxnet import gluon
 import numpy as np
-from dgl import DGLGraph
-from dgl.data import register_data_args, load_data
+
+import dgl
+from dgl.data import register_data_args
+from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
 from gat import GAT
 from utils import EarlyStopping
 
@@ -34,33 +36,38 @@ def evaluate(model, features, labels, mask):
 
 def main(args):
     # load and preprocess dataset
-    data = load_data(args)
-
-    features = mx.nd.array(data.features)
-    labels = mx.nd.array(data.labels)
-    mask = mx.nd.array(np.where(data.train_mask == 1))
-    test_mask = mx.nd.array(np.where(data.test_mask == 1))
-    val_mask = mx.nd.array(np.where(data.val_mask == 1))
-    in_feats = features.shape[1]
-    n_classes = data.num_labels
-    n_edges = data.graph.number_of_edges()
+    if args.dataset == 'cora':
+        data = CoraGraphDataset()
+    elif args.dataset == 'citeseer':
+        data = CiteseerGraphDataset()
+    elif args.dataset == 'pubmed':
+        data = PubmedGraphDataset()
+    else:
+        raise ValueError('Unknown dataset: {}'.format(args.dataset))
 
+    g = data[0]
     if args.gpu < 0:
-        ctx = mx.cpu()
+        cuda = False
+        ctx = mx.cpu(0)
     else:
+        cuda = True
         ctx = mx.gpu(args.gpu)
-        features = features.as_in_context(ctx)
-        labels = labels.as_in_context(ctx)
-        mask = mask.as_in_context(ctx)
-        test_mask = test_mask.as_in_context(ctx)
-        val_mask = val_mask.as_in_context(ctx)
-    # create graph
-    g = data.graph
-    # add self-loop
-    g.remove_edges_from(nx.selfloop_edges(g))
-    g = DGLGraph(g)
-    g.add_edges(g.nodes(), g.nodes())
         g = g.to(ctx)
+
+    features = g.ndata['feat']
+    labels = mx.nd.array(g.ndata['label'], dtype="float32", ctx=ctx)
+    mask = g.ndata['train_mask']
+    mask = mx.nd.array(np.nonzero(mask.asnumpy())[0], ctx=ctx)
+    val_mask = g.ndata['val_mask']
+    val_mask =  mx.nd.array(np.nonzero(val_mask.asnumpy())[0], ctx=ctx)
+    test_mask = g.ndata['test_mask']
+    test_mask =  mx.nd.array(np.nonzero(test_mask.asnumpy())[0], ctx=ctx)
+    in_feats = features.shape[1]
+    n_classes = data.num_labels
+    n_edges = data.graph.number_of_edges()
+
+    g = dgl.remove_self_loop(g)
+    g = dgl.add_self_loop(g)
     # create model
     heads = ([args.num_heads] * args.num_layers) + [args.num_out_heads]
     model = GAT(g,

examples/mxnet/gcn/gcn_concat.py

@@ -11,8 +11,8 @@ import mxnet as mx
 from mxnet import gluon
 import dgl
 import dgl.function as fn
-from dgl import DGLGraph
-from dgl.data import register_data_args, load_data
+from dgl.data import register_data_args
+from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
 
 
 class GCNLayer(gluon.Block):
@@ -75,16 +75,29 @@ def evaluate(model, features, labels, mask):
 
 def main(args):
     # load and preprocess dataset
-    data = load_data(args)
+    if args.dataset == 'cora':
+        data = CoraGraphDataset()
+    elif args.dataset == 'citeseer':
+        data = CiteseerGraphDataset()
+    elif args.dataset == 'pubmed':
+        data = PubmedGraphDataset()
+    else:
+        raise ValueError('Unknown dataset: {}'.format(args.dataset))
 
-    if args.self_loop:
-        data.graph.add_edges_from([(i,i) for i in range(len(data.graph))])
+    g = data[0]
+    if args.gpu < 0:
+        cuda = False
+        ctx = mx.cpu(0)
+    else:
+        cuda = True
+        ctx = mx.gpu(args.gpu)
+        g = g.to(ctx)
 
-    features = mx.nd.array(data.features)
-    labels = mx.nd.array(data.labels)
-    train_mask = mx.nd.array(data.train_mask)
-    val_mask = mx.nd.array(data.val_mask)
-    test_mask = mx.nd.array(data.test_mask)
+    features = g.ndata['feat']
+    labels = mx.nd.array(g.ndata['label'], dtype="float32", ctx=ctx)
+    train_mask = g.ndata['train_mask']
+    val_mask = g.ndata['val_mask']
+    test_mask = g.ndata['test_mask']
     in_feats = features.shape[1]
     n_classes = data.num_labels
     n_edges = data.graph.number_of_edges()
@@ -99,21 +112,10 @@ def main(args):
               val_mask.sum().asscalar(),
               test_mask.sum().asscalar()))
 
-    if args.gpu < 0:
-        cuda = False
-        ctx = mx.cpu(0)
-    else:
-        cuda = True
-        ctx = mx.gpu(args.gpu)
-
-    features = features.as_in_context(ctx)
-    labels = labels.as_in_context(ctx)
-    train_mask = train_mask.as_in_context(ctx)
-    val_mask = val_mask.as_in_context(ctx)
-    test_mask = test_mask.as_in_context(ctx)
-
-    # create GCN model
-    g = DGLGraph(data.graph)
+    # add self loop
+    if args.self_loop:
+        g = dgl.remove_self_loop(g)
+        g = dgl.add_self_loop(g)
     # normalization
     in_degs = g.in_degrees().astype('float32')
     out_degs = g.out_degrees().astype('float32')

examples/mxnet/gcn/train.py

@@ -6,7 +6,8 @@ import mxnet as mx
 from mxnet import gluon
 
 import dgl
-from dgl.data import register_data_args, load_data
+from dgl.data import register_data_args
+from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
 
 from gcn import GCN
 #from gcn_mp import GCN
@@ -19,13 +20,29 @@ def evaluate(model, features, labels, mask):
 
 def main(args):
     # load and preprocess dataset
-    data = load_data(args)
-    features = mx.nd.array(data.features)
-    labels = mx.nd.array(data.labels)
-    train_mask = mx.nd.array(data.train_mask)
-    val_mask = mx.nd.array(data.val_mask)
-    test_mask = mx.nd.array(data.test_mask)
+    if args.dataset == 'cora':
+        data = CoraGraphDataset()
+    elif args.dataset == 'citeseer':
+        data = CiteseerGraphDataset()
+    elif args.dataset == 'pubmed':
+        data = PubmedGraphDataset()
+    else:
+        raise ValueError('Unknown dataset: {}'.format(args.dataset))
+
+    g = data[0]
+    if args.gpu < 0:
+        cuda = False
+        ctx = mx.cpu(0)
+    else:
+        cuda = True
+        ctx = mx.gpu(args.gpu)
+        g = g.to(ctx)
 
+    features = g.ndata['feat']
+    labels = mx.nd.array(g.ndata['label'], dtype="float32", ctx=ctx)
+    train_mask = g.ndata['train_mask']
+    val_mask = g.ndata['val_mask']
+    test_mask = g.ndata['test_mask']
     in_feats = features.shape[1]
     n_classes = data.num_labels
     n_edges = data.graph.number_of_edges()
@@ -40,25 +57,10 @@ def main(args):
               val_mask.sum().asscalar(),
               test_mask.sum().asscalar()))
 
-    if args.gpu < 0:
-        cuda = False
-        ctx = mx.cpu(0)
-    else:
-        cuda = True
-        ctx = mx.gpu(args.gpu)
-
-    features = features.as_in_context(ctx)
-    labels = labels.as_in_context(ctx)
-    train_mask = train_mask.as_in_context(ctx)
-    val_mask = val_mask.as_in_context(ctx)
-    test_mask = test_mask.as_in_context(ctx)
-
-    # create GCN model
-    g = data.graph
+    # add self loop
     if args.self_loop:
-        g.remove_edges_from(nx.selfloop_edges(g))
-        g.add_edges_from(zip(g.nodes(), g.nodes()))
-    g = dgl.graph(g).to(ctx)
+        g = dgl.remove_self_loop(g)
+        g = dgl.add_self_loop(g)
     # normalization
     degs = g.in_degrees().astype('float32')
     norm = mx.nd.power(degs, -0.5)

examples/mxnet/graphsage/main.py

@@ -11,8 +11,9 @@ import networkx as nx
 import mxnet as mx
 from mxnet import nd, gluon
 from mxnet.gluon import nn
-from dgl import DGLGraph
-from dgl.data import register_data_args, load_data
+import dgl
+from dgl.data import register_data_args
+from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
 from dgl.nn.mxnet.conv import SAGEConv
 
 
@@ -52,13 +53,29 @@ def evaluate(model, features, labels, mask):
 
 def main(args):
     # load and preprocess dataset
-    data = load_data(args)
-    features = nd.array(data.features)
-    labels = nd.array(data.labels)
-    train_mask = nd.array(data.train_mask)
-    val_mask = nd.array(data.val_mask)
-    test_mask = nd.array(data.test_mask)
+    if args.dataset == 'cora':
+        data = CoraGraphDataset()
+    elif args.dataset == 'citeseer':
+        data = CiteseerGraphDataset()
+    elif args.dataset == 'pubmed':
+        data = PubmedGraphDataset()
+    else:
+        raise ValueError('Unknown dataset: {}'.format(args.dataset))
+
+    g = data[0]
+    if args.gpu < 0:
+        cuda = False
+        ctx = mx.cpu(0)
+    else:
+        cuda = True
+        ctx = mx.gpu(args.gpu)
+        g = g.to(ctx)
 
+    features = g.ndata['feat']
+    labels = mx.nd.array(g.ndata['label'], dtype="float32", ctx=ctx)
+    train_mask = g.ndata['train_mask']
+    val_mask = g.ndata['val_mask']
+    test_mask = g.ndata['test_mask']
     in_feats = features.shape[1]
     n_classes = data.num_labels
     n_edges = data.graph.number_of_edges()
@@ -73,22 +90,9 @@ def main(args):
            val_mask.sum().asscalar(),
            test_mask.sum().asscalar()))
 
-    if args.gpu < 0:
-        ctx = mx.cpu(0)
-    else:
-        ctx = mx.gpu(args.gpu)
-        print("use cuda:", args.gpu)
-
-    features = features.as_in_context(ctx)
-    labels = labels.as_in_context(ctx)
-    train_mask = train_mask.as_in_context(ctx)
-    val_mask = val_mask.as_in_context(ctx)
-    test_mask = test_mask.as_in_context(ctx)
-
-    # graph preprocess and calculate normalization factor
-    g = data.graph
-    g.remove_edges_from(nx.selfloop_edges(g))
-    g = DGLGraph(g)
+    # add self loop
+    g = dgl.remove_self_loop(g)
+    g = dgl.add_self_loop(g)
     n_edges = g.number_of_edges()
 
     # create GraphSAGE model

examples/mxnet/monet/citation.py

@@ -5,8 +5,9 @@ import networkx as nx
 import mxnet as mx
 from mxnet import gluon, nd
 from mxnet.gluon import nn
-from dgl import DGLGraph
-from dgl.data import register_data_args, load_data
+import dgl
+from dgl.data import register_data_args
+from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
 from dgl.nn.mxnet.conv import GMMConv
 
 
@@ -59,13 +60,29 @@ def evaluate(model, features, pseudo, labels, mask):
 
 def main(args):
     # load and preprocess dataset
-    data = load_data(args)
-    features = nd.array(data.features)
-    labels = nd.array(data.labels)
-    train_mask = nd.array(data.train_mask)
-    val_mask = nd.array(data.val_mask)
-    test_mask = nd.array(data.test_mask)
+    if args.dataset == 'cora':
+        data = CoraGraphDataset()
+    elif args.dataset == 'citeseer':
+        data = CiteseerGraphDataset()
+    elif args.dataset == 'pubmed':
+        data = PubmedGraphDataset()
+    else:
+        raise ValueError('Unknown dataset: {}'.format(args.dataset))
+
+    g = data[0]
+    if args.gpu < 0:
+        cuda = False
+        ctx = mx.cpu(0)
+    else:
+        cuda = True
+        ctx = mx.gpu(args.gpu)
+        g = g.to(ctx)
 
+    features = g.ndata['feat']
+    labels = mx.nd.array(g.ndata['label'], dtype="float32", ctx=ctx)
+    train_mask = g.ndata['train_mask']
+    val_mask = g.ndata['val_mask']
+    test_mask = g.ndata['test_mask']
     in_feats = features.shape[1]
     n_classes = data.num_labels
     n_edges = data.graph.number_of_edges()
@@ -80,29 +97,19 @@ def main(args):
            val_mask.sum().asscalar(),
            test_mask.sum().asscalar()))
 
-    if args.gpu < 0:
-        ctx = mx.cpu(0)
-    else:
-        ctx = mx.gpu(args.gpu)
-        print("use cuda:", args.gpu)
-
-    features = features.as_in_context(ctx)
-    labels = labels.as_in_context(ctx)
-    train_mask = train_mask.as_in_context(ctx)
-    val_mask = val_mask.as_in_context(ctx)
-    test_mask = test_mask.as_in_context(ctx)
-
-    # graph preprocess and calculate normalization factor
-    g = data.graph
-    g.remove_edges_from(nx.selfloop_edges(g))
-    g = DGLGraph(g)
+    # add self loop
+    g = dgl.remove_self_loop(g)
+    g = dgl.add_self_loop(g)
+
     n_edges = g.number_of_edges()
     us, vs = g.edges()
+    us = us.asnumpy()
+    vs = vs.asnumpy()
     pseudo = []
     for i in range(g.number_of_edges()):
         pseudo.append([
-            1 / np.sqrt(g.in_degree(us[i].asscalar())),
-            1 / np.sqrt(g.in_degree(vs[i].asscalar()))
+            1 / np.sqrt(g.in_degree(us[i])),
+            1 / np.sqrt(g.in_degree(vs[i]))
         ])
     pseudo = nd.array(pseudo, ctx=ctx)
 

examples/mxnet/sgc/sgc.py

@@ -10,8 +10,9 @@ import numpy as np
 import mxnet as mx
 from mxnet import nd, gluon
 from mxnet.gluon import nn
-from dgl import DGLGraph
-from dgl.data import register_data_args, load_data
+import dgl
+from dgl.data import register_data_args
+from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
 from dgl.nn.mxnet.conv import SGConv
 
 
@@ -22,13 +23,29 @@ def evaluate(model, g, features, labels, mask):
 
 def main(args):
     # load and preprocess dataset
-    data = load_data(args)
-    features = nd.array(data.features)
-    labels = nd.array(data.labels)
-    train_mask = nd.array(data.train_mask)
-    val_mask = nd.array(data.val_mask)
-    test_mask = nd.array(data.test_mask)
+    if args.dataset == 'cora':
+        data = CoraGraphDataset()
+    elif args.dataset == 'citeseer':
+        data = CiteseerGraphDataset()
+    elif args.dataset == 'pubmed':
+        data = PubmedGraphDataset()
+    else:
+        raise ValueError('Unknown dataset: {}'.format(args.dataset))
+
+    g = data[0]
+    if args.gpu < 0:
+        cuda = False
+        ctx = mx.cpu(0)
+    else:
+        cuda = True
+        ctx = mx.gpu(args.gpu)
+        g = g.to(ctx)
 
+    features = g.ndata['feat']
+    labels = mx.nd.array(g.ndata['label'], dtype="float32", ctx=ctx)
+    train_mask = g.ndata['train_mask']
+    val_mask = g.ndata['val_mask']
+    test_mask = g.ndata['test_mask']
     in_feats = features.shape[1]
     n_classes = data.num_labels
     n_edges = data.graph.number_of_edges()
@@ -43,22 +60,9 @@ def main(args):
               val_mask.sum().asscalar(),
               test_mask.sum().asscalar()))
 
-    if args.gpu < 0:
-        ctx = mx.cpu(0)
-    else:
-        ctx = mx.gpu(args.gpu)
-
-    features = features.as_in_context(ctx)
-    labels = labels.as_in_context(ctx)
-    train_mask = train_mask.as_in_context(ctx)
-    val_mask = val_mask.as_in_context(ctx)
-    test_mask = test_mask.as_in_context(ctx)
-
-    # graph preprocess and calculate normalization factor
-    g = DGLGraph(data.graph)
-    n_edges = g.number_of_edges()
     # add self loop
-    g.add_edges(g.nodes(), g.nodes())
+    g = dgl.remove_self_loop(g)
+    g = dgl.add_self_loop(g)
 
     # create SGC model
     model = SGConv(in_feats,

examples/mxnet/tagcn/train.py

@@ -4,8 +4,9 @@ import networkx as nx
 import mxnet as mx
 from mxnet import gluon
 
-from dgl import DGLGraph
-from dgl.data import register_data_args, load_data
+import dgl
+from dgl.data import register_data_args
+from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
 
 from tagcn import TAGCN
 
@@ -16,12 +17,29 @@ def evaluate(model, features, labels, mask):
 
 def main(args):
     # load and preprocess dataset
-    data = load_data(args)
-    features = mx.nd.array(data.features)
-    labels = mx.nd.array(data.labels)
-    train_mask = mx.nd.array(data.train_mask)
-    val_mask = mx.nd.array(data.val_mask)
-    test_mask = mx.nd.array(data.test_mask)
+    if args.dataset == 'cora':
+        data = CoraGraphDataset()
+    elif args.dataset == 'citeseer':
+        data = CiteseerGraphDataset()
+    elif args.dataset == 'pubmed':
+        data = PubmedGraphDataset()
+    else:
+        raise ValueError('Unknown dataset: {}'.format(args.dataset))
+
+    g = data[0]
+    if args.gpu < 0:
+        cuda = False
+        ctx = mx.cpu(0)
+    else:
+        cuda = True
+        ctx = mx.gpu(args.gpu)
+        g = g.to(ctx)
+
+    features = g.ndata['feat']
+    labels = mx.nd.array(g.ndata['label'], dtype="float32", ctx=ctx)
+    train_mask = g.ndata['train_mask']
+    val_mask = g.ndata['val_mask']
+    test_mask = g.ndata['test_mask']
     in_feats = features.shape[1]
     n_classes = data.num_labels
     n_edges = data.graph.number_of_edges()
@@ -36,26 +54,9 @@ def main(args):
               val_mask.sum().asscalar(),
               test_mask.sum().asscalar()))
 
-    if args.gpu < 0:
-        cuda = False
-        ctx = mx.cpu(0)
-    else:
-        cuda = True
-        ctx = mx.gpu(args.gpu)
-
-    features = features.as_in_context(ctx)
-    labels = labels.as_in_context(ctx)
-    train_mask = train_mask.as_in_context(ctx)
-    val_mask = val_mask.as_in_context(ctx)
-    test_mask = test_mask.as_in_context(ctx)
-
-    # graph preprocess and calculate normalization factor
-    g = data.graph
     # add self loop
-    if args.self_loop:
-        g.remove_edges_from(nx.selfloop_edges(g))
-        g.add_edges_from(zip(g.nodes(), g.nodes()))
-    g = DGLGraph(g)
+    g = dgl.remove_self_loop(g)
+    g = dgl.add_self_loop(g)
 
     # create TAGCN model
     model = TAGCN(g,

examples/pytorch/appnp/train.py

@@ -3,8 +3,8 @@ import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from dgl import DGLGraph
-from dgl.data import register_data_args, load_data
+from dgl.data import register_data_args
+from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
 import dgl
 from appnp import APPNP
 
@@ -22,17 +22,27 @@ def evaluate(model, features, labels, mask):
 
 def main(args):
     # load and preprocess dataset
-    data = load_data(args)
-    features = torch.FloatTensor(data.features)
-    labels = torch.LongTensor(data.labels)
-    if hasattr(torch, 'BoolTensor'):
-        train_mask = torch.BoolTensor(data.train_mask)
-        val_mask = torch.BoolTensor(data.val_mask)
-        test_mask = torch.BoolTensor(data.test_mask)
+    if args.dataset == 'cora':
+        data = CoraGraphDataset()
+    elif args.dataset == 'citeseer':
+        data = CiteseerGraphDataset()
+    elif args.dataset == 'pubmed':
+        data = PubmedGraphDataset()
     else:
-        train_mask = torch.ByteTensor(data.train_mask)
-        val_mask = torch.ByteTensor(data.val_mask)
-        test_mask = torch.ByteTensor(data.test_mask)
+        raise ValueError('Unknown dataset: {}'.format(args.dataset))
+
+    g = data[0]
+    if args.gpu < 0:
+        cuda = False
+    else:
+        cuda = True
+        g = g.to(args.gpu)
+
+    features = g.ndata['feat']
+    labels = g.ndata['label']
+    train_mask = g.ndata['train_mask']
+    val_mask = g.ndata['val_mask']
+    test_mask = g.ndata['test_mask']
     in_feats = features.shape[1]
     n_classes = data.num_labels
     n_edges = data.graph.number_of_edges()
@@ -47,27 +57,10 @@ def main(args):
            val_mask.int().sum().item(),
            test_mask.int().sum().item()))
 
-    if args.gpu < 0:
-        cuda = False
-    else:
-        cuda = True
-        torch.cuda.set_device(args.gpu)
-        features = features.cuda()
-        labels = labels.cuda()
-        train_mask = train_mask.cuda()
-        val_mask = val_mask.cuda()
-        test_mask = test_mask.cuda()
-
-    # graph preprocess and calculate normalization factor
-    g = DGLGraph(data.graph)
     n_edges = g.number_of_edges()
     # add self loop
-    g.add_edges(g.nodes(), g.nodes())
-    g.set_n_initializer(dgl.init.zero_initializer)
-    g.set_e_initializer(dgl.init.zero_initializer)
-
-    if args.gpu >= 0:
-        g = g.to(args.gpu)
+    g = dgl.remove_self_loop(g)
+    g = dgl.add_self_loop(g)
 
     # create APPNP model
     model = APPNP(g,

examples/pytorch/gat/train.py

@@ -14,8 +14,10 @@ import networkx as nx
 import time
 import torch
 import torch.nn.functional as F
-from dgl import DGLGraph
-from dgl.data import register_data_args, load_data
+import dgl
+from dgl.data import register_data_args
+from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
+
 from gat import GAT
 from utils import EarlyStopping
 
@@ -37,17 +39,27 @@ def evaluate(model, features, labels, mask):
 
 def main(args):
     # load and preprocess dataset
-    data = load_data(args)
-    features = torch.FloatTensor(data.features)
-    labels = torch.LongTensor(data.labels)
-    if hasattr(torch, 'BoolTensor'):
-        train_mask = torch.BoolTensor(data.train_mask)
-        val_mask = torch.BoolTensor(data.val_mask)
-        test_mask = torch.BoolTensor(data.test_mask)
+    if args.dataset == 'cora':
+        data = CoraGraphDataset()
+    elif args.dataset == 'citeseer':
+        data = CiteseerGraphDataset()
+    elif args.dataset == 'pubmed':
+        data = PubmedGraphDataset()
     else:
-        train_mask = torch.ByteTensor(data.train_mask)
-        val_mask = torch.ByteTensor(data.val_mask)
-        test_mask = torch.ByteTensor(data.test_mask)
+        raise ValueError('Unknown dataset: {}'.format(args.dataset))
+
+    g = data[0]
+    if args.gpu < 0:
+        cuda = False
+    else:
+        cuda = True
+        g = g.to(args.gpu)
+
+    features = g.ndata['feat']
+    labels = g.ndata['label']
+    train_mask = g.ndata['train_mask']
+    val_mask = g.ndata['val_mask']
+    test_mask = g.ndata['test_mask']
     num_feats = features.shape[1]
     n_classes = data.num_labels
     n_edges = data.graph.number_of_edges()
@@ -62,24 +74,9 @@ def main(args):
            val_mask.int().sum().item(),
            test_mask.int().sum().item()))
 
-    if args.gpu < 0:
-        cuda = False
-    else:
-        cuda = True
-        torch.cuda.set_device(args.gpu)
-        features = features.cuda()
-        labels = labels.cuda()
-        train_mask = train_mask.cuda()
-        val_mask = val_mask.cuda()
-        test_mask = test_mask.cuda()
-
-    g = data.graph
     # add self loop
-    g.remove_edges_from(nx.selfloop_edges(g))
-    g = DGLGraph(g)
-    g.add_edges(g.nodes(), g.nodes())
-    if cuda:
-        g = g.to(args.gpu)
+    g = dgl.remove_self_loop(g)
+    g = dgl.add_self_loop(g)
     n_edges = g.number_of_edges()
     # create model
     heads = ([args.num_heads] * args.num_layers) + [args.num_out_heads]

examples/pytorch/gcn/gcn_mp.py

@@ -11,8 +11,9 @@ import networkx as nx
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from dgl import DGLGraph
-from dgl.data import register_data_args, load_data
+import dgl
+from dgl.data import register_data_args
+from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
 
 
 def gcn_msg(edge):
@@ -116,17 +117,27 @@ def evaluate(model, features, labels, mask):
 
 def main(args):
     # load and preprocess dataset
-    data = load_data(args)
-    features = torch.FloatTensor(data.features)
-    labels = torch.LongTensor(data.labels)
-    if hasattr(torch, 'BoolTensor'):
-        train_mask = torch.BoolTensor(data.train_mask)
-        val_mask = torch.BoolTensor(data.val_mask)
-        test_mask = torch.BoolTensor(data.test_mask)
+    if args.dataset == 'cora':
+        data = CoraGraphDataset()
+    elif args.dataset == 'citeseer':
+        data = CiteseerGraphDataset()
+    elif args.dataset == 'pubmed':
+        data = PubmedGraphDataset()
     else:
-        train_mask = torch.ByteTensor(data.train_mask)
-        val_mask = torch.ByteTensor(data.val_mask)
-        test_mask = torch.ByteTensor(data.test_mask)
+        raise ValueError('Unknown dataset: {}'.format(args.dataset))
+
+    g = data[0]
+    if args.gpu < 0:
+        cuda = False
+    else:
+        cuda = True
+        g = g.to(args.gpu)
+
+    features = g.ndata['feat']
+    labels = g.ndata['label']
+    train_mask = g.ndata['train_mask']
+    val_mask = g.ndata['val_mask']
+    test_mask = g.ndata['test_mask']
     in_feats = features.shape[1]
     n_classes = data.num_labels
     n_edges = data.graph.number_of_edges()
@@ -141,24 +152,11 @@ def main(args):
               val_mask.int().sum().item(),
               test_mask.int().sum().item()))
 
-    if args.gpu < 0:
-        cuda = False
-    else:
-        cuda = True
-        torch.cuda.set_device(args.gpu)
-        features = features.cuda()
-        labels = labels.cuda()
-        train_mask = train_mask.cuda()
-        val_mask = val_mask.cuda()
-        test_mask = test_mask.cuda()
-
-    # graph preprocess and calculate normalization factor
-    g = data.graph
-    g.remove_edges_from(nx.selfloop_edges(g))
-    g = DGLGraph(g)
     # add self loop
-    g.add_edges(g.nodes(), g.nodes())
+    g = dgl.remove_self_loop(g)
+    g = dgl.add_self_loop(g)
     n_edges = g.number_of_edges()
+
     # normalization
     degs = g.in_degrees().float()
     norm = torch.pow(degs, -0.5)

examples/pytorch/gcn/train.py

@@ -4,8 +4,9 @@ import networkx as nx
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from dgl import DGLGraph
-from dgl.data import register_data_args, load_data
+import dgl
+from dgl.data import register_data_args
+from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
 
 from gcn import GCN
 #from gcn_mp import GCN
@@ -23,17 +24,27 @@ def evaluate(model, features, labels, mask):
 
 def main(args):
     # load and preprocess dataset
-    data = load_data(args)
-    features = torch.FloatTensor(data.features)
-    labels = torch.LongTensor(data.labels)
-    if hasattr(torch, 'BoolTensor'):
-        train_mask = torch.BoolTensor(data.train_mask)
-        val_mask = torch.BoolTensor(data.val_mask)
-        test_mask = torch.BoolTensor(data.test_mask)
+    if args.dataset == 'cora':
+        data = CoraGraphDataset()
+    elif args.dataset == 'citeseer':
+        data = CiteseerGraphDataset()
+    elif args.dataset == 'pubmed':
+        data = PubmedGraphDataset()
     else:
-        train_mask = torch.ByteTensor(data.train_mask)
-        val_mask = torch.ByteTensor(data.val_mask)
-        test_mask = torch.ByteTensor(data.test_mask)
+        raise ValueError('Unknown dataset: {}'.format(args.dataset))
+
+    g = data[0]
+    if args.gpu < 0:
+        cuda = False
+    else:
+        cuda = True
+        g = g.to(args.gpu)
+
+    features = g.ndata['feat']
+    labels = g.ndata['label']
+    train_mask = g.ndata['train_mask']
+    val_mask = g.ndata['val_mask']
+    test_mask = g.ndata['test_mask']
     in_feats = features.shape[1]
     n_classes = data.num_labels
     n_edges = data.graph.number_of_edges()
@@ -48,27 +59,12 @@ def main(args):
               val_mask.int().sum().item(),
               test_mask.int().sum().item()))
 
-    if args.gpu < 0:
-        cuda = False
-    else:
-        cuda = True
-        torch.cuda.set_device(args.gpu)
-        features = features.cuda()
-        labels = labels.cuda()
-        train_mask = train_mask.cuda()
-        val_mask = val_mask.cuda()
-        test_mask = test_mask.cuda()
-
-    # graph preprocess and calculate normalization factor
-    g = data.graph
     # add self loop
     if args.self_loop:
-        g.remove_edges_from(nx.selfloop_edges(g))
-        g.add_edges_from(zip(g.nodes(), g.nodes()))
-    g = DGLGraph(g)
+        g = dgl.remove_self_loop(g)
+        g = dgl.add_self_loop(g)
     n_edges = g.number_of_edges()
-    if cuda:
-        g = g.to(args.gpu)
+
     # normalization
     degs = g.in_degrees().float()
     norm = torch.pow(degs, -0.5)

examples/pytorch/sgc/sgc.py

@@ -11,8 +11,9 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import dgl.function as fn
-from dgl import DGLGraph
-from dgl.data import register_data_args, load_data
+import dgl
+from dgl.data import register_data_args
+from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
 from dgl.nn.pytorch.conv import SGConv
 
 
@@ -27,17 +28,27 @@ def evaluate(model, g, features, labels, mask):
 
 def main(args):
     # load and preprocess dataset
-    data = load_data(args)
-    features = torch.FloatTensor(data.features)
-    labels = torch.LongTensor(data.labels)
-    if hasattr(torch, 'BoolTensor'):
-        train_mask = torch.BoolTensor(data.train_mask)
-        val_mask = torch.BoolTensor(data.val_mask)
-        test_mask = torch.BoolTensor(data.test_mask)
+    if args.dataset == 'cora':
+        data = CoraGraphDataset()
+    elif args.dataset == 'citeseer':
+        data = CiteseerGraphDataset()
+    elif args.dataset == 'pubmed':
+        data = PubmedGraphDataset()
     else:
-        train_mask = torch.ByteTensor(data.train_mask)
-        val_mask = torch.ByteTensor(data.val_mask)
-        test_mask = torch.ByteTensor(data.test_mask)
+        raise ValueError('Unknown dataset: {}'.format(args.dataset))
+
+    g = data[0]
+    if args.gpu < 0:
+        cuda = False
+    else:
+        cuda = True
+        g = g.to(args.gpu)
+
+    features = g.ndata['feat']
+    labels = g.ndata['label']
+    train_mask = g.ndata['train_mask']
+    val_mask = g.ndata['val_mask']
+    test_mask = g.ndata['test_mask']
     in_feats = features.shape[1]
     n_classes = data.num_labels
     n_edges = data.graph.number_of_edges()
@@ -52,22 +63,10 @@ def main(args):
               val_mask.int().sum().item(),
               test_mask.int().sum().item()))
 
-    if args.gpu < 0:
-        cuda = False
-    else:
-        cuda = True
-        torch.cuda.set_device(args.gpu)
-        features = features.cuda()
-        labels = labels.cuda()
-        train_mask = train_mask.cuda()
-        val_mask = val_mask.cuda()
-        test_mask = test_mask.cuda()
-
-    # graph preprocess and calculate normalization factor
-    g = DGLGraph(data.graph)
     n_edges = g.number_of_edges()
     # add self loop
-    g.add_edges(g.nodes(), g.nodes())
+    g = dgl.remove_self_loop(g)
+    g = dgl.add_self_loop(g)
 
     # create SGC model
     model = SGConv(in_feats,
@@ -78,7 +77,6 @@ def main(args):
 
     if cuda:
         model.cuda()
-        g = g.to(args.gpu)
     loss_fcn = torch.nn.CrossEntropyLoss()
 
     # use optimizer

examples/tensorflow/dgi/train.py

@@ -4,8 +4,9 @@ import numpy as np
 import networkx as nx
 import tensorflow as tf
 from tensorflow.keras import layers
-from dgl import DGLGraph
-from dgl.data import register_data_args, load_data
+import dgl
+from dgl.data import register_data_args
+from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
 from dgi import DGI, Classifier
 
 
@@ -20,28 +21,36 @@ def evaluate(model, features, labels, mask):
 
 def main(args):
     # load and preprocess dataset
-    data = load_data(args)
+    if args.dataset == 'cora':
+        data = CoraGraphDataset()
+    elif args.dataset == 'citeseer':
+        data = CiteseerGraphDataset()
+    elif args.dataset == 'pubmed':
+        data = PubmedGraphDataset()
+    else:
+        raise ValueError('Unknown dataset: {}'.format(args.dataset))
+
+    g = data[0]
     if args.gpu < 0:
         device = "/cpu:0"
     else:
         device = "/gpu:{}".format(args.gpu)
+        g = g.to(device)
+
     with tf.device(device):
-        features = tf.convert_to_tensor(data.features, dtype=tf.float32)
-        labels = tf.convert_to_tensor(data.labels, dtype=tf.int64)
-        train_mask = tf.convert_to_tensor(data.train_mask, dtype=tf.bool)
-        val_mask = tf.convert_to_tensor(data.val_mask, dtype=tf.bool)
-        test_mask = tf.convert_to_tensor(data.test_mask, dtype=tf.bool)
+        features = g.ndata['feat']
+        labels = g.ndata['label']
+        train_mask = g.ndata['train_mask']
+        val_mask = g.ndata['val_mask']
+        test_mask = g.ndata['test_mask']
         in_feats = features.shape[1]
         n_classes = data.num_labels
         n_edges = data.graph.number_of_edges()
 
-        # graph preprocess
-        g = data.graph
         # add self loop
         if args.self_loop:
-            g.remove_edges_from(nx.selfloop_edges(g))
-            g.add_edges_from(zip(g.nodes(), g.nodes()))
-        g = DGLGraph(g)
+            g = dgl.remove_self_loop(g)
+            g = dgl.add_self_loop(g)
         n_edges = g.number_of_edges()
 
         # create DGI model

examples/tensorflow/gat/README.md

@@ -14,6 +14,7 @@ Dependencies
 
 ```bash
 pip install tensorflow requests
+DGLBACKEND=tensorflow
 ```
 
 How to run

examples/tensorflow/gat/train.py

@@ -15,8 +15,9 @@ import numpy as np
 import networkx as nx
 import time
 import tensorflow as tf
-from dgl import DGLGraph
-from dgl.data import register_data_args, load_data
+import dgl
+from dgl.data import register_data_args
+from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
 from gat import GAT
 from utils import EarlyStopping
 
@@ -35,20 +36,28 @@ def evaluate(model, features, labels, mask):
 
 def main(args):
     # load and preprocess dataset
-    data = load_data(args)
+    if args.dataset == 'cora':
+        data = CoraGraphDataset()
+    elif args.dataset == 'citeseer':
+        data = CiteseerGraphDataset()
+    elif args.dataset == 'pubmed':
+        data = PubmedGraphDataset()
+    else:
+        raise ValueError('Unknown dataset: {}'.format(args.dataset))
 
+    g = data[0]
     if args.gpu < 0:
         device = "/cpu:0"
     else:
         device = "/gpu:{}".format(args.gpu)
+        g = g.to(device)
 
     with tf.device(device):
-
-        features = tf.convert_to_tensor(data.features, dtype=tf.float32)
-        labels = tf.convert_to_tensor(data.labels, dtype=tf.int64)
-        train_mask = tf.convert_to_tensor(data.train_mask, dtype=tf.bool)
-        val_mask = tf.convert_to_tensor(data.val_mask, dtype=tf.bool)
-        test_mask = tf.convert_to_tensor(data.test_mask, dtype=tf.bool)
+        features = g.ndata['feat']
+        labels = g.ndata['label']
+        train_mask = g.ndata['train_mask']
+        val_mask = g.ndata['val_mask']
+        test_mask = g.ndata['test_mask']
         num_feats = features.shape[1]
         n_classes = data.num_labels
         n_edges = data.graph.number_of_edges()
@@ -63,11 +72,8 @@ def main(args):
                val_mask.numpy().sum(),
                test_mask.numpy().sum()))
 
-        g = data.graph
-        # add self loop
-        g.remove_edges_from(nx.selfloop_edges(g))
-        g = DGLGraph(g).to(device)
-        g.add_edges(g.nodes(), g.nodes())
+        g = dgl.remove_self_loop(g)
+        g = dgl.add_self_loop(g)
         n_edges = g.number_of_edges()
         # create model
         heads = ([args.num_heads] * args.num_layers) + [args.num_out_heads]

examples/tensorflow/gcn/README.md

@@ -12,6 +12,7 @@ Dependencies
 
 ``bash
 pip install tensorflow requests
+export DGLBACKEND=tensorflow
 ``
 
 Codes

examples/tensorflow/gcn/gcn_builtin.py

@@ -4,9 +4,10 @@ import math
 import numpy as np
 import networkx as nx
 import tensorflow as tf
-from dgl import DGLGraph
 import dgl.function as fn
-from dgl.data import register_data_args, load_data
+import dgl
+from dgl.data import register_data_args
+from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
 from tensorflow.keras import layers
 
 
@@ -94,19 +95,28 @@ def evaluate(model, features, labels, mask):
 
 def main(args):
     # load and preprocess dataset
-    data = load_data(args)
+    if args.dataset == 'cora':
+        data = CoraGraphDataset()
+    elif args.dataset == 'citeseer':
+        data = CiteseerGraphDataset()
+    elif args.dataset == 'pubmed':
+        data = PubmedGraphDataset()
+    else:
+        raise ValueError('Unknown dataset: {}'.format(args.dataset))
 
+    g = data[0]
     if args.gpu < 0:
         device = "/cpu:0"
     else:
         device = "/gpu:{}".format(args.gpu)
+        g = g.to(device)
 
     with tf.device(device):
-        features = tf.convert_to_tensor(data.features, dtype=tf.float32)
-        labels = tf.convert_to_tensor(data.labels, dtype=tf.int64)
-        train_mask = tf.convert_to_tensor(data.train_mask, dtype=tf.bool)
-        val_mask = tf.convert_to_tensor(data.val_mask, dtype=tf.bool)
-        test_mask = tf.convert_to_tensor(data.test_mask, dtype=tf.bool)
+        features = g.ndata['feat']
+        labels = g.ndata['label']
+        train_mask = g.ndata['train_mask']
+        val_mask = g.ndata['val_mask']
+        test_mask = g.ndata['test_mask']
         in_feats = features.shape[1]
         n_classes = data.num_labels
         n_edges = data.graph.number_of_edges()
@@ -121,12 +131,9 @@ def main(args):
                val_mask.numpy().sum(),
                test_mask.numpy().sum()))
 
-        # graph preprocess and calculate normalization factor
-        g = data.graph
-        g.remove_edges_from(nx.selfloop_edges(g))
-        g = DGLGraph(g)
-        # # add self loop
-        g.add_edges(g.nodes(), g.nodes())
+        # add self loop
+        g = dgl.remove_self_loop(g)
+        g = dgl.add_self_loop(g)
         n_edges = g.number_of_edges()
         # # normalization
         degs = tf.cast(tf.identity(g.in_degrees()), dtype=tf.float32)

examples/tensorflow/gcn/gcn_mp.py

@@ -4,9 +4,9 @@ import math
 import numpy as np
 import networkx as nx
 import tensorflow as tf
-from dgl import DGLGraph
-import dgl.function as fn
-from dgl.data import register_data_args, load_data
+import dgl
+from dgl.data import register_data_args
+from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
 from tensorflow.keras import layers
 
 def gcn_msg(edge):
@@ -100,19 +100,28 @@ def evaluate(model, features, labels, mask):
 
 def main(args):
     # load and preprocess dataset
-    data = load_data(args)
+    if args.dataset == 'cora':
+        data = CoraGraphDataset()
+    elif args.dataset == 'citeseer':
+        data = CiteseerGraphDataset()
+    elif args.dataset == 'pubmed':
+        data = PubmedGraphDataset()
+    else:
+        raise ValueError('Unknown dataset: {}'.format(args.dataset))
 
+    g = data[0]
     if args.gpu < 0:
         device = "/cpu:0"
     else:
         device = "/gpu:{}".format(args.gpu)
+        g = g.to(device)
 
     with tf.device(device):
-        features = tf.convert_to_tensor(data.features, dtype=tf.float32)
-        labels = tf.convert_to_tensor(data.labels, dtype=tf.int64)
-        train_mask = tf.convert_to_tensor(data.train_mask, dtype=tf.bool)
-        val_mask = tf.convert_to_tensor(data.val_mask, dtype=tf.bool)
-        test_mask = tf.convert_to_tensor(data.test_mask, dtype=tf.bool)
+        features = g.ndata['feat']
+        labels = g.ndata['label']
+        train_mask = g.ndata['train_mask']
+        val_mask = g.ndata['val_mask']
+        test_mask = g.ndata['test_mask']
         in_feats = features.shape[1]
         n_classes = data.num_labels
         n_edges = data.graph.number_of_edges()
@@ -127,12 +136,11 @@ def main(args):
                val_mask.numpy().sum(),
                test_mask.numpy().sum()))
 
-        # graph preprocess and calculate normalization factor
-        g = data.graph
-        g.remove_edges_from(nx.selfloop_edges(g))
-        g = DGLGraph(g)
-        # # add self loop
-        g.add_edges(g.nodes(), g.nodes())
+        # add self loop
+        if args.self_loop:
+            g = dgl.remove_self_loop(g)
+            g = dgl.add_self_loop(g)
+        n_edges = g.number_of_edges()
         n_edges = g.number_of_edges()
         # # normalization
         degs = tf.cast(tf.identity(g.in_degrees()), dtype=tf.float32)
@@ -203,6 +211,8 @@ if __name__ == '__main__':
                         help="number of hidden gcn layers")
     parser.add_argument("--weight-decay", type=float, default=5e-4,
                         help="Weight for L2 loss")
+    parser.add_argument("--self-loop", action='store_true',
+                        help="graph self-loop (default=False)")
     args = parser.parse_args()
     print(args)
 

examples/tensorflow/gcn/train.py

@@ -3,8 +3,9 @@ import time
 import numpy as np
 import networkx as nx
 import tensorflow as tf
-from dgl import DGLGraph
-from dgl.data import register_data_args, load_data
+import dgl
+from dgl.data import register_data_args
+from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
 from gcn import GCN
 
 
@@ -19,19 +20,28 @@ def evaluate(model, features, labels, mask):
 
 def main(args):
     # load and preprocess dataset
-    data = load_data(args)
+    if args.dataset == 'cora':
+        data = CoraGraphDataset()
+    elif args.dataset == 'citeseer':
+        data = CiteseerGraphDataset()
+    elif args.dataset == 'pubmed':
+        data = PubmedGraphDataset()
+    else:
+        raise ValueError('Unknown dataset: {}'.format(args.dataset))
 
+    g = data[0]
     if args.gpu < 0:
         device = "/cpu:0"
     else:
         device = "/gpu:{}".format(args.gpu)
+        g = g.to(device)
 
     with tf.device(device):
-        features = tf.convert_to_tensor(data.features, dtype=tf.float32)
-        labels = tf.convert_to_tensor(data.labels, dtype=tf.int64)
-        train_mask = tf.convert_to_tensor(data.train_mask, dtype=tf.bool)
-        val_mask = tf.convert_to_tensor(data.val_mask, dtype=tf.bool)
-        test_mask = tf.convert_to_tensor(data.test_mask, dtype=tf.bool)
+        features = g.ndata['feat']
+        labels = g.ndata['label']
+        train_mask = g.ndata['train_mask']
+        val_mask = g.ndata['val_mask']
+        test_mask = g.ndata['test_mask']
         in_feats = features.shape[1]
         n_classes = data.num_labels
         n_edges = data.graph.number_of_edges()
@@ -46,12 +56,10 @@ def main(args):
                val_mask.numpy().sum(),
                test_mask.numpy().sum()))
 
-        # graph preprocess and calculate normalization factor
-        g = data.graph
+        # add self loop
         if args.self_loop:
-            g.remove_edges_from(nx.selfloop_edges(g))
-            g.add_edges_from(zip(g.nodes(), g.nodes()))
-        g = DGLGraph(g).to(device)
+            g = dgl.remove_self_loop(g)
+            g = dgl.add_self_loop(g)
         n_edges = g.number_of_edges()
         # normalization
         degs = tf.cast(tf.identity(g.in_degrees()), dtype=tf.float32)