gcn.py 1.1 KB
Newer Older
1
2
3
4
5
6
"""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
Ziyue Huang's avatar
Ziyue Huang committed
7
8
9
"""
import mxnet as mx
from mxnet import gluon
10

Ziyue Huang's avatar
Ziyue Huang committed
11
import dgl
12
from dgl.nn.mxnet import GraphConv
Ziyue Huang's avatar
Ziyue Huang committed
13

14

Ziyue Huang's avatar
Ziyue Huang committed
15
class GCN(gluon.Block):
16
17
18
    def __init__(
        self, g, in_feats, n_hidden, n_classes, n_layers, activation, dropout
    ):
Ziyue Huang's avatar
Ziyue Huang committed
19
        super(GCN, self).__init__()
20
        self.g = g
Ziyue Huang's avatar
Ziyue Huang committed
21
22
        self.layers = gluon.nn.Sequential()
        # input layer
23
        self.layers.add(GraphConv(in_feats, n_hidden, activation=activation))
Ziyue Huang's avatar
Ziyue Huang committed
24
25
        # hidden layers
        for i in range(n_layers - 1):
26
27
28
            self.layers.add(
                GraphConv(n_hidden, n_hidden, activation=activation)
            )
Ziyue Huang's avatar
Ziyue Huang committed
29
        # output layer
30
31
        self.layers.add(GraphConv(n_hidden, n_classes))
        self.dropout = gluon.nn.Dropout(rate=dropout)
Ziyue Huang's avatar
Ziyue Huang committed
32
33
34

    def forward(self, features):
        h = features
35
36
37
        for i, layer in enumerate(self.layers):
            if i != 0:
                h = self.dropout(h)
38
            h = layer(self.g, h)
Ziyue Huang's avatar
Ziyue Huang committed
39
        return h