gcn_batch.py 6.01 KB
Newer Older
Da Zheng's avatar
Da Zheng committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
"""
Semi-Supervised Classification with Graph Convolutional Networks
Paper: https://arxiv.org/abs/1609.02907
Code: https://github.com/tkipf/gcn

GCN with batch processing
"""
import argparse
import numpy as np
import time
import mxnet as mx
from mxnet import gluon
import dgl
from dgl import DGLGraph
from dgl.data import register_data_args, load_data
16
from functools import partial
Da Zheng's avatar
Da Zheng committed
17

18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
def gcn_msg(edge, normalization=None):
    # print('h', edge.src['h'].shape, edge.src['out_degree'])
    msg = edge.src['h']
    if normalization == 'sym':
        msg = msg / edge.src['out_degree'].sqrt().reshape((-1,1))
    return {'m': msg}


def gcn_reduce(node, normalization=None):
    # print('m', node.mailbox['m'].shape, node.data['in_degree'])
    accum = mx.nd.sum(node.mailbox['m'], 1)
    if normalization == 'sym':
        accum = accum / node.data['in_degree'].sqrt().reshape((-1,1))
    elif normalization == 'left':
        accum = accum / node.data['in_degree'].reshape((-1,1))
    return {'accum': accum}
Da Zheng's avatar
Da Zheng committed
34
35
36


class NodeUpdateModule(gluon.Block):
37
    def __init__(self, out_feats, activation=None, dropout=0):
Da Zheng's avatar
Da Zheng committed
38
39
        super(NodeUpdateModule, self).__init__()
        self.linear = gluon.nn.Dense(out_feats, activation=activation)
40
        self.dropout = dropout
Da Zheng's avatar
Da Zheng committed
41
42

    def forward(self, node):
43
44
45
46
47
        accum = self.linear(node.data['accum'])
        if self.dropout:
            accum = mx.nd.Dropout(accum, p=self.dropout)
        return {'h': mx.nd.concat(node.data['h'], accum, dim=1)}

Da Zheng's avatar
Da Zheng committed
48
49
50
51
52
53
54
55
56

class GCN(gluon.Block):
    def __init__(self,
                 g,
                 in_feats,
                 n_hidden,
                 n_classes,
                 n_layers,
                 activation,
57
58
59
                 dropout,
                 normalization,
                 ):
Da Zheng's avatar
Da Zheng committed
60
61
62
        super(GCN, self).__init__()
        self.g = g
        self.dropout = dropout
63
64
65
66
67
68
69
70
71
72
73
74

        self.inp_layer = gluon.nn.Dense(n_hidden, activation)

        self.conv_layers = gluon.nn.Sequential()
        for i in range(n_layers):
            self.conv_layers.add(NodeUpdateModule(n_hidden, activation, dropout))

        self.out_layer = gluon.nn.Dense(n_classes)

        self.gcn_msg = partial(gcn_msg, normalization=normalization)
        self.gcn_reduce = partial(gcn_reduce, normalization=normalization)

Da Zheng's avatar
Da Zheng committed
75
76

    def forward(self, features):
77
78
79
80
81
82
83
84
85
86
87
        emb_inp = [features, self.inp_layer(features)]
        if self.dropout:
            emb_inp[-1] = mx.nd.Dropout(emb_inp[-1], p=self.dropout)

        self.g.ndata['h'] = mx.nd.concat(*emb_inp, dim=1)
        for layer in self.conv_layers:
            self.g.update_all(self.gcn_msg, self.gcn_reduce, layer)

        emb_out = self.g.ndata.pop('h')
        return self.out_layer(emb_out)

Da Zheng's avatar
Da Zheng committed
88
89
90
91
92

def main(args):
    # load and preprocess dataset
    data = load_data(args)

93
94
95
    if args.self_loop:
        data.graph.add_edges_from([(i,i) for i in range(len(data.graph))])

Da Zheng's avatar
Da Zheng committed
96
97
98
    features = mx.nd.array(data.features)
    labels = mx.nd.array(data.labels)
    mask = mx.nd.array(data.train_mask)
