-
Notifications
You must be signed in to change notification settings - Fork 4
/
model.py
61 lines (47 loc) · 1.52 KB
/
model.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
# -*- coding: utf-8 -*-
import torch
import torch.nn as nn
def gcn_message(edges):
"""
compute a batch of message called 'msg' using the source nodes' feature 'h'
"""
return {'msg': edges.src['h']}
def gcn_reduce(nodes):
"""
compute the new 'h' features by summing received 'msg' in each node's mailbox.
"""
return {'h': torch.sum(nodes.mailbox['msg'], dim=1)}
class GCNLayer(nn.Module):
"""
Define the GCNLayer module.
"""
def __init__(self, in_feats, out_feats):
super(GCNLayer, self).__init__()
self.linear = nn.Linear(in_feats, out_feats)
def forward(self, g, inputs):
# g is the graph and the inputs is the input node features
# first set the node features
g.ndata['h'] = inputs
# trigger message passing on all edges
g.send(g.edges(), gcn_message)
# trigger aggregation at all nodes
g.recv(g.nodes(), gcn_reduce)
# get the result node features
h = g.ndata.pop('h')
# perform linear transformation
return self.linear(h)
class GCN(nn.Module):
"""
Define a 2-layer GCN model.
"""
def __init__(self, in_feats, hidden_size, num_classes):
super(GCN, self).__init__()
self.gcn1 = GCNLayer(in_feats, hidden_size)
self.gcn2 = GCNLayer(hidden_size, num_classes)
def forward(self, g, inputs):
h = self.gcn1(g, inputs)
h = torch.relu(h)
h = self.gcn2(g, h)
return h
if __name__ == '__main__':
net = GCN(34, 5, 2)