[Example][Refactor] Refactor GCN example (#4160)

* Refactor GCN example

* Refactor GCN based on graphsage

* Readme update

* Minor update

* update

* Remove user-defined GCN implementation

* README update

* Update

* Update CONTRIBUTORS.md

* update task_example_test
Co-authored-by: Xin Yao <xiny@nvidia.com>

CONTRIBUTORS.md

@@ -63,3 +63,4 @@ Contributors
 * [Shaked Brody](https://github.com/shakedbr) from Technion
 * [Jiahui Liu](https://github.com/paoxiaode) from Nvidia
 * [Neil Dickson](https://github.com/ndickson-nvidia) from Nvidia
+* [Chang Liu](https://github.com/chang-l) from Nvidia

examples/pytorch/gcn/README.md

@@ -2,33 +2,21 @@ Graph Convolutional Networks (GCN)
 ============
 
 - Paper link: [https://arxiv.org/abs/1609.02907](https://arxiv.org/abs/1609.02907)
-- Author's code repo: [https://github.com/tkipf/gcn](https://github.com/tkipf/gcn). Note that the original code is 
-implemented with Tensorflow for the paper. 
+- Author's code repo: [https://github.com/tkipf/gcn](https://github.com/tkipf/gcn).
 
-Dependencies
-------------
-- PyTorch 0.4.1+
-- requests
-
-``bash
-pip install torch requests
-``
-
-Codes
------
-The folder contains three implementations of GCN:
-- `gcn.py` uses DGL's predefined graph convolution module.
-- `gcn_mp.py` uses user-defined message and reduce functions.
-Modify `train.py` to switch between different implementations.
-
-Results
+How to run
 -------
 
-Run with following (available dataset: "cora", "citeseer", "pubmed")
+### DGL built-in GraphConv module
+
+Run with the following (available dataset: "cora", "citeseer", "pubmed")
 ```bash
-python3 train.py --dataset cora --gpu 0 --self-loop
+python3 train.py --dataset cora
 ```
 
-* cora: ~0.810 (0.79-0.83) (paper: 0.815)
-* citeseer: 0.707 (paper: 0.703)
-* pubmed: 0.792 (paper: 0.790)
+Summary
+-------
+* cora: ~0.810 (paper: 0.815)
+* citeseer: ~0.707 (paper: 0.703)
+* pubmed: ~0.792 (paper: 0.790)
+

examples/pytorch/gcn/gcn.py

-"""GCN using DGL nn package
-
-References:
-- Semi-Supervised Classification with Graph Convolutional Networks
-- Paper: https://arxiv.org/abs/1609.02907
-- Code: https://github.com/tkipf/gcn
-"""
-import torch
-import torch.nn as nn
-from dgl.nn.pytorch import GraphConv
-
-class GCN(nn.Module):
-    def __init__(self,
-                 g,
-                 in_feats,
-                 n_hidden,
-                 n_classes,
-                 n_layers,
-                 activation,
-                 dropout):
-        super(GCN, self).__init__()
-        self.g = g
-        self.layers = nn.ModuleList()
-        # input layer
-        self.layers.append(GraphConv(in_feats, n_hidden, activation=activation))
-        # hidden layers
-        for i in range(n_layers - 1):
-            self.layers.append(GraphConv(n_hidden, n_hidden, activation=activation))
-        # output layer
-        self.layers.append(GraphConv(n_hidden, n_classes))
-        self.dropout = nn.Dropout(p=dropout)
-
-    def forward(self, features):
-        h = features
-        for i, layer in enumerate(self.layers):
-            if i != 0:
-                h = self.dropout(h)
-            h = layer(self.g, h)
-        return h

examples/pytorch/gcn/gcn_mp.py

-"""GCN using basic message passing
-
-References:
-- Semi-Supervised Classification with Graph Convolutional Networks
-- Paper: https://arxiv.org/abs/1609.02907
-- Code: https://github.com/tkipf/gcn
-"""
-import argparse, time, math
-import numpy as np
-import networkx as nx
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import dgl
-from dgl.data import register_data_args
-from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
-
-
-def gcn_msg(edge):
-    msg = edge.src['h'] * edge.src['norm']
-    return {'m': msg}
-
-
-def gcn_reduce(node):
-    accum = torch.sum(node.mailbox['m'], 1) * node.data['norm']
-    return {'h': accum}
-
-
-class NodeApplyModule(nn.Module):
-    def __init__(self, out_feats, activation=None, bias=True):
-        super(NodeApplyModule, self).__init__()
-        if bias:
-            self.bias = nn.Parameter(torch.Tensor(out_feats))
-        else:
-            self.bias = None
-        self.activation = activation
-        self.reset_parameters()
-
-    def reset_parameters(self):
-        if self.bias is not None:
-            stdv = 1. / math.sqrt(self.bias.size(0))
-            self.bias.data.uniform_(-stdv, stdv)
-
-    def forward(self, nodes):
-        h = nodes.data['h']
-        if self.bias is not None:
-            h = h + self.bias
-        if self.activation:
-            h = self.activation(h)
-        return {'h': h}
-
-
-class GCNLayer(nn.Module):
-    def __init__(self,
-                 g,
-                 in_feats,
-                 out_feats,
-                 activation,
-                 dropout,
-                 bias=True):
-        super(GCNLayer, self).__init__()
-        self.g = g
-        self.weight = nn.Parameter(torch.Tensor(in_feats, out_feats))
-        if dropout:
-            self.dropout = nn.Dropout(p=dropout)
-        else:
-            self.dropout = 0.
-        self.node_update = NodeApplyModule(out_feats, activation, bias)
-        self.reset_parameters()
-
-    def reset_parameters(self):
-        stdv = 1. / math.sqrt(self.weight.size(1))
-        self.weight.data.uniform_(-stdv, stdv)
-
-    def forward(self, h):
-        if self.dropout:
-            h = self.dropout(h)
-        self.g.ndata['h'] = torch.mm(h, self.weight)
-        self.g.update_all(gcn_msg, gcn_reduce, self.node_update)
-        h = self.g.ndata.pop('h')
-        return h
-
-class GCN(nn.Module):
-    def __init__(self,
-                 g,
-                 in_feats,
-                 n_hidden,
-                 n_classes,
-                 n_layers,
-                 activation,
-                 dropout):
-        super(GCN, self).__init__()
-        self.layers = nn.ModuleList()
-        # input layer
-        self.layers.append(GCNLayer(g, in_feats, n_hidden, activation, dropout))
-        # hidden layers
-        for i in range(n_layers - 1):
-            self.layers.append(GCNLayer(g, n_hidden, n_hidden, activation, dropout))
-        # output layer
-        self.layers.append(GCNLayer(g, n_hidden, n_classes, None, dropout))
-
-    def forward(self, features):
-        h = features
-        for layer in self.layers:
-            h = layer(h)
-        return h
-
-def evaluate(model, features, labels, mask):
-    model.eval()
-    with torch.no_grad():
-        logits = model(features)
-        logits = logits[mask]
-        labels = labels[mask]
-        _, indices = torch.max(logits, dim=1)
-        correct = torch.sum(indices == labels)
-        return correct.item() * 1.0 / len(labels)
-
-def main(args):
-    # load and preprocess dataset
-    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
-    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 = g.number_of_edges()
-    print("""----Data statistics------'
-      #Edges %d
-      #Classes %d
-      #Train samples %d
-      #Val samples %d
-      #Test samples %d""" %
-          (n_edges, n_classes,
-              train_mask.int().sum().item(),
-              val_mask.int().sum().item(),
-              test_mask.int().sum().item()))
-
-    # add self loop
-    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)
-    norm[torch.isinf(norm)] = 0
-    if cuda:
-        norm = norm.cuda()
-    g.ndata['norm'] = norm.unsqueeze(1)
-
-    # create GCN model
-    model = GCN(g,
-                in_feats,
-                args.n_hidden,
-                n_classes,
-                args.n_layers,
-                F.relu,
-                args.dropout)
-
-    if cuda:
-        model.cuda()
-    loss_fcn = torch.nn.CrossEntropyLoss()
-
-    # use optimizer
-    optimizer = torch.optim.Adam(model.parameters(),
-                                 lr=args.lr,
-                                 weight_decay=args.weight_decay)
-
-    # initialize graph
-    dur = []
-    for epoch in range(args.n_epochs):
-        model.train()
-        if epoch >= 3:
-            t0 = time.time()
-        # forward
-        logits = model(features)
-        loss = loss_fcn(logits[train_mask], labels[train_mask])
-
-        optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
-
-        if epoch >= 3:
-            dur.append(time.time() - t0)
-
-        acc = evaluate(model, features, labels, val_mask)
-        print("Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | Accuracy {:.4f} | "
-              "ETputs(KTEPS) {:.2f}". format(epoch, np.mean(dur), loss.item(),
-                                             acc, n_edges / np.mean(dur) / 1000))
-
-    print()
-    acc = evaluate(model, features, labels, test_mask)
-    print("Test Accuracy {:.4f}".format(acc))
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='GCN')
-    register_data_args(parser)
-    parser.add_argument("--dropout", type=float, default=0.5,
-            help="dropout probability")
-    parser.add_argument("--gpu", type=int, default=-1,
-            help="gpu")
-    parser.add_argument("--lr", type=float, default=1e-2,
-            help="learning rate")
-    parser.add_argument("--n-epochs", type=int, default=200,
-            help="number of training epochs")
-    parser.add_argument("--n-hidden", type=int, default=16,
-            help="number of hidden gcn units")
-    parser.add_argument("--n-layers", type=int, default=1,
-            help="number of hidden gcn layers")
-    parser.add_argument("--weight-decay", type=float, default=5e-4,
-            help="Weight for L2 loss")
-    args = parser.parse_args()
-    print(args)
-
-    main(args)

examples/pytorch/gcn/train.py

-import argparse
-import time
-import numpy as np
 import torch
+import torch.nn as nn
 import torch.nn.functional as F
 import dgl
+import dgl.nn as dglnn
 from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
+from dgl import AddSelfLoop
+import argparse
 
-from gcn import GCN
-#from gcn_mp import GCN
-#from gcn_spmv import GCN
-
-
-def evaluate(model, features, labels, mask):
+class GCN(nn.Module):
+    def __init__(self, in_size, hid_size, out_size):
+        super().__init__()
+        self.layers = nn.ModuleList()
+        # two-layer GCN
+        self.layers.append(dglnn.GraphConv(in_size, hid_size, activation=F.relu))
+        self.layers.append(dglnn.GraphConv(hid_size, out_size))
+        self.dropout = nn.Dropout(0.5)
+
+    def forward(self, g, features):
+        h = features
+        for i, layer in enumerate(self.layers):
+            if i != 0:
+                h = self.dropout(h)
+            h = layer(g, h)
+        return h
+    
+def evaluate(g, features, labels, mask, model):
     model.eval()
     with torch.no_grad():
-        logits = model(features)
+        logits = model(g, features)
         logits = logits[mask]
         labels = labels[mask]
         _, indices = torch.max(logits, dim=1)
@@ -22,124 +35,65 @@ def evaluate(model, features, labels, mask):
         return correct.item() * 1.0 / len(labels)
 
 
-def main(args):
+def train(g, features, labels, masks, model):
+    # define train/val samples, loss function and optimizer
+    train_mask = masks[0]
+    val_mask = masks[1]
+    loss_fcn = nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-2, weight_decay=5e-4)
+
+    # training loop
+    for epoch in range(200):
+        model.train()
+        logits = model(g, features)
+        loss = loss_fcn(logits[train_mask], labels[train_mask])
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        acc = evaluate(g, features, labels, val_mask, model)
+        print("Epoch {:05d} | Loss {:.4f} | Accuracy {:.4f} "
+              . format(epoch, loss.item(), acc))
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset", type=str, default="cora",
+                        help="Dataset name ('cora', 'citeseer', 'pubmed').")
+    args = parser.parse_args()
+    print(f'Training with DGL built-in GraphConv module.')
+ 
     # load and preprocess dataset
+    transform = AddSelfLoop()  # by default, it will first remove self-loops to prevent duplication
     if args.dataset == 'cora':
-        data = CoraGraphDataset()
+        data = CoraGraphDataset(transform=transform)
     elif args.dataset == 'citeseer':
-        data = CiteseerGraphDataset()
+        data = CiteseerGraphDataset(transform=transform)
     elif args.dataset == 'pubmed':
-        data = PubmedGraphDataset()
+        data = PubmedGraphDataset(transform=transform)
     else:
         raise ValueError('Unknown dataset: {}'.format(args.dataset))
-
     g = data[0]
-    if args.gpu < 0:
-        cuda = False
-    else:
-        cuda = True
-        g = g.int().to(args.gpu)
-
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    g = g.int().to(device)
     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 = g.number_of_edges()
-    print("""----Data statistics------'
-      #Edges %d
-      #Classes %d
-      #Train samples %d
-      #Val samples %d
-      #Test samples %d""" %
-          (n_edges, n_classes,
-              train_mask.int().sum().item(),
-              val_mask.int().sum().item(),
-              test_mask.int().sum().item()))
-
-    # 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()
+    masks = g.ndata['train_mask'], g.ndata['val_mask'], g.ndata['test_mask']
         
     # normalization
     degs = g.in_degrees().float()
-    norm = torch.pow(degs, -0.5)
+    norm = torch.pow(degs, -0.5).to(device)
     norm[torch.isinf(norm)] = 0
-    if cuda:
-        norm = norm.cuda()
     g.ndata['norm'] = norm.unsqueeze(1)
 
     # create GCN model    
-    model = GCN(g,
-                in_feats,
-                args.n_hidden,
-                n_classes,
-                args.n_layers,
-                F.relu,
-                args.dropout)
-
-    if cuda:
-        model.cuda()
-    loss_fcn = torch.nn.CrossEntropyLoss()
-
-    # use optimizer
-    optimizer = torch.optim.Adam(model.parameters(),
-                                 lr=args.lr,
-                                 weight_decay=args.weight_decay)
-
-    # initialize graph
-    dur = []
-    for epoch in range(args.n_epochs):
-        model.train()
-        if epoch >= 3:
-            t0 = time.time()
-        # forward
-        logits = model(features)
-        loss = loss_fcn(logits[train_mask], labels[train_mask])
-
-        optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
-
-        if epoch >= 3:
-            dur.append(time.time() - t0)
-
-        acc = evaluate(model, features, labels, val_mask)
-        print("Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | Accuracy {:.4f} | "
-              "ETputs(KTEPS) {:.2f}". format(epoch, np.mean(dur), loss.item(),
-                                             acc, n_edges / np.mean(dur) / 1000))
-
-    print()
-    acc = evaluate(model, features, labels, test_mask)
-    print("Test accuracy {:.2%}".format(acc))
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='GCN')
-    parser.add_argument("--dataset", type=str, default="cora",
-                        help="Dataset name ('cora', 'citeseer', 'pubmed').")
-    parser.add_argument("--dropout", type=float, default=0.5,
-                        help="dropout probability")
-    parser.add_argument("--gpu", type=int, default=-1,
-                        help="gpu")
-    parser.add_argument("--lr", type=float, default=1e-2,
-                        help="learning rate")
-    parser.add_argument("--n-epochs", type=int, default=200,
-                        help="number of training epochs")
-    parser.add_argument("--n-hidden", type=int, default=16,
-                        help="number of hidden gcn units")
-    parser.add_argument("--n-layers", type=int, default=1,
-                        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)")
-    parser.set_defaults(self_loop=False)
-    args = parser.parse_args()
-    print(args)
-
-    main(args)
+    in_size = features.shape[1]
+    out_size = data.num_classes
+    model = GCN(in_size, 16, out_size).to(device)
+
+    # model training
+    print('Training...')
+    train(g, features, labels, masks, model)
+    
+    # test the model
+    print('Testing...')
+    acc = evaluate(g, features, labels, masks[2], model)
+    print("Test accuracy {:.4f}".format(acc))

tests/scripts/task_example_test.bat

@@ -24,7 +24,7 @@ SET DGL_DOWNLOAD_DIR=!CD!
 
 PUSHD !GCN_EXAMPLE_DIR!
 python pagerank.py || GOTO :FAIL
-python gcn\gcn.py --dataset cora --gpu !DEV! || GOTO :FAIL
+python gcn\train.py --dataset cora || GOTO :FAIL
 POPD
 ENDLOCAL
 EXIT /B

tests/scripts/task_example_test.sh

@@ -39,7 +39,7 @@ export DGL_DOWNLOAD_DIR=${PWD}
 pushd $GCN_EXAMPLE_DIR> /dev/null
 
 python3 pagerank.py || fail "run pagerank.py on $1"
-python3 gcn/gcn.py --dataset cora --gpu $dev || fail "run gcn/gcn.py on $1"
+python3 gcn/train.py --dataset cora || fail "run gcn/train.py on $1"
 python3 lda/lda_model.py || fail "run lda/lda_model.py on $1"
 
 popd > /dev/null