99
100
101
102
103
    in_degree  = mx.nd.array([data.graph.in_degree(i)
                              for i in range(len(data.graph))])
    out_degree = mx.nd.array([data.graph.out_degree(i)
                              for i in range(len(data.graph))])

Da Zheng's avatar
Da Zheng committed
104
105
106
107
108
109
110
111
112
113
114
115
    in_feats = features.shape[1]
    n_classes = data.num_labels
    n_edges = data.graph.number_of_edges()

    if args.gpu <= 0:
        cuda = False
        ctx = mx.cpu(0)
    else:
        cuda = True
        features = features.as_in_context(mx.gpu(0))
        labels = labels.as_in_context(mx.gpu(0))
        mask = mask.as_in_context(mx.gpu(0))
116
117
        in_degree = in_degree.as_in_context(mx.gpu(0))
        out_degree = out_degree.as_in_context(mx.gpu(0))
Da Zheng's avatar
Da Zheng committed
118
119
120
121
        ctx = mx.gpu(0)

    # create GCN model
    g = DGLGraph(data.graph)
122
123
124
    g.ndata['in_degree']  = in_degree
    g.ndata['out_degree'] = out_degree

Da Zheng's avatar
Da Zheng committed
125
126
127
128
129
130
    model = GCN(g,
                in_feats,
                args.n_hidden,
                n_classes,
                args.n_layers,
                'relu',
131
132
133
                args.dropout,
                args.normalization,
                )
Da Zheng's avatar
Da Zheng committed
134
    model.initialize(ctx=ctx)
135
    loss_fcn = gluon.loss.SoftmaxCELoss()
Da Zheng's avatar
Da Zheng committed
136
137
138
139
140
141
142
143
144
145
146

    # use optimizer
    trainer = gluon.Trainer(model.collect_params(), 'adam', {'learning_rate': args.lr})

    # initialize graph
    dur = []
    for epoch in range(args.n_epochs):
        if epoch >= 3:
            t0 = time.time()
        # forward
        with mx.autograd.record():
147
148
            pred = model(features)
            loss = loss_fcn(pred, labels, mask)
Da Zheng's avatar
Da Zheng committed
149
150
151
152
153
154
155
156
157
158

        #optimizer.zero_grad()
        loss.backward()
        trainer.step(features.shape[0])

        if epoch >= 3:
            dur.append(time.time() - t0)
            print("Epoch {:05d} | Loss {:.4f} | Time(s) {:.4f} | ETputs(KTEPS) {:.2f}".format(
                epoch, loss.asnumpy()[0], np.mean(dur), n_edges / np.mean(dur) / 1000))

159
160
161
162
    # test set accuracy
    pred = model(features)
    accuracy = (pred*100).softmax().pick(labels).mean()
    print("Final accuracy {:.2%}".format(accuracy.mean().asscalar()))
163
    return accuracy.mean().asscalar()
164
165


Da Zheng's avatar
Da Zheng committed
166
167
168
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='GCN')
    register_data_args(parser)
169
    parser.add_argument("--dropout", type=float, default=0.5,
Da Zheng's avatar
Da Zheng committed
170
171
172
173
174
175
176
177
178
            help="dropout probability")
    parser.add_argument("--gpu", type=int, default=-1,
            help="gpu")
    parser.add_argument("--lr", type=float, default=1e-3,
            help="learning rate")
    parser.add_argument("--n-epochs", type=int, default=20,
            help="number of training epochs")
    parser.add_argument("--n-hidden", type=int, default=16,
            help="number of hidden gcn units")
179
    parser.add_argument("--n-layers", type=int, default=2,
Da Zheng's avatar
Da Zheng committed
180
            help="number of hidden gcn layers")
181
182
183
184
185
    parser.add_argument("--normalization",
            choices=['sym','left'], default=None,
            help="graph normalization types (default=None)")
    parser.add_argument("--self-loop", action='store_true',
            help="graph self-loop (default=False)")
Da Zheng's avatar
Da Zheng committed
186
187
    args = parser.parse_args()

188
189
    print(args)

Da Zheng's avatar
Da Zheng committed
190
    main(args)