Skip to content

GitLab

Commit fde4f581 authored by Minjie Wang's avatar Minjie Wang
Browse files

gcn runnable

parent 2be55fb5
Hide whitespace changes
Inline Side-by-side
Showing with 36 additions and 212 deletions
examples/pytorch/gcn/README.md
View file @ fde4f581
...@@ -4,43 +4,9 @@ Graph Convolutional Networks (GCN) ...@@ -4,43 +4,9 @@ Graph Convolutional Networks (GCN)
Paper link: [https://arxiv.org/abs/1609.02907](https://arxiv.org/abs/1609.02907) 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) Author's code repo: [https://github.com/tkipf/gcn](https://github.com/tkipf/gcn)
The folder contains three different implementations using DGL. The folder contains two different implementations using DGL.
Naive GCN (gcn.py) Batched GCN (gcn.py)
-------
The model is defined in the finest granularity (aka on *one* edge and *one* node).
* The message function `gcn_msg` computes the message for one edge. It simply returns the `h` representation of the source node.
```python
def gcn_msg(src, edge):
# src['h'] is a tensor of shape (D,). D is the feature length.
return src['h']
```
* The reduce function `gcn_reduce` accumulates the incoming messages for one node. The `msgs` argument is a list of all the messages. In GCN, the incoming messages are summed up.
```python
def gcn_reduce(node, msgs):
# msgs is a list of in-coming messages.
return sum(msgs)
```
* The update function `NodeUpdateModule` computes the new new node representation `h` using non-linear transformation on the reduced messages.
```python
class NodeUpdateModule(nn.Module):
def __init__(self, in_feats, out_feats, activation=None):
super(NodeUpdateModule, self).__init__()
self.linear = nn.Linear(in_feats, out_feats)
self.activation = activation
def forward(self, node, accum):
# accum is a tensor of shape (D,).
h = self.linear(accum)
if self.activation:
h = self.activation(h)
return {'h' : h}
```
After defining the functions on each node/edge, the message passing is triggered by calling `update_all` on the DGLGraph object (in GCN module).
Batched GCN (gcn_batch.py)
----------- -----------
Defining the model on only one node and edge makes it hard to fully utilize GPUs. As a result, we allow users to define model on a *batch of* nodes and edges. Defining the model on only one node and edge makes it hard to fully utilize GPUs. As a result, we allow users to define model on a *batch of* nodes and edges.
......
examples/pytorch/gcn/gcn.py
View file @ fde4f581
...@@ -2,6 +2,8 @@ ...@@ -2,6 +2,8 @@
Semi-Supervised Classification with Graph Convolutional Networks Semi-Supervised Classification with Graph Convolutional Networks
Paper: https://arxiv.org/abs/1609.02907 Paper: https://arxiv.org/abs/1609.02907
Code: https://github.com/tkipf/gcn Code: https://github.com/tkipf/gcn
GCN with batch processing
""" """
import argparse import argparse
import numpy as np import numpy as np
...@@ -9,14 +11,15 @@ import time ...@@ -9,14 +11,15 @@ import time
import torch import torch
import torch.nn as nn import torch.nn as nn
import torch.nn.functional as F import torch.nn.functional as F
import dgl
from dgl import DGLGraph from dgl import DGLGraph
from dgl.data import register_data_args, load_data from dgl.data import register_data_args, load_data
def gcn_msg(src, edge): def gcn_msg(src, edge):
return src['h'] return src
def gcn_reduce(node, msgs): def gcn_reduce(node, msgs):
return {'h' : sum(msgs)} return torch.sum(msgs, 1)
class NodeApplyModule(nn.Module): class NodeApplyModule(nn.Module):
def __init__(self, in_feats, out_feats, activation=None): def __init__(self, in_feats, out_feats, activation=None):
...@@ -25,14 +28,14 @@ class NodeApplyModule(nn.Module): ...@@ -25,14 +28,14 @@ class NodeApplyModule(nn.Module):
self.activation = activation self.activation = activation
def forward(self, node): def forward(self, node):
h = self.linear(node['h']) h = self.linear(node)
if self.activation: if self.activation:
h = self.activation(h) h = self.activation(h)
return {'h' : h} return h
class GCN(nn.Module): class GCN(nn.Module):
def __init__(self, def __init__(self,
nx_graph, g,
in_feats, in_feats,
n_hidden, n_hidden,
n_classes, n_classes,
...@@ -40,7 +43,7 @@ class GCN(nn.Module): ...@@ -40,7 +43,7 @@ class GCN(nn.Module):
activation, activation,
dropout): dropout):
super(GCN, self).__init__() super(GCN, self).__init__()
self.g = DGLGraph(nx_graph) self.g = g
self.dropout = dropout self.dropout = dropout
# input layer # input layer
self.layers = nn.ModuleList([NodeApplyModule(in_feats, n_hidden, activation)]) self.layers = nn.ModuleList([NodeApplyModule(in_feats, n_hidden, activation)])
...@@ -50,31 +53,24 @@ class GCN(nn.Module): ...@@ -50,31 +53,24 @@ class GCN(nn.Module):
# output layer # output layer
self.layers.append(NodeApplyModule(n_hidden, n_classes)) self.layers.append(NodeApplyModule(n_hidden, n_classes))
def forward(self, features, train_nodes): def forward(self, features):
for n, feat in features.items(): self.g.set_n_repr(features)
self.g.nodes[n]['h'] = feat
for layer in self.layers: for layer in self.layers:
# apply dropout # apply dropout
if self.dropout: if self.dropout:
self.g.nodes[n]['h'] = F.dropout(g.nodes[n]['h'], p=self.dropout) val = F.dropout(self.g.get_n_repr(), p=self.dropout)
self.g.set_n_repr(val)
self.g.update_all(gcn_msg, gcn_reduce, layer) self.g.update_all(gcn_msg, gcn_reduce, layer)
return torch.cat([torch.unsqueeze(self.g.nodes[n]['h'], 0) for n in train_nodes]) return self.g.pop_n_repr()
def main(args): def main(args):
# load and preprocess dataset # load and preprocess dataset
data = load_data(args) data = load_data(args)
# features of each samples features = torch.FloatTensor(data.features)
features = {} labels = torch.LongTensor(data.labels)
labels = [] mask = torch.ByteTensor(data.train_mask)
train_nodes = [] in_feats = features.shape[1]
for n in data.graph.nodes():
features[n] = torch.FloatTensor(data.features[n, :])
if data.train_mask[n] == 1:
train_nodes.append(n)
labels.append(data.labels[n])
labels = torch.LongTensor(labels)
in_feats = data.features.shape[1]
n_classes = data.num_labels n_classes = data.num_labels
n_edges = data.graph.number_of_edges() n_edges = data.graph.number_of_edges()
...@@ -83,11 +79,13 @@ def main(args): ...@@ -83,11 +79,13 @@ def main(args):
else: else:
cuda = True cuda = True
torch.cuda.set_device(args.gpu) torch.cuda.set_device(args.gpu)
features = {k : v.cuda() for k, v in features.items()} features = features.cuda()
labels = labels.cuda() labels = labels.cuda()
mask = mask.cuda()
# create GCN model # create GCN model
model = GCN(data.graph, g = DGLGraph(data.graph)
model = GCN(g,
in_feats, in_feats,
args.n_hidden, args.n_hidden,
n_classes, n_classes,
...@@ -107,9 +105,9 @@ def main(args): ...@@ -107,9 +105,9 @@ def main(args):
if epoch >= 3: if epoch >= 3:
t0 = time.time() t0 = time.time()
# forward # forward
logits = model(features, train_nodes) logits = model(features)
logp = F.log_softmax(logits, 1) logp = F.log_softmax(logits, 1)
loss = F.nll_loss(logp, labels) loss = F.nll_loss(logp[mask], labels[mask])
optimizer.zero_grad() optimizer.zero_grad()
loss.backward() loss.backward()
...@@ -130,7 +128,7 @@ if __name__ == '__main__': ...@@ -130,7 +128,7 @@ if __name__ == '__main__':
help="gpu") help="gpu")
parser.add_argument("--lr", type=float, default=1e-3, parser.add_argument("--lr", type=float, default=1e-3,
help="learning rate") help="learning rate")
parser.add_argument("--n-epochs", type=int, default=10, parser.add_argument("--n-epochs", type=int, default=20,
help="number of training epochs") help="number of training epochs")
parser.add_argument("--n-hidden", type=int, default=16, parser.add_argument("--n-hidden", type=int, default=16,
help="number of hidden gcn units") help="number of hidden gcn units")
......
examples/pytorch/gcn/gcn_batch.py deleted 100644 → 0
View file @ 2be55fb5
"""
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 torch
import torch.nn as nn
import torch.nn.functional as F
import dgl
from dgl import DGLGraph
from dgl.data import register_data_args, load_data
def gcn_msg(src, edge):
return src
def gcn_reduce(node, msgs):
return torch.sum(msgs, 1)
class NodeApplyModule(nn.Module):
def __init__(self, in_feats, out_feats, activation=None):
super(NodeApplyModule, self).__init__()
self.linear = nn.Linear(in_feats, out_feats)
self.activation = activation
def forward(self, node):
h = self.linear(node)
if self.activation:
h = self.activation(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.g = g
self.dropout = dropout
# input layer
self.layers = nn.ModuleList([NodeApplyModule(in_feats, n_hidden, activation)])
# hidden layers
for i in range(n_layers - 1):
self.layers.append(NodeApplyModule(n_hidden, n_hidden, activation))
# output layer
self.layers.append(NodeApplyModule(n_hidden, n_classes))
def forward(self, features):
self.g.set_n_repr(features)
for layer in self.layers:
# apply dropout
if self.dropout:
val = F.dropout(self.g.get_n_repr(), p=self.dropout)
self.g.set_n_repr(val)
self.g.update_all(gcn_msg, gcn_reduce, layer, batchable=True)
return self.g.pop_n_repr()
def main(args):
# load and preprocess dataset
data = load_data(args)
features = torch.FloatTensor(data.features)
labels = torch.LongTensor(data.labels)
mask = torch.ByteTensor(data.train_mask)
in_feats = features.shape[1]
n_classes = data.num_labels
n_edges = data.graph.number_of_edges()
if args.gpu < 0:
cuda = False
else:
cuda = True
torch.cuda.set_device(args.gpu)
features = features.cuda()
labels = labels.cuda()
mask = mask.cuda()
# create GCN model
g = DGLGraph(data.graph)
model = GCN(g,
in_feats,
args.n_hidden,
n_classes,
args.n_layers,
F.relu,
args.dropout)
if cuda:
model.cuda()
# use optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
# initialize graph
dur = []
for epoch in range(args.n_epochs):
if epoch >= 3:
t0 = time.time()
# forward
logits = model(features)
logp = F.log_softmax(logits, 1)
loss = F.nll_loss(logp[mask], labels[mask])
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch >= 3:
dur.append(time.time() - t0)
print("Epoch {:05d} | Loss {:.4f} | Time(s) {:.4f} | ETputs(KTEPS) {:.2f}".format(
epoch, loss.item(), np.mean(dur), n_edges / np.mean(dur) / 1000))
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='GCN')
register_data_args(parser)
parser.add_argument("--dropout", type=float, default=0,
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")
parser.add_argument("--n-layers", type=int, default=1,
help="number of hidden gcn layers")
args = parser.parse_args()
print(args)
main(args)
examples/pytorch/gcn/gcn_spmv.py
View file @ fde4f581
...@@ -55,7 +55,7 @@ class GCN(nn.Module): ...@@ -55,7 +55,7 @@ class GCN(nn.Module):
if self.dropout: if self.dropout:
val = F.dropout(self.g.get_n_repr(), p=self.dropout) val = F.dropout(self.g.get_n_repr(), p=self.dropout)
self.g.set_n_repr(val) self.g.set_n_repr(val)
self.g.update_all(fn.copy_src(), fn.sum(), layer, batchable=True) self.g.update_all(fn.copy_src(), fn.sum(), layer)
return self.g.pop_n_repr() return self.g.pop_n_repr()
def main(args): def main(args):
......
include/dgl/graph_op.h
View file @ fde4f581
...@@ -10,10 +10,16 @@ class GraphOp { ...@@ -10,10 +10,16 @@ class GraphOp {
public: public:
/*! /*!
* \brief Return the line graph. * \brief Return the line graph.
*
* If i~j and j~i are two edges in original graph G, then
* (i,j)~(j,i) and (j,i)~(i,j) are the "backtracking" edges on
* the line graph.
*
* \param graph The input graph. * \param graph The input graph.
* \param backtracking Whether the backtracking edges are included or not
* \return the line graph * \return the line graph
*/ */
static Graph LineGraph(const Graph* graph); static Graph LineGraph(const Graph* graph, bool backtracking);
/*! /*!
* \brief Return a disjoint union of the input graphs. * \brief Return a disjoint union of the input graphs.
* *
......
python/dgl/batch.py
View file @ fde4f581
...@@ -143,12 +143,10 @@ def unbatch(graph): ...@@ -143,12 +143,10 @@ def unbatch(graph):
node_frames = [FrameRef() for i in range(bsize)] node_frames = [FrameRef() for i in range(bsize)]
edge_frames = [FrameRef() for i in range(bsize)] edge_frames = [FrameRef() for i in range(bsize)]
for attr, col in graph._node_frame.items(): for attr, col in graph._node_frame.items():
# TODO: device context
col_splits = F.unpack(col, bn) col_splits = F.unpack(col, bn)
for i in range(bsize): for i in range(bsize):
node_frames[i][attr] = col_splits[i] node_frames[i][attr] = col_splits[i]
for attr, col in graph._edge_frame.items(): for attr, col in graph._edge_frame.items():
# TODO: device context
col_splits = F.unpack(col, be) col_splits = F.unpack(col, be)
for i in range(bsize): for i in range(bsize):
edge_frames[i][attr] = col_splits[i] edge_frames[i][attr] = col_splits[i]
......
python/dgl/graph.py
View file @ fde4f581
...@@ -31,15 +31,11 @@ class DGLGraph(object): ...@@ -31,15 +31,11 @@ class DGLGraph(object):
Node feature storage. Node feature storage.
edge_frame : FrameRef edge_frame : FrameRef
Edge feature storage. Edge feature storage.
attr : keyword arguments, optional
Attributes to add to graph as key=value pairs.
""" """
def __init__(self, def __init__(self,
graph_data=None, graph_data=None,
node_frame=None, node_frame=None,
edge_frame=None, edge_frame=None):
**attr):
# TODO: keyword attr
# graph # graph
self._graph = create_graph_index(graph_data) self._graph = create_graph_index(graph_data)
# frame # frame
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment