1_first.py 11.3 KB
Newer Older
Minjie Wang's avatar
Minjie Wang committed
1
"""
Minjie Wang's avatar
Minjie Wang committed
2
3
.. currentmodule:: dgl

4
DGL at a Glance
Minjie Wang's avatar
Minjie Wang committed
5
6
=========================

Minjie Wang's avatar
Minjie Wang committed
7
8
**Author**: `Minjie Wang <https://jermainewang.github.io/>`_, Quan Gan, `Jake
Zhao <https://cs.nyu.edu/~jakezhao/>`_, Zheng Zhang
9

Minjie Wang's avatar
Minjie Wang committed
10
11
12
13
DGL is a Python package dedicated to deep learning on graphs, built atop
existing tensor DL frameworks (e.g. Pytorch, MXNet) and simplifying the
implementation of graph-based neural networks.

Minjie Wang's avatar
Minjie Wang committed
14
The goal of this tutorial:
15

Minjie Wang's avatar
Minjie Wang committed
16
- Understand how DGL enables computation on graph from a high level.
17
- Train a simple graph neural network in DGL to classify nodes in a graph.
Minjie Wang's avatar
Minjie Wang committed
18
19

At the end of this tutorial, we hope you get a brief feeling of how DGL works.
Minjie Wang's avatar
Minjie Wang committed
20

Minjie Wang's avatar
Minjie Wang committed
21
22
*This tutorial assumes basic familiarity with pytorch.*
"""
Minjie Wang's avatar
Minjie Wang committed
23

Minjie Wang's avatar
Minjie Wang committed
24
###############################################################################
Minjie Wang's avatar
Minjie Wang committed
25
26
# Step 0: Problem description
# ---------------------------
27
#
Minjie Wang's avatar
Minjie Wang committed
28
29
30
# We start with the well-known "Zachary's karate club" problem. The karate club
# is a social network which captures 34 members and document pairwise links
# between members who interact outside the club.  The club later divides into
31
# two communities led by the instructor (node 0) and the club president (node
Minjie Wang's avatar
Minjie Wang committed
32
33
# 33). The network is visualized as follows with the color indicating the
# community:
34
#
Minjie Wang's avatar
Minjie Wang committed
35
# .. image:: https://s3.us-east-2.amazonaws.com/dgl.ai/tutorial/img/karate-club.png
36
#    :align: center
Minjie Wang's avatar
Minjie Wang committed
37
#
Minjie Wang's avatar
Minjie Wang committed
38
39
# The task is to predict which side (0 or 33) each member tends to join given
# the social network itself.
40
41
42


###############################################################################
Minjie Wang's avatar
Minjie Wang committed
43
44
45
# Step 1: Creating a graph in DGL
# -------------------------------
# Creating the graph for Zachary's karate club goes as follows:
Minjie Wang's avatar
Minjie Wang committed
46
47

import dgl
48

49
def build_karate_club_graph():
50
    g = dgl.DGLGraph()
51
52
    # add 34 nodes into the graph; nodes are labeled from 0~33
    g.add_nodes(34)
Minjie Wang's avatar
Minjie Wang committed
53
    # all 78 edges as a list of tuples
54
55
56
57
58
59
60
61
62
63
64
65
66
    edge_list = [(1, 0), (2, 0), (2, 1), (3, 0), (3, 1), (3, 2),
        (4, 0), (5, 0), (6, 0), (6, 4), (6, 5), (7, 0), (7, 1),
        (7, 2), (7, 3), (8, 0), (8, 2), (9, 2), (10, 0), (10, 4),
        (10, 5), (11, 0), (12, 0), (12, 3), (13, 0), (13, 1), (13, 2),
        (13, 3), (16, 5), (16, 6), (17, 0), (17, 1), (19, 0), (19, 1),
        (21, 0), (21, 1), (25, 23), (25, 24), (27, 2), (27, 23),
        (27, 24), (28, 2), (29, 23), (29, 26), (30, 1), (30, 8),
        (31, 0), (31, 24), (31, 25), (31, 28), (32, 2), (32, 8),
        (32, 14), (32, 15), (32, 18), (32, 20), (32, 22), (32, 23),
        (32, 29), (32, 30), (32, 31), (33, 8), (33, 9), (33, 13),
        (33, 14), (33, 15), (33, 18), (33, 19), (33, 20), (33, 22),
        (33, 23), (33, 26), (33, 27), (33, 28), (33, 29), (33, 30),
        (33, 31), (33, 32)]
Minjie Wang's avatar
Minjie Wang committed
67
    # add edges two lists of nodes: src and dst
68
69
    src, dst = tuple(zip(*edge_list))
    g.add_edges(src, dst)
Minjie Wang's avatar
Minjie Wang committed
70
    # edges are directional in DGL; make them bi-directional
71
    g.add_edges(dst, src)
72
73
74
75

    return g

###############################################################################
Minjie Wang's avatar
Minjie Wang committed
76
# We can print out the number of nodes and edges in our newly constructed graph:
77

78
79
80
G = build_karate_club_graph()
print('We have %d nodes.' % G.number_of_nodes())
print('We have %d edges.' % G.number_of_edges())
81

82
###############################################################################
Minjie Wang's avatar
Minjie Wang committed
83
# We can also visualize the graph by converting it to a `networkx
84
# <https://networkx.github.io/documentation/stable/>`_ graph:
85

86
import networkx as nx
Minjie Wang's avatar
Minjie Wang committed
87
88
89
90
91
92
# Since the actual graph is undirected, we convert it for visualization
# purpose.
nx_G = G.to_networkx().to_undirected()
# Kamada-Kawaii layout usually looks pretty for arbitrary graphs
pos = nx.kamada_kawai_layout(nx_G)
nx.draw(nx_G, pos, with_labels=True, node_color=[[.7, .7, .7]])
93
94

###############################################################################
Minjie Wang's avatar
Minjie Wang committed
95
96
97
98
99
100
101
102
103
104
# Step 2: assign features to nodes or edges
# --------------------------------------------
# Graph neural networks associate features with nodes and edges for training.
# For our classification example, we assign each node's an input feature as a one-hot vector:
# node :math:`v_i`'s feature vector is :math:`[0,\ldots,1,\dots,0]`,
# where the :math:`i^{th}` position is one.
#
# In DGL, we can add features for all nodes at once, using a feature tensor that
# batches node features along the first dimension. This code below adds the one-hot
# feature for all nodes:
105
106
107
108

import torch

G.ndata['feat'] = torch.eye(34)
109
110
111


###############################################################################
112
# We can print out the node features to verify:
113

114
115
# print out node 2's input feature
print(G.nodes[2].data['feat'])
116

117
118
# print out node 10 and 11's input features
print(G.nodes[[10, 11]].data['feat'])
119
120

###############################################################################
Minjie Wang's avatar
Minjie Wang committed
121
122
123
124
125
126
# Step 3: define a Graph Convolutional Network (GCN)
# --------------------------------------------------
# To perform node classification, we use the Graph Convolutional Network
# (GCN) developed by `Kipf and Welling <https://arxiv.org/abs/1609.02907>`_. Here
# we provide the simpliest definition of a GCN framework, but we recommend the
# reader to read the original paper for more details.
Minjie Wang's avatar
Minjie Wang committed
127
#
Minjie Wang's avatar
Minjie Wang committed
128
129
130
131
132
# - At layer :math:`l`, each node :math:`v_i^l` carries a feature vector :math:`h_i^l`.
# - Each layer of the GCN tries to aggregate the features from :math:`u_i^{l}` where
#   :math:`u_i`'s are neighborhood nodes to :math:`v` into the next layer representation at
#   :math:`v_i^{l+1}`. This is followed by an affine transformation with some
#   non-linearity.
Minjie Wang's avatar
Minjie Wang committed
133
#
Minjie Wang's avatar
Minjie Wang committed
134
135
136
# The above definition of GCN fits into a **message-passing** paradigm: each
# node will update its own feature with information sent from neighboring
# nodes. A graphical demonstration is displayed below.
137
#
138
# .. image:: https://s3.us-east-2.amazonaws.com/dgl.ai/tutorial/1_first/mailbox.png
Minjie Wang's avatar
Minjie Wang committed
139
140
#    :alt: mailbox
#    :align: center
141
#
Minjie Wang's avatar
Minjie Wang committed
142
# Now, we show that the GCN layer can be easily implemented in DGL.
143
144
145
146
147

import torch.nn as nn
import torch.nn.functional as F

# Define the message & reduce function
Minjie Wang's avatar
Minjie Wang committed
148
# NOTE: we ignore the GCN's normalization constant c_ij for this tutorial.
149
150
def gcn_message(edges):
    # The argument is a batch of edges.
Minjie Wang's avatar
Minjie Wang committed
151
    # This computes a (batch of) message called 'msg' using the source node's feature 'h'.
152
153
154
155
    return {'msg' : edges.src['h']}

def gcn_reduce(nodes):
    # The argument is a batch of nodes.
Minjie Wang's avatar
Minjie Wang committed
156
    # This computes the new 'h' features by summing received 'msg' in each node's mailbox.
157
158
159
160
161
162
163
164
165
166
167
168
    return {'h' : torch.sum(nodes.mailbox['msg'], dim=1)}

# Define the GCNLayer module
class GCNLayer(nn.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
Minjie Wang's avatar
Minjie Wang committed
169
        # trigger message passing on all edges 
170
        g.send(g.edges(), gcn_message)
Minjie Wang's avatar
Minjie Wang committed
171
        # trigger aggregation at all nodes
172
173
174
175
176
177
178
        g.recv(g.nodes(), gcn_reduce)
        # get the result node features
        h = g.ndata.pop('h')
        # perform linear transformation
        return self.linear(h)

###############################################################################
Minjie Wang's avatar
Minjie Wang committed
179
180
181
182
# In general, the nodes send information computed via the *message functions*,
# and aggregates incoming information with the *reduce functions*.
#
# We then define a deeper GCN model that contains two GCN layers:
183
184

# Define a 2-layer GCN model
Minjie Wang's avatar
Minjie Wang committed
185
class GCN(nn.Module):
186
    def __init__(self, in_feats, hidden_size, num_classes):
Minjie Wang's avatar
Minjie Wang committed
187
        super(GCN, self).__init__()
188
189
190
191
192
193
194
195
        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
Minjie Wang's avatar
Minjie Wang committed
196
197
198
199
# The first layer transforms input features of size of 34 to a hidden size of 5.
# The second layer transforms the hidden layer and produces output features of
# size 2, corresponding to the two groups of the karate club.
net = GCN(34, 5, 2)
200
201

###############################################################################
Minjie Wang's avatar
Minjie Wang committed
202
203
# Step 4: data preparation and initialization
# -------------------------------------------
204
#
Minjie Wang's avatar
Minjie Wang committed
205
206
207
# We use one-hot vectors to initialize the node features. Since this is a
# semi-supervised setting, only the instructor (node 0) and the club president
# (node 33) are assigned labels. The implementation is available as follow.
208
209
210
211

inputs = torch.eye(34)
labeled_nodes = torch.tensor([0, 33])  # only the instructor and the president nodes are labeled
labels = torch.tensor([0, 1])  # their labels are different
212

213
###############################################################################
Minjie Wang's avatar
Minjie Wang committed
214
215
216
217
218
# Step 5: train then visualize
# ----------------------------
# The training loop is exactly the same as other PyTorch models.
# We (1) create an optimizer, (2) feed the inputs to the model,
# (3) calculate the loss and (4) use autograd to optimize the model.
219

220
221
222
223
224
225
226
227
228
optimizer = torch.optim.Adam(net.parameters(), lr=0.01)
all_logits = []
for epoch in range(30):
    logits = net(G, inputs)
    # we save the logits for visualization later
    all_logits.append(logits.detach())
    logp = F.log_softmax(logits, 1)
    # we only compute loss for labeled nodes
    loss = F.nll_loss(logp[labeled_nodes], labels)
229

230
231
232
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
233

234
    print('Epoch %d | Loss: %.4f' % (epoch, loss.item()))
235

236
###############################################################################
Minjie Wang's avatar
Minjie Wang committed
237
238
239
240
241
242
243
# This is a rather toy example, so it does not even have a validation or test
# set. Instead, Since the model produces an output feature of size 2 for each node, we can
# visualize by plotting the output feature in a 2D space.
# The following code animates the training process from initial guess
# (where the nodes are not classified correctly at all) to the end
# (where the nodes are linearly separable).

244
245
import matplotlib.animation as animation
import matplotlib.pyplot as plt
246

247
248
249
250
251
252
253
254
255
256
257
258
259
260
def draw(i):
    cls1color = '#00FFFF'
    cls2color = '#FF00FF'
    pos = {}
    colors = []
    for v in range(34):
        pos[v] = all_logits[i][v].numpy()
        cls = pos[v].argmax()
        colors.append(cls1color if cls else cls2color)
    ax.cla()
    ax.axis('off')
    ax.set_title('Epoch: %d' % i)
    nx.draw_networkx(nx_G.to_undirected(), pos, node_color=colors,
            with_labels=True, node_size=300, ax=ax)
261

262
263
264
265
266
fig = plt.figure(dpi=150)
fig.clf()
ax = fig.subplots()
draw(0)  # draw the prediction of the first epoch
plt.close()
267

268
269
270
271
272
273
274
275
###############################################################################
# .. image:: https://s3.us-east-2.amazonaws.com/dgl.ai/tutorial/1_first/karate0.png
#    :height: 300px
#    :width: 400px
#    :align: center

###############################################################################
# The following animation shows how the model correctly predicts the community
Minjie Wang's avatar
Minjie Wang committed
276
# after a series of training epochs.
277
278
279
280
281
282
283
284

ani = animation.FuncAnimation(fig, draw, frames=len(all_logits), interval=200)

###############################################################################
# .. image:: https://s3.us-east-2.amazonaws.com/dgl.ai/tutorial/1_first/karate.gif
#    :height: 300px
#    :width: 400px
#    :align: center
285
286
287
288

###############################################################################
# Next steps
# ----------
Minjie Wang's avatar
Minjie Wang committed
289
# 
Minjie Wang's avatar
Minjie Wang committed
290
291
# In the :doc:`next tutorial <2_basics>`, we will go through some more basics
# of DGL, such as reading and writing node/edge features.