chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,205 @@
|
||||
"""
|
||||
.. _model-gcn:
|
||||
|
||||
Graph Convolutional Network
|
||||
====================================
|
||||
|
||||
**Author:** `Qi Huang <https://github.com/HQ01>`_, `Minjie Wang <https://jermainewang.github.io/>`_,
|
||||
Yu Gai, Quan Gan, Zheng Zhang
|
||||
|
||||
.. warning::
|
||||
|
||||
The tutorial aims at gaining insights into the paper, with code as a mean
|
||||
of explanation. The implementation thus is NOT optimized for running
|
||||
efficiency. For recommended implementation, please refer to the `official
|
||||
examples <https://github.com/dmlc/dgl/tree/master/examples>`_.
|
||||
|
||||
This is a gentle introduction of using DGL to implement Graph Convolutional
|
||||
Networks (Kipf & Welling et al., `Semi-Supervised Classification with Graph
|
||||
Convolutional Networks <https://arxiv.org/pdf/1609.02907.pdf>`_). We explain
|
||||
what is under the hood of the :class:`~dgl.nn.GraphConv` module.
|
||||
The reader is expected to learn how to define a new GNN layer using DGL's
|
||||
message passing APIs.
|
||||
"""
|
||||
|
||||
###############################################################################
|
||||
# Model Overview
|
||||
# ------------------------------------------
|
||||
# GCN from the perspective of message passing
|
||||
# ```````````````````````````````````````````````
|
||||
# We describe a layer of graph convolutional neural network from a message
|
||||
# passing perspective; the math can be found `here <math_>`_.
|
||||
# It boils down to the following step, for each node :math:`u`:
|
||||
#
|
||||
# 1) Aggregate neighbors' representations :math:`h_{v}` to produce an
|
||||
# intermediate representation :math:`\hat{h}_u`. 2) Transform the aggregated
|
||||
# representation :math:`\hat{h}_{u}` with a linear projection followed by a
|
||||
# non-linearity: :math:`h_{u} = f(W_{u} \hat{h}_u)`.
|
||||
#
|
||||
# We will implement step 1 with DGL message passing, and step 2 by
|
||||
# PyTorch ``nn.Module``.
|
||||
#
|
||||
# GCN implementation with DGL
|
||||
# ``````````````````````````````````````````
|
||||
# We first define the message and reduce function as usual. Since the
|
||||
# aggregation on a node :math:`u` only involves summing over the neighbors'
|
||||
# representations :math:`h_v`, we can simply use builtin functions:
|
||||
|
||||
import os
|
||||
|
||||
os.environ["DGLBACKEND"] = "pytorch"
|
||||
import dgl
|
||||
import dgl.function as fn
|
||||
import torch as th
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from dgl import DGLGraph
|
||||
|
||||
gcn_msg = fn.copy_u(u="h", out="m")
|
||||
gcn_reduce = fn.sum(msg="m", out="h")
|
||||
|
||||
###############################################################################
|
||||
# We then proceed to define the GCNLayer module. A GCNLayer essentially performs
|
||||
# message passing on all the nodes then applies a fully-connected layer.
|
||||
#
|
||||
# .. note::
|
||||
#
|
||||
# This is showing how to implement a GCN from scratch. DGL provides a more
|
||||
# efficient :class:`builtin GCN layer module <dgl.nn.pytorch.conv.GraphConv>`.
|
||||
#
|
||||
|
||||
|
||||
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, feature):
|
||||
# Creating a local scope so that all the stored ndata and edata
|
||||
# (such as the `'h'` ndata below) are automatically popped out
|
||||
# when the scope exits.
|
||||
with g.local_scope():
|
||||
g.ndata["h"] = feature
|
||||
g.update_all(gcn_msg, gcn_reduce)
|
||||
h = g.ndata["h"]
|
||||
return self.linear(h)
|
||||
|
||||
|
||||
###############################################################################
|
||||
# The forward function is essentially the same as any other commonly seen NNs
|
||||
# model in PyTorch. We can initialize GCN like any ``nn.Module``. For example,
|
||||
# let's define a simple neural network consisting of two GCN layers. Suppose we
|
||||
# are training the classifier for the cora dataset (the input feature size is
|
||||
# 1433 and the number of classes is 7). The last GCN layer computes node embeddings,
|
||||
# so the last layer in general does not apply activation.
|
||||
|
||||
|
||||
class Net(nn.Module):
|
||||
def __init__(self):
|
||||
super(Net, self).__init__()
|
||||
self.layer1 = GCNLayer(1433, 16)
|
||||
self.layer2 = GCNLayer(16, 7)
|
||||
|
||||
def forward(self, g, features):
|
||||
x = F.relu(self.layer1(g, features))
|
||||
x = self.layer2(g, x)
|
||||
return x
|
||||
|
||||
|
||||
net = Net()
|
||||
print(net)
|
||||
|
||||
###############################################################################
|
||||
# We load the cora dataset using DGL's built-in data module.
|
||||
|
||||
from dgl.data import CoraGraphDataset
|
||||
|
||||
|
||||
def load_cora_data():
|
||||
dataset = CoraGraphDataset()
|
||||
g = dataset[0]
|
||||
features = g.ndata["feat"]
|
||||
labels = g.ndata["label"]
|
||||
train_mask = g.ndata["train_mask"]
|
||||
test_mask = g.ndata["test_mask"]
|
||||
return g, features, labels, train_mask, test_mask
|
||||
|
||||
|
||||
###############################################################################
|
||||
# When a model is trained, we can use the following method to evaluate
|
||||
# the performance of the model on the test dataset:
|
||||
|
||||
|
||||
def evaluate(model, g, features, labels, mask):
|
||||
model.eval()
|
||||
with th.no_grad():
|
||||
logits = model(g, features)
|
||||
logits = logits[mask]
|
||||
labels = labels[mask]
|
||||
_, indices = th.max(logits, dim=1)
|
||||
correct = th.sum(indices == labels)
|
||||
return correct.item() * 1.0 / len(labels)
|
||||
|
||||
|
||||
###############################################################################
|
||||
# We then train the network as follows:
|
||||
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
|
||||
g, features, labels, train_mask, test_mask = load_cora_data()
|
||||
# Add edges between each node and itself to preserve old node representations
|
||||
g.add_edges(g.nodes(), g.nodes())
|
||||
optimizer = th.optim.Adam(net.parameters(), lr=1e-2)
|
||||
dur = []
|
||||
for epoch in range(50):
|
||||
if epoch >= 3:
|
||||
t0 = time.time()
|
||||
net.train()
|
||||
logits = net(g, features)
|
||||
logp = F.log_softmax(logits, 1)
|
||||
loss = F.nll_loss(logp[train_mask], labels[train_mask])
|
||||
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
if epoch >= 3:
|
||||
dur.append(time.time() - t0)
|
||||
acc = evaluate(net, g, features, labels, test_mask)
|
||||
print(
|
||||
"Epoch {:05d} | Loss {:.4f} | Test Acc {:.4f} | Time(s) {:.4f}".format(
|
||||
epoch, loss.item(), acc, np.mean(dur)
|
||||
)
|
||||
)
|
||||
###############################################################################
|
||||
# .. _math:
|
||||
#
|
||||
# GCN in one formula
|
||||
# ------------------
|
||||
# Mathematically, the GCN model follows this formula:
|
||||
#
|
||||
# :math:`H^{(l+1)} = \sigma(\tilde{D}^{-\frac{1}{2}}\tilde{A}\tilde{D}^{-\frac{1}{2}}H^{(l)}W^{(l)})`
|
||||
#
|
||||
# Here, :math:`H^{(l)}` denotes the :math:`l^{th}` layer in the network,
|
||||
# :math:`\sigma` is the non-linearity, and :math:`W` is the weight matrix for
|
||||
# this layer. :math:`\tilde{D}` and :math:`\tilde{A}` are separately the degree
|
||||
# and adjacency matrices for the graph. With the superscript ~, we are referring
|
||||
# to the variant where we add additional edges between each node and itself to
|
||||
# preserve its old representation in graph convolutions. The shape of the input
|
||||
# :math:`H^{(0)}` is :math:`N \times D`, where :math:`N` is the number of nodes
|
||||
# and :math:`D` is the number of input features. We can chain up multiple
|
||||
# layers as such to produce a node-level representation output with shape
|
||||
# :math:`N \times F`, where :math:`F` is the dimension of the output node
|
||||
# feature vector.
|
||||
#
|
||||
# The equation can be efficiently implemented using sparse matrix
|
||||
# multiplication kernels (such as Kipf's
|
||||
# `pygcn <https://github.com/tkipf/pygcn>`_ code). The above DGL implementation
|
||||
# in fact has already used this trick due to the use of builtin functions.
|
||||
#
|
||||
# Note that the tutorial code implements a simplified version of GCN where we
|
||||
# replace :math:`\tilde{D}^{-\frac{1}{2}}\tilde{A}\tilde{D}^{-\frac{1}{2}}` with
|
||||
# :math:`\tilde{A}`. For a full implementation, see our example
|
||||
# `here <https://github.com/dmlc/dgl/tree/master/examples/pytorch/gcn>`_.
|
||||
@@ -0,0 +1,417 @@
|
||||
"""
|
||||
.. _model-rgcn:
|
||||
|
||||
Relational Graph Convolutional Network
|
||||
================================================
|
||||
|
||||
**Author:** Lingfan Yu, Mufei Li, Zheng Zhang
|
||||
|
||||
.. warning::
|
||||
|
||||
The tutorial aims at gaining insights into the paper, with code as a mean
|
||||
of explanation. The implementation thus is NOT optimized for running
|
||||
efficiency. For recommended implementation, please refer to the `official
|
||||
examples <https://github.com/dmlc/dgl/tree/master/examples>`_.
|
||||
|
||||
In this tutorial, you learn how to implement a relational graph convolutional
|
||||
network (R-GCN). This type of network is one effort to generalize GCN
|
||||
to handle different relationships between entities in a knowledge base. To
|
||||
learn more about the research behind R-GCN, see `Modeling Relational Data with Graph Convolutional
|
||||
Networks <https://arxiv.org/pdf/1703.06103.pdf>`_
|
||||
|
||||
The straightforward graph convolutional network (GCN) exploits
|
||||
structural information of a dataset (that is, the graph connectivity) in order to
|
||||
improve the extraction of node representations. Graph edges are left as
|
||||
untyped.
|
||||
|
||||
A knowledge graph is made up of a collection of triples in the form
|
||||
subject, relation, object. Edges thus encode important information and
|
||||
have their own embeddings to be learned. Furthermore, there may exist
|
||||
multiple edges among any given pair.
|
||||
|
||||
"""
|
||||
###############################################################################
|
||||
# A brief introduction to R-GCN
|
||||
# ---------------------------
|
||||
# In *statistical relational learning* (SRL), there are two fundamental
|
||||
# tasks:
|
||||
#
|
||||
# - **Entity classification** - Where you assign types and categorical
|
||||
# properties to entities.
|
||||
# - **Link prediction** - Where you recover missing triples.
|
||||
#
|
||||
# In both cases, missing information is expected to be recovered from the
|
||||
# neighborhood structure of the graph. For example, the R-GCN
|
||||
# paper cited earlier provides the following example. Knowing that Mikhail Baryshnikov was educated at the Vaganova Academy
|
||||
# implies both that Mikhail Baryshnikov should have the label person, and
|
||||
# that the triple (Mikhail Baryshnikov, lived in, Russia) must belong to the
|
||||
# knowledge graph.
|
||||
#
|
||||
# R-GCN solves these two problems using a common graph convolutional network. It's
|
||||
# extended with multi-edge encoding to compute embedding of the entities, but
|
||||
# with different downstream processing.
|
||||
#
|
||||
# - Entity classification is done by attaching a softmax classifier at the
|
||||
# final embedding of an entity (node). Training is through loss of standard
|
||||
# cross-entropy.
|
||||
# - Link prediction is done by reconstructing an edge with an autoencoder
|
||||
# architecture, using a parameterized score function. Training uses negative
|
||||
# sampling.
|
||||
#
|
||||
# This tutorial focuses on the first task, entity classification, to show how to generate entity
|
||||
# representation. `Complete
|
||||
# code <https://github.com/dmlc/dgl/tree/master/examples/pytorch/rgcn>`_
|
||||
# for both tasks is found in the DGL Github repository.
|
||||
#
|
||||
# Key ideas of R-GCN
|
||||
# -------------------
|
||||
# Recall that in GCN, the hidden representation for each node :math:`i` at
|
||||
# :math:`(l+1)^{th}` layer is computed by:
|
||||
#
|
||||
# .. math:: h_i^{l+1} = \sigma\left(\sum_{j\in N_i}\frac{1}{c_i} W^{(l)} h_j^{(l)}\right)~~~~~~~~~~(1)\\
|
||||
#
|
||||
# where :math:`c_i` is a normalization constant.
|
||||
#
|
||||
# The key difference between R-GCN and GCN is that in R-GCN, edges can
|
||||
# represent different relations. In GCN, weight :math:`W^{(l)}` in equation
|
||||
# :math:`(1)` is shared by all edges in layer :math:`l`. In contrast, in
|
||||
# R-GCN, different edge types use different weights and only edges of the
|
||||
# same relation type :math:`r` are associated with the same projection weight
|
||||
# :math:`W_r^{(l)}`.
|
||||
#
|
||||
# So the hidden representation of entities in :math:`(l+1)^{th}` layer in
|
||||
# R-GCN can be formulated as the following equation:
|
||||
#
|
||||
# .. math:: h_i^{l+1} = \sigma\left(W_0^{(l)}h_i^{(l)}+\sum_{r\in R}\sum_{j\in N_i^r}\frac{1}{c_{i,r}}W_r^{(l)}h_j^{(l)}\right)~~~~~~~~~~(2)\\
|
||||
#
|
||||
# where :math:`N_i^r` denotes the set of neighbor indices of node :math:`i`
|
||||
# under relation :math:`r\in R` and :math:`c_{i,r}` is a normalization
|
||||
# constant. In entity classification, the R-GCN paper uses
|
||||
# :math:`c_{i,r}=|N_i^r|`.
|
||||
#
|
||||
# The problem of applying the above equation directly is the rapid growth of
|
||||
# the number of parameters, especially with highly multi-relational data. In
|
||||
# order to reduce model parameter size and prevent overfitting, the original
|
||||
# paper proposes to use basis decomposition.
|
||||
#
|
||||
# .. math:: W_r^{(l)}=\sum\limits_{b=1}^B a_{rb}^{(l)}V_b^{(l)}~~~~~~~~~~(3)\\
|
||||
#
|
||||
# Therefore, the weight :math:`W_r^{(l)}` is a linear combination of basis
|
||||
# transformation :math:`V_b^{(l)}` with coefficients :math:`a_{rb}^{(l)}`.
|
||||
# The number of bases :math:`B` is much smaller than the number of relations
|
||||
# in the knowledge base.
|
||||
#
|
||||
# .. note::
|
||||
# Another weight regularization, block-decomposition, is implemented in
|
||||
# the `link prediction <link-prediction_>`_.
|
||||
#
|
||||
# Implement R-GCN in DGL
|
||||
# ----------------------
|
||||
#
|
||||
# An R-GCN model is composed of several R-GCN layers. The first R-GCN layer
|
||||
# also serves as input layer and takes in features (for example, description texts)
|
||||
# that are associated with node entity and project to hidden space. In this tutorial,
|
||||
# we only use the entity ID as an entity feature.
|
||||
#
|
||||
# R-GCN layers
|
||||
# ~~~~~~~~~~~~
|
||||
#
|
||||
# For each node, an R-GCN layer performs the following steps:
|
||||
#
|
||||
# - Compute outgoing message using node representation and weight matrix
|
||||
# associated with the edge type (message function)
|
||||
# - Aggregate incoming messages and generate new node representations (reduce
|
||||
# and apply function)
|
||||
#
|
||||
# The following code is the definition of an R-GCN hidden layer.
|
||||
#
|
||||
# .. note::
|
||||
# Each relation type is associated with a different weight. Therefore,
|
||||
# the full weight matrix has three dimensions: relation, input_feature,
|
||||
# output_feature.
|
||||
#
|
||||
# .. note::
|
||||
#
|
||||
# This is showing how to implement an R-GCN from scratch. DGL provides a more
|
||||
# efficient :class:`builtin R-GCN layer module <dgl.nn.pytorch.conv.RelGraphConv>`.
|
||||
#
|
||||
|
||||
import os
|
||||
|
||||
os.environ["DGLBACKEND"] = "pytorch"
|
||||
from functools import partial
|
||||
|
||||
import dgl
|
||||
import dgl.function as fn
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from dgl import DGLGraph
|
||||
|
||||
|
||||
class RGCNLayer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_feat,
|
||||
out_feat,
|
||||
num_rels,
|
||||
num_bases=-1,
|
||||
bias=None,
|
||||
activation=None,
|
||||
is_input_layer=False,
|
||||
):
|
||||
super(RGCNLayer, self).__init__()
|
||||
self.in_feat = in_feat
|
||||
self.out_feat = out_feat
|
||||
self.num_rels = num_rels
|
||||
self.num_bases = num_bases
|
||||
self.bias = bias
|
||||
self.activation = activation
|
||||
self.is_input_layer = is_input_layer
|
||||
|
||||
# sanity check
|
||||
if self.num_bases <= 0 or self.num_bases > self.num_rels:
|
||||
self.num_bases = self.num_rels
|
||||
# weight bases in equation (3)
|
||||
self.weight = nn.Parameter(
|
||||
torch.Tensor(self.num_bases, self.in_feat, self.out_feat)
|
||||
)
|
||||
if self.num_bases < self.num_rels:
|
||||
# linear combination coefficients in equation (3)
|
||||
self.w_comp = nn.Parameter(
|
||||
torch.Tensor(self.num_rels, self.num_bases)
|
||||
)
|
||||
# add bias
|
||||
if self.bias:
|
||||
self.bias = nn.Parameter(torch.Tensor(out_feat))
|
||||
# init trainable parameters
|
||||
nn.init.xavier_uniform_(
|
||||
self.weight, gain=nn.init.calculate_gain("relu")
|
||||
)
|
||||
if self.num_bases < self.num_rels:
|
||||
nn.init.xavier_uniform_(
|
||||
self.w_comp, gain=nn.init.calculate_gain("relu")
|
||||
)
|
||||
if self.bias:
|
||||
nn.init.xavier_uniform_(
|
||||
self.bias, gain=nn.init.calculate_gain("relu")
|
||||
)
|
||||
|
||||
def forward(self, g):
|
||||
if self.num_bases < self.num_rels:
|
||||
# generate all weights from bases (equation (3))
|
||||
weight = self.weight.view(
|
||||
self.in_feat, self.num_bases, self.out_feat
|
||||
)
|
||||
weight = torch.matmul(self.w_comp, weight).view(
|
||||
self.num_rels, self.in_feat, self.out_feat
|
||||
)
|
||||
else:
|
||||
weight = self.weight
|
||||
if self.is_input_layer:
|
||||
|
||||
def message_func(edges):
|
||||
# for input layer, matrix multiply can be converted to be
|
||||
# an embedding lookup using source node id
|
||||
embed = weight.view(-1, self.out_feat)
|
||||
index = edges.data[dgl.ETYPE] * self.in_feat + edges.src["id"]
|
||||
return {"msg": embed[index] * edges.data["norm"]}
|
||||
|
||||
else:
|
||||
|
||||
def message_func(edges):
|
||||
w = weight[edges.data[dgl.ETYPE]]
|
||||
msg = torch.bmm(edges.src["h"].unsqueeze(1), w).squeeze()
|
||||
msg = msg * edges.data["norm"]
|
||||
return {"msg": msg}
|
||||
|
||||
def apply_func(nodes):
|
||||
h = nodes.data["h"]
|
||||
if self.bias:
|
||||
h = h + self.bias
|
||||
if self.activation:
|
||||
h = self.activation(h)
|
||||
return {"h": h}
|
||||
|
||||
g.update_all(message_func, fn.sum(msg="msg", out="h"), apply_func)
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Full R-GCN model defined
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
num_nodes,
|
||||
h_dim,
|
||||
out_dim,
|
||||
num_rels,
|
||||
num_bases=-1,
|
||||
num_hidden_layers=1,
|
||||
):
|
||||
super(Model, self).__init__()
|
||||
self.num_nodes = num_nodes
|
||||
self.h_dim = h_dim
|
||||
self.out_dim = out_dim
|
||||
self.num_rels = num_rels
|
||||
self.num_bases = num_bases
|
||||
self.num_hidden_layers = num_hidden_layers
|
||||
|
||||
# create rgcn layers
|
||||
self.build_model()
|
||||
|
||||
# create initial features
|
||||
self.features = self.create_features()
|
||||
|
||||
def build_model(self):
|
||||
self.layers = nn.ModuleList()
|
||||
# input to hidden
|
||||
i2h = self.build_input_layer()
|
||||
self.layers.append(i2h)
|
||||
# hidden to hidden
|
||||
for _ in range(self.num_hidden_layers):
|
||||
h2h = self.build_hidden_layer()
|
||||
self.layers.append(h2h)
|
||||
# hidden to output
|
||||
h2o = self.build_output_layer()
|
||||
self.layers.append(h2o)
|
||||
|
||||
# initialize feature for each node
|
||||
def create_features(self):
|
||||
features = torch.arange(self.num_nodes)
|
||||
return features
|
||||
|
||||
def build_input_layer(self):
|
||||
return RGCNLayer(
|
||||
self.num_nodes,
|
||||
self.h_dim,
|
||||
self.num_rels,
|
||||
self.num_bases,
|
||||
activation=F.relu,
|
||||
is_input_layer=True,
|
||||
)
|
||||
|
||||
def build_hidden_layer(self):
|
||||
return RGCNLayer(
|
||||
self.h_dim,
|
||||
self.h_dim,
|
||||
self.num_rels,
|
||||
self.num_bases,
|
||||
activation=F.relu,
|
||||
)
|
||||
|
||||
def build_output_layer(self):
|
||||
return RGCNLayer(
|
||||
self.h_dim,
|
||||
self.out_dim,
|
||||
self.num_rels,
|
||||
self.num_bases,
|
||||
activation=partial(F.softmax, dim=1),
|
||||
)
|
||||
|
||||
def forward(self, g):
|
||||
if self.features is not None:
|
||||
g.ndata["id"] = self.features
|
||||
for layer in self.layers:
|
||||
layer(g)
|
||||
return g.ndata.pop("h")
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Handle dataset
|
||||
# ~~~~~~~~~~~~~~~~
|
||||
# This tutorial uses Institute for Applied Informatics and Formal Description Methods (AIFB) dataset from R-GCN paper.
|
||||
|
||||
# load graph data
|
||||
dataset = dgl.data.rdf.AIFBDataset()
|
||||
g = dataset[0]
|
||||
category = dataset.predict_category
|
||||
train_mask = g.nodes[category].data.pop("train_mask")
|
||||
test_mask = g.nodes[category].data.pop("test_mask")
|
||||
train_idx = torch.nonzero(train_mask, as_tuple=False).squeeze()
|
||||
test_idx = torch.nonzero(test_mask, as_tuple=False).squeeze()
|
||||
labels = g.nodes[category].data.pop("label")
|
||||
num_rels = len(g.canonical_etypes)
|
||||
num_classes = dataset.num_classes
|
||||
# normalization factor
|
||||
for cetype in g.canonical_etypes:
|
||||
g.edges[cetype].data["norm"] = dgl.norm_by_dst(g, cetype).unsqueeze(1)
|
||||
category_id = g.ntypes.index(category)
|
||||
|
||||
###############################################################################
|
||||
# Create graph and model
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
# configurations
|
||||
n_hidden = 16 # number of hidden units
|
||||
n_bases = -1 # use number of relations as number of bases
|
||||
n_hidden_layers = 0 # use 1 input layer, 1 output layer, no hidden layer
|
||||
n_epochs = 25 # epochs to train
|
||||
lr = 0.01 # learning rate
|
||||
l2norm = 0 # L2 norm coefficient
|
||||
|
||||
# create graph
|
||||
g = dgl.to_homogeneous(g, edata=["norm"])
|
||||
node_ids = torch.arange(g.num_nodes())
|
||||
target_idx = node_ids[g.ndata[dgl.NTYPE] == category_id]
|
||||
|
||||
# create model
|
||||
model = Model(
|
||||
g.num_nodes(),
|
||||
n_hidden,
|
||||
num_classes,
|
||||
num_rels,
|
||||
num_bases=n_bases,
|
||||
num_hidden_layers=n_hidden_layers,
|
||||
)
|
||||
|
||||
###############################################################################
|
||||
# Training loop
|
||||
# ~~~~~~~~~~~~~~~~
|
||||
|
||||
# optimizer
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=l2norm)
|
||||
|
||||
print("start training...")
|
||||
model.train()
|
||||
for epoch in range(n_epochs):
|
||||
optimizer.zero_grad()
|
||||
logits = model.forward(g)
|
||||
logits = logits[target_idx]
|
||||
loss = F.cross_entropy(logits[train_idx], labels[train_idx])
|
||||
loss.backward()
|
||||
|
||||
optimizer.step()
|
||||
|
||||
train_acc = torch.sum(logits[train_idx].argmax(dim=1) == labels[train_idx])
|
||||
train_acc = train_acc.item() / len(train_idx)
|
||||
val_loss = F.cross_entropy(logits[test_idx], labels[test_idx])
|
||||
val_acc = torch.sum(logits[test_idx].argmax(dim=1) == labels[test_idx])
|
||||
val_acc = val_acc.item() / len(test_idx)
|
||||
print(
|
||||
"Epoch {:05d} | ".format(epoch)
|
||||
+ "Train Accuracy: {:.4f} | Train Loss: {:.4f} | ".format(
|
||||
train_acc, loss.item()
|
||||
)
|
||||
+ "Validation Accuracy: {:.4f} | Validation loss: {:.4f}".format(
|
||||
val_acc, val_loss.item()
|
||||
)
|
||||
)
|
||||
###############################################################################
|
||||
# .. _link-prediction:
|
||||
#
|
||||
# The second task, link prediction
|
||||
# --------------------------------
|
||||
# So far, you have seen how to use DGL to implement entity classification with an
|
||||
# R-GCN model. In the knowledge base setting, representation generated by
|
||||
# R-GCN can be used to uncover potential relationships between nodes. In the
|
||||
# R-GCN paper, the authors feed the entity representations generated by R-GCN
|
||||
# into the `DistMult <https://arxiv.org/pdf/1412.6575.pdf>`_ prediction model
|
||||
# to predict possible relationships.
|
||||
#
|
||||
# The implementation is similar to that presented here, but with an extra DistMult layer
|
||||
# stacked on top of the R-GCN layers. You can find the complete
|
||||
# implementation of link prediction with R-GCN in our `Github Python code
|
||||
# example <https://github.com/dmlc/dgl/blob/master/examples/pytorch/rgcn/link.py>`_.
|
||||
@@ -0,0 +1,653 @@
|
||||
"""
|
||||
.. _model-line-graph:
|
||||
|
||||
Line Graph Neural Network
|
||||
=========================
|
||||
|
||||
**Author**: `Qi Huang <https://github.com/HQ01>`_, Yu Gai,
|
||||
`Minjie Wang <https://jermainewang.github.io/>`_, Zheng Zhang
|
||||
|
||||
.. warning::
|
||||
|
||||
The tutorial aims at gaining insights into the paper, with code as a mean
|
||||
of explanation. The implementation thus is NOT optimized for running
|
||||
efficiency. For recommended implementation, please refer to the `official
|
||||
examples <https://github.com/dmlc/dgl/tree/master/examples>`_.
|
||||
|
||||
"""
|
||||
|
||||
###########################################################################################
|
||||
#
|
||||
# In this tutorial, you learn how to solve community detection tasks by implementing a line
|
||||
# graph neural network (LGNN). Community detection, or graph clustering, consists of partitioning
|
||||
# the vertices in a graph into clusters in which nodes are more similar to
|
||||
# one another.
|
||||
#
|
||||
# In the :doc:`Graph convolutinal network tutorial <1_gcn>`, you learned how to classify the nodes of an input
|
||||
# graph in a semi-supervised setting. You used a graph convolutional neural network (GCN)
|
||||
# as an embedding mechanism for graph features.
|
||||
#
|
||||
# To generalize a graph neural network (GNN) into supervised community detection, a line-graph based
|
||||
# variation of GNN is introduced in the research paper
|
||||
# `Supervised Community Detection with Line Graph Neural Networks <https://arxiv.org/abs/1705.08415>`__.
|
||||
# One of the highlights of the model is
|
||||
# to augment the straightforward GNN architecture so that it operates on
|
||||
# a line graph of edge adjacencies, defined with a non-backtracking operator.
|
||||
#
|
||||
# A line graph neural network (LGNN) shows how DGL can implement an advanced graph algorithm by
|
||||
# mixing basic tensor operations, sparse-matrix multiplication, and message-
|
||||
# passing APIs.
|
||||
#
|
||||
# In the following sections, you learn about community detection, line
|
||||
# graphs, LGNN, and its implementation.
|
||||
#
|
||||
# Supervised community detection task with the Cora dataset
|
||||
# --------------------------------------------
|
||||
# Community detection
|
||||
# ~~~~~~~~~~~~~~~~~~~~
|
||||
# In a community detection task, you cluster similar nodes instead of
|
||||
# labeling them. The node similarity is typically described as having higher inner
|
||||
# density within each cluster.
|
||||
#
|
||||
# What's the difference between community detection and node classification？
|
||||
# Comparing to node classification, community detection focuses on retrieving
|
||||
# cluster information in the graph, rather than assigning a specific label to
|
||||
# a node. For example, as long as a node is clustered with its community
|
||||
# members, it doesn't matter whether the node is assigned as "community A",
|
||||
# or "community B", while assigning all "great movies" to label "bad movies"
|
||||
# will be a disaster in a movie network classification task.
|
||||
#
|
||||
# What's the difference then, between a community detection algorithm and
|
||||
# other clustering algorithm such as k-means? Community detection algorithm operates on
|
||||
# graph-structured data. Comparing to k-means, community detection leverages
|
||||
# graph structure, instead of simply clustering nodes based on their
|
||||
# features.
|
||||
#
|
||||
# Cora dataset
|
||||
# ~~~~~
|
||||
# To be consistent with the GCN tutorial,
|
||||
# you use the `Cora dataset <https://linqs.soe.ucsc.edu/data>`__
|
||||
# to illustrate a simple community detection task. Cora is a scientific publication dataset,
|
||||
# with 2708 papers belonging to seven
|
||||
# different machine learning fields. Here, you formulate Cora as a
|
||||
# directed graph, with each node being a paper, and each edge being a
|
||||
# citation link (A->B means A cites B). Here is a visualization of the whole
|
||||
# Cora dataset.
|
||||
#
|
||||
# .. figure:: https://i.imgur.com/X404Byc.png
|
||||
# :alt: cora
|
||||
# :height: 400px
|
||||
# :width: 500px
|
||||
# :align: center
|
||||
#
|
||||
# Cora naturally contains seven classes, and statistics below show that each
|
||||
# class does satisfy our assumption of community, i.e. nodes of same class
|
||||
# class have higher connection probability among them than with nodes of different class.
|
||||
# The following code snippet verifies that there are more intra-class edges
|
||||
# than inter-class.
|
||||
|
||||
import os
|
||||
|
||||
os.environ["DGLBACKEND"] = "pytorch"
|
||||
import dgl
|
||||
import torch
|
||||
import torch as th
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from dgl.data import citation_graph as citegrh
|
||||
|
||||
data = citegrh.load_cora()
|
||||
|
||||
G = data[0]
|
||||
labels = th.tensor(G.ndata["label"])
|
||||
|
||||
# find all the nodes labeled with class 0
|
||||
label0_nodes = th.nonzero(labels == 0, as_tuple=False).squeeze()
|
||||
# find all the edges pointing to class 0 nodes
|
||||
src, _ = G.in_edges(label0_nodes)
|
||||
src_labels = labels[src]
|
||||
# find all the edges whose both endpoints are in class 0
|
||||
intra_src = th.nonzero(src_labels == 0, as_tuple=False)
|
||||
print("Intra-class edges percent: %.4f" % (len(intra_src) / len(src_labels)))
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
###########################################################################################
|
||||
# Binary community subgraph from Cora with a test dataset
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# Without loss of generality, in this tutorial you limit the scope of the
|
||||
# task to binary community detection.
|
||||
#
|
||||
# .. note::
|
||||
#
|
||||
# To create a practice binary-community dataset from Cora, first extract
|
||||
# all two-class pairs from the original Cora seven classes. For each pair, you
|
||||
# treat each class as one community, and find the largest subgraph that
|
||||
# at least contains one cross-community edge as the training example. As
|
||||
# a result, there are a total of 21 training samples in this small dataset.
|
||||
#
|
||||
# With the following code, you can visualize one of the training samples and its community structure.
|
||||
|
||||
import networkx as nx
|
||||
|
||||
train_set = dgl.data.CoraBinary()
|
||||
G1, pmpd1, label1 = train_set[1]
|
||||
nx_G1 = G1.to_networkx()
|
||||
|
||||
|
||||
def visualize(labels, g):
|
||||
pos = nx.spring_layout(g, seed=1)
|
||||
plt.figure(figsize=(8, 8))
|
||||
plt.axis("off")
|
||||
nx.draw_networkx(
|
||||
g,
|
||||
pos=pos,
|
||||
node_size=50,
|
||||
cmap=plt.get_cmap("coolwarm"),
|
||||
node_color=labels,
|
||||
edge_color="k",
|
||||
arrows=False,
|
||||
width=0.5,
|
||||
style="dotted",
|
||||
with_labels=False,
|
||||
)
|
||||
|
||||
|
||||
visualize(label1, nx_G1)
|
||||
|
||||
###########################################################################################
|
||||
# To learn more, go the original research paper to see how to generalize
|
||||
# to multiple communities case.
|
||||
#
|
||||
# Community detection in a supervised setting
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# The community detection problem could be tackled with both supervised and
|
||||
# unsupervised approaches. You can formulate
|
||||
# community detection in a supervised setting as follows:
|
||||
#
|
||||
# - Each training example consists of :math:`(G, L)`, where :math:`G` is a
|
||||
# directed graph :math:`(V, E)`. For each node :math:`v` in :math:`V`, we
|
||||
# assign a ground truth community label :math:`z_v \in \{0,1\}`.
|
||||
# - The parameterized model :math:`f(G, \theta)` predicts a label set
|
||||
# :math:`\tilde{Z} = f(G)` for nodes :math:`V`.
|
||||
# - For each example :math:`(G,L)`, the model learns to minimize a specially
|
||||
# designed loss function (equivariant loss) :math:`L_{equivariant} =
|
||||
# (\tilde{Z}，Z)`
|
||||
#
|
||||
# .. note::
|
||||
#
|
||||
# In this supervised setting, the model naturally predicts a label for
|
||||
# each community. However, community assignment should be equivariant to
|
||||
# label permutations. To achieve this, in each forward process, we take
|
||||
# the minimum among losses calculated from all possible permutations of
|
||||
# labels.
|
||||
#
|
||||
# Mathematically, this means
|
||||
# :math:`L_{equivariant} = \underset{\pi \in S_c} {min}-\log(\hat{\pi}, \pi)`,
|
||||
# where :math:`S_c` is the set of all permutations of labels, and
|
||||
# :math:`\hat{\pi}` is the set of predicted labels,
|
||||
# :math:`- \log(\hat{\pi},\pi)` denotes negative log likelihood.
|
||||
#
|
||||
# For instance, for a sample graph with node :math:`\{1,2,3,4\}` and
|
||||
# community assignment :math:`\{A, A, A, B\}`, with each node's label
|
||||
# :math:`l \in \{0,1\}`,The group of all possible permutations
|
||||
# :math:`S_c = \{\{0,0,0,1\}, \{1,1,1,0\}\}`.
|
||||
#
|
||||
# Line graph neural network key ideas
|
||||
# ------------------------------------
|
||||
# An key innovation in this topic is the use of a line graph.
|
||||
# Unlike models in previous tutorials, message passing happens not only on the
|
||||
# original graph, e.g. the binary community subgraph from Cora, but also on the
|
||||
# line graph associated with the original graph.
|
||||
#
|
||||
# What is a line-graph?
|
||||
# ~~~~~~~~~~~~~~~~~~~~~
|
||||
# In graph theory, line graph is a graph representation that encodes the
|
||||
# edge adjacency structure in the original graph.
|
||||
#
|
||||
# Specifically, a line-graph :math:`L(G)` turns an edge of the original graph `G`
|
||||
# into a node. This is illustrated with the graph below (taken from the
|
||||
# research paper).
|
||||
#
|
||||
# .. figure:: https://i.imgur.com/4WO5jEm.png
|
||||
# :alt: lg
|
||||
# :align: center
|
||||
#
|
||||
# Here, :math:`e_{A}:= （i\rightarrow j）` and :math:`e_{B}:= (j\rightarrow k)`
|
||||
# are two edges in the original graph :math:`G`. In line graph :math:`G_L`,
|
||||
# they correspond to nodes :math:`v^{l}_{A}, v^{l}_{B}`.
|
||||
#
|
||||
# The next natural question is, how to connect nodes in line-graph？ How to
|
||||
# connect two edges? Here, we use the following connection rule:
|
||||
#
|
||||
# Two nodes :math:`v^{l}_{A}`, :math:`v^{l}_{B}` in `lg` are connected if
|
||||
# the corresponding two edges :math:`e_{A}, e_{B}` in `g` share one and only
|
||||
# one node:
|
||||
# :math:`e_{A}`'s destination node is :math:`e_{B}`'s source node
|
||||
# (:math:`j`).
|
||||
#
|
||||
# .. note::
|
||||
#
|
||||
# Mathematically, this definition corresponds to a notion called non-backtracking
|
||||
# operator:
|
||||
# :math:`B_{(i \rightarrow j), (\hat{i} \rightarrow \hat{j})}`
|
||||
# :math:`= \begin{cases}
|
||||
# 1 \text{ if } j = \hat{i}, \hat{j} \neq i\\
|
||||
# 0 \text{ otherwise} \end{cases}`
|
||||
# where an edge is formed if :math:`B_{node1, node2} = 1`.
|
||||
#
|
||||
#
|
||||
# One layer in LGNN, algorithm structure
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
#
|
||||
# LGNN chains together a series of line graph neural network layers. The graph
|
||||
# representation :math:`x` and its line graph companion :math:`y` evolve with
|
||||
# the dataflow as follows.
|
||||
#
|
||||
# .. figure:: https://i.imgur.com/bZGGIGp.png
|
||||
# :alt: alg
|
||||
# :align: center
|
||||
#
|
||||
# At the :math:`k`-th layer, the :math:`i`-th neuron of the :math:`l`-th
|
||||
# channel updates its embedding :math:`x^{(k+1)}_{i,l}` with:
|
||||
#
|
||||
# .. math::
|
||||
# \begin{split}
|
||||
# x^{(k+1)}_{i,l} ={}&\rho[x^{(k)}_{i}\theta^{(k)}_{1,l}
|
||||
# +(Dx^{(k)})_{i}\theta^{(k)}_{2,l} \\
|
||||
# &+\sum^{J-1}_{j=0}(A^{2^{j}}x^{k})_{i}\theta^{(k)}_{3+j,l}\\
|
||||
# &+[\{\text{Pm},\text{Pd}\}y^{(k)}]_{i}\theta^{(k)}_{3+J,l}] \\
|
||||
# &+\text{skip-connection}
|
||||
# \qquad i \in V, l = 1,2,3, ... b_{k+1}/2
|
||||
# \end{split}
|
||||
#
|
||||
# Then, the line-graph representation :math:`y^{(k+1)}_{i,l}` with,
|
||||
#
|
||||
# .. math::
|
||||
#
|
||||
# \begin{split}
|
||||
# y^{(k+1)}_{i',l^{'}} = {}&\rho[y^{(k)}_{i^{'}}\gamma^{(k)}_{1,l^{'}}+
|
||||
# (D_{L(G)}y^{(k)})_{i^{'}}\gamma^{(k)}_{2,l^{'}}\\
|
||||
# &+\sum^{J-1}_{j=0}(A_{L(G)}^{2^{j}}y^{k})_{i}\gamma^{(k)}_{3+j,l^{'}}\\
|
||||
# &+[\{\text{Pm},\text{Pd}\}^{T}x^{(k+1)}]_{i^{'}}\gamma^{(k)}_{3+J,l^{'}}]\\
|
||||
# &+\text{skip-connection}
|
||||
# \qquad i^{'} \in V_{l}, l^{'} = 1,2,3, ... b^{'}_{k+1}/2
|
||||
# \end{split}
|
||||
#
|
||||
# Where :math:`\text{skip-connection}` refers to performing the same operation without the non-linearity
|
||||
# :math:`\rho`, and with linear projection :math:`\theta_\{\frac{b_{k+1}}{2} + 1, ..., b_{k+1}-1, b_{k+1}\}`
|
||||
# and :math:`\gamma_\{\frac{b_{k+1}}{2} + 1, ..., b_{k+1}-1, b_{k+1}\}`.
|
||||
#
|
||||
# Implement LGNN in DGL
|
||||
# ---------------------
|
||||
# Even though the equations in the previous section might seem intimidating,
|
||||
# it helps to understand the following information before you implement the LGNN.
|
||||
#
|
||||
# The two equations are symmetric and can be implemented as two instances
|
||||
# of the same class with different parameters.
|
||||
# The first equation operates on graph representation :math:`x`,
|
||||
# whereas the second operates on line-graph
|
||||
# representation :math:`y`. Let us denote this abstraction as :math:`f`. Then
|
||||
# the first is :math:`f(x,y; \theta_x)`, and the second
|
||||
# is :math:`f(y,x, \theta_y)`. That is, they are parameterized to compute
|
||||
# representations of the original graph and its
|
||||
# companion line graph, respectively.
|
||||
#
|
||||
# Each equation consists of four terms. Take the first one as an example, which follows.
|
||||
#
|
||||
# - :math:`x^{(k)}\theta^{(k)}_{1,l}`, a linear projection of previous
|
||||
# layer's output :math:`x^{(k)}`, denote as :math:`\text{prev}(x)`.
|
||||
# - :math:`(Dx^{(k)})\theta^{(k)}_{2,l}`, a linear projection of degree
|
||||
# operator on :math:`x^{(k)}`, denote as :math:`\text{deg}(x)`.
|
||||
# - :math:`\sum^{J-1}_{j=0}(A^{2^{j}}x^{(k)})\theta^{(k)}_{3+j,l}`,
|
||||
# a summation of :math:`2^{j}` adjacency operator on :math:`x^{(k)}`,
|
||||
# denote as :math:`\text{radius}(x)`
|
||||
# - :math:`[\{Pm,Pd\}y^{(k)}]\theta^{(k)}_{3+J,l}`, fusing another
|
||||
# graph's embedding information using incidence matrix
|
||||
# :math:`\{Pm, Pd\}`, followed with a linear projection,
|
||||
# denote as :math:`\text{fuse}(y)`.
|
||||
#
|
||||
# Each of the terms are performed again with different
|
||||
# parameters, and without the nonlinearity after the sum.
|
||||
# Therefore, :math:`f` could be written as:
|
||||
#
|
||||
# .. math::
|
||||
# \begin{split}
|
||||
# f(x^{(k)},y^{(k)}) = {}\rho[&\text{prev}(x^{(k-1)}) + \text{deg}(x^{(k-1)}) +\text{radius}(x^{k-1})
|
||||
# +\text{fuse}(y^{(k)})]\\
|
||||
# +&\text{prev}(x^{(k-1)}) + \text{deg}(x^{(k-1)}) +\text{radius}(x^{k-1}) +\text{fuse}(y^{(k)})
|
||||
# \end{split}
|
||||
#
|
||||
# Two equations are chained-up in the following order:
|
||||
#
|
||||
# .. math::
|
||||
# \begin{split}
|
||||
# x^{(k+1)} = {}& f(x^{(k)}, y^{(k)})\\
|
||||
# y^{(k+1)} = {}& f(y^{(k)}, x^{(k+1)})
|
||||
# \end{split}
|
||||
#
|
||||
# Keep in mind the listed observations in this overview and proceed to implementation.
|
||||
# An important point is that you use different strategies for the noted terms.
|
||||
#
|
||||
# .. note::
|
||||
# You can understand :math:`\{Pm, Pd\}` more thoroughly with this explanation.
|
||||
# Roughly speaking, there is a relationship between how :math:`g` and
|
||||
# :math:`lg` (the line graph) work together with loopy brief propagation.
|
||||
# Here, you implement :math:`\{Pm, Pd\}` as a SciPy COO sparse matrix in the dataset,
|
||||
# and stack them as tensors when batching. Another batching solution is to
|
||||
# treat :math:`\{Pm, Pd\}` as the adjacency matrix of a bipartite graph, which maps
|
||||
# line graph's feature to graph's, and vice versa.
|
||||
#
|
||||
# Implementing :math:`\text{prev}` and :math:`\text{deg}` as tensor operation
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# Linear projection and degree operation are both simply matrix
|
||||
# multiplication. Write them as PyTorch tensor operations.
|
||||
#
|
||||
# In ``__init__``, you define the projection variables.
|
||||
#
|
||||
# ::
|
||||
#
|
||||
# self.linear_prev = nn.Linear(in_feats, out_feats)
|
||||
# self.linear_deg = nn.Linear(in_feats, out_feats)
|
||||
#
|
||||
#
|
||||
# In ``forward()``, :math:`\text{prev}` and :math:`\text{deg}` are the same
|
||||
# as any other PyTorch tensor operations.
|
||||
#
|
||||
# ::
|
||||
#
|
||||
# prev_proj = self.linear_prev(feat_a)
|
||||
# deg_proj = self.linear_deg(deg * feat_a)
|
||||
#
|
||||
# Implementing :math:`\text{radius}` as message passing in DGL
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# As discussed in GCN tutorial, you can formulate one adjacency operator as
|
||||
# doing one-step message passing. As a generalization, :math:`2^j` adjacency
|
||||
# operations can be formulated as performing :math:`2^j` step of message
|
||||
# passing. Therefore, the summation is equivalent to summing nodes'
|
||||
# representation of :math:`2^j, j=0, 1, 2..` step message passing, i.e.
|
||||
# gathering information in :math:`2^{j}` neighborhood of each node.
|
||||
#
|
||||
# In ``__init__``, define the projection variables used in each
|
||||
# :math:`2^j` steps of message passing.
|
||||
#
|
||||
# ::
|
||||
#
|
||||
# self.linear_radius = nn.ModuleList(
|
||||
# [nn.Linear(in_feats, out_feats) for i in range(radius)])
|
||||
#
|
||||
# In ``__forward__``, use following function ``aggregate_radius()`` to
|
||||
# gather data from multiple hops. This can be seen in the following code.
|
||||
# Note that the ``update_all`` is called multiple times.
|
||||
|
||||
# Return a list containing features gathered from multiple radius.
|
||||
import dgl.function as fn
|
||||
|
||||
|
||||
def aggregate_radius(radius, g, z):
|
||||
# initializing list to collect message passing result
|
||||
z_list = []
|
||||
g.ndata["z"] = z
|
||||
# pulling message from 1-hop neighbourhood
|
||||
g.update_all(fn.copy_u(u="z", out="m"), fn.sum(msg="m", out="z"))
|
||||
z_list.append(g.ndata["z"])
|
||||
for i in range(radius - 1):
|
||||
for j in range(2**i):
|
||||
# pulling message from 2^j neighborhood
|
||||
g.update_all(fn.copy_u(u="z", out="m"), fn.sum(msg="m", out="z"))
|
||||
z_list.append(g.ndata["z"])
|
||||
return z_list
|
||||
|
||||
|
||||
#########################################################################
|
||||
# Implementing :math:`\text{fuse}` as sparse matrix multiplication
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# :math:`\{Pm, Pd\}` is a sparse matrix with only two non-zero entries on
|
||||
# each column. Therefore, you construct it as a sparse matrix in the dataset,
|
||||
# and implement :math:`\text{fuse}` as a sparse matrix multiplication.
|
||||
#
|
||||
# in ``__forward__``:
|
||||
#
|
||||
# ::
|
||||
#
|
||||
# fuse = self.linear_fuse(th.mm(pm_pd, feat_b))
|
||||
#
|
||||
# Completing :math:`f(x, y)`
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# Finally, the following shows how to sum up all the terms together, pass it to skip connection, and
|
||||
# batch norm.
|
||||
#
|
||||
# ::
|
||||
#
|
||||
# result = prev_proj + deg_proj + radius_proj + fuse
|
||||
#
|
||||
# Pass result to skip connection.
|
||||
#
|
||||
# ::
|
||||
#
|
||||
# result = th.cat([result[:, :n], F.relu(result[:, n:])], 1)
|
||||
#
|
||||
# Then pass the result to batch norm.
|
||||
#
|
||||
# ::
|
||||
#
|
||||
# result = self.bn(result) #Batch Normalization.
|
||||
#
|
||||
#
|
||||
# Here is the complete code for one LGNN layer's abstraction :math:`f(x,y)`
|
||||
class LGNNCore(nn.Module):
|
||||
def __init__(self, in_feats, out_feats, radius):
|
||||
super(LGNNCore, self).__init__()
|
||||
self.out_feats = out_feats
|
||||
self.radius = radius
|
||||
|
||||
self.linear_prev = nn.Linear(in_feats, out_feats)
|
||||
self.linear_deg = nn.Linear(in_feats, out_feats)
|
||||
self.linear_radius = nn.ModuleList(
|
||||
[nn.Linear(in_feats, out_feats) for i in range(radius)]
|
||||
)
|
||||
self.linear_fuse = nn.Linear(in_feats, out_feats)
|
||||
self.bn = nn.BatchNorm1d(out_feats)
|
||||
|
||||
def forward(self, g, feat_a, feat_b, deg, pm_pd):
|
||||
# term "prev"
|
||||
prev_proj = self.linear_prev(feat_a)
|
||||
# term "deg"
|
||||
deg_proj = self.linear_deg(deg * feat_a)
|
||||
|
||||
# term "radius"
|
||||
# aggregate 2^j-hop features
|
||||
hop2j_list = aggregate_radius(self.radius, g, feat_a)
|
||||
# apply linear transformation
|
||||
hop2j_list = [
|
||||
linear(x) for linear, x in zip(self.linear_radius, hop2j_list)
|
||||
]
|
||||
radius_proj = sum(hop2j_list)
|
||||
|
||||
# term "fuse"
|
||||
fuse = self.linear_fuse(th.mm(pm_pd, feat_b))
|
||||
|
||||
# sum them together
|
||||
result = prev_proj + deg_proj + radius_proj + fuse
|
||||
|
||||
# skip connection and batch norm
|
||||
n = self.out_feats // 2
|
||||
result = th.cat([result[:, :n], F.relu(result[:, n:])], 1)
|
||||
result = self.bn(result)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
##############################################################################################################
|
||||
# Chain-up LGNN abstractions as an LGNN layer
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# To implement:
|
||||
#
|
||||
# .. math::
|
||||
# \begin{split}
|
||||
# x^{(k+1)} = {}& f(x^{(k)}, y^{(k)})\\
|
||||
# y^{(k+1)} = {}& f(y^{(k)}, x^{(k+1)})
|
||||
# \end{split}
|
||||
#
|
||||
# Chain-up two ``LGNNCore`` instances, as in the example code, with different parameters in the forward pass.
|
||||
class LGNNLayer(nn.Module):
|
||||
def __init__(self, in_feats, out_feats, radius):
|
||||
super(LGNNLayer, self).__init__()
|
||||
self.g_layer = LGNNCore(in_feats, out_feats, radius)
|
||||
self.lg_layer = LGNNCore(in_feats, out_feats, radius)
|
||||
|
||||
def forward(self, g, lg, x, lg_x, deg_g, deg_lg, pm_pd):
|
||||
next_x = self.g_layer(g, x, lg_x, deg_g, pm_pd)
|
||||
pm_pd_y = th.transpose(pm_pd, 0, 1)
|
||||
next_lg_x = self.lg_layer(lg, lg_x, x, deg_lg, pm_pd_y)
|
||||
return next_x, next_lg_x
|
||||
|
||||
|
||||
########################################################################################
|
||||
# Chain-up LGNN layers
|
||||
# ~~~~~~~~~~~~~~~~~~~~
|
||||
# Define an LGNN with three hidden layers, as in the following example.
|
||||
class LGNN(nn.Module):
|
||||
def __init__(self, radius):
|
||||
super(LGNN, self).__init__()
|
||||
self.layer1 = LGNNLayer(1, 16, radius) # input is scalar feature
|
||||
self.layer2 = LGNNLayer(16, 16, radius) # hidden size is 16
|
||||
self.layer3 = LGNNLayer(16, 16, radius)
|
||||
self.linear = nn.Linear(16, 2) # predice two classes
|
||||
|
||||
def forward(self, g, lg, pm_pd):
|
||||
# compute the degrees
|
||||
deg_g = g.in_degrees().float().unsqueeze(1)
|
||||
deg_lg = lg.in_degrees().float().unsqueeze(1)
|
||||
# use degree as the input feature
|
||||
x, lg_x = deg_g, deg_lg
|
||||
x, lg_x = self.layer1(g, lg, x, lg_x, deg_g, deg_lg, pm_pd)
|
||||
x, lg_x = self.layer2(g, lg, x, lg_x, deg_g, deg_lg, pm_pd)
|
||||
x, lg_x = self.layer3(g, lg, x, lg_x, deg_g, deg_lg, pm_pd)
|
||||
return self.linear(x)
|
||||
|
||||
|
||||
#########################################################################################
|
||||
# Training and inference
|
||||
# -----------------------
|
||||
# First load the data.
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
training_loader = DataLoader(
|
||||
train_set, batch_size=1, collate_fn=train_set.collate_fn, drop_last=True
|
||||
)
|
||||
|
||||
#######################################################################################
|
||||
# Next, define the main training loop. Note that each training sample contains
|
||||
# three objects: A :class:`~dgl.DGLGraph`, a SciPy sparse matrix ``pmpd``, and a label
|
||||
# array in ``numpy.ndarray``. Generate the line graph by using this command:
|
||||
#
|
||||
# ::
|
||||
#
|
||||
# lg = g.line_graph(backtracking=False)
|
||||
#
|
||||
# Note that ``backtracking=False`` is required to correctly simulate non-backtracking
|
||||
# operation. We also define a utility function to convert the SciPy sparse matrix to
|
||||
# torch sparse tensor.
|
||||
|
||||
# Create the model
|
||||
model = LGNN(radius=3)
|
||||
# define the optimizer
|
||||
optimizer = th.optim.Adam(model.parameters(), lr=1e-2)
|
||||
|
||||
# A utility function to convert a scipy.coo_matrix to torch.SparseFloat
|
||||
def sparse2th(mat):
|
||||
value = mat.data
|
||||
indices = th.LongTensor([mat.row, mat.col])
|
||||
tensor = th.sparse.FloatTensor(
|
||||
indices, th.from_numpy(value).float(), mat.shape
|
||||
)
|
||||
return tensor
|
||||
|
||||
|
||||
# Train for 20 epochs
|
||||
for i in range(20):
|
||||
all_loss = []
|
||||
all_acc = []
|
||||
for [g, pmpd, label] in training_loader:
|
||||
# Generate the line graph.
|
||||
lg = g.line_graph(backtracking=False)
|
||||
# Create torch tensors
|
||||
pmpd = sparse2th(pmpd)
|
||||
label = th.from_numpy(label)
|
||||
|
||||
# Forward
|
||||
z = model(g, lg, pmpd)
|
||||
|
||||
# Calculate loss:
|
||||
# Since there are only two communities, there are only two permutations
|
||||
# of the community labels.
|
||||
loss_perm1 = F.cross_entropy(z, label)
|
||||
loss_perm2 = F.cross_entropy(z, 1 - label)
|
||||
loss = th.min(loss_perm1, loss_perm2)
|
||||
|
||||
# Calculate accuracy:
|
||||
_, pred = th.max(z, 1)
|
||||
acc_perm1 = (pred == label).float().mean()
|
||||
acc_perm2 = (pred == 1 - label).float().mean()
|
||||
acc = th.max(acc_perm1, acc_perm2)
|
||||
all_loss.append(loss.item())
|
||||
all_acc.append(acc.item())
|
||||
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
niters = len(all_loss)
|
||||
print(
|
||||
"Epoch %d | loss %.4f | accuracy %.4f"
|
||||
% (i, sum(all_loss) / niters, sum(all_acc) / niters)
|
||||
)
|
||||
#######################################################################################
|
||||
# Visualize training progress
|
||||
# -----------------------------
|
||||
# You can visualize the network's community prediction on one training example,
|
||||
# together with the ground truth. Start this with the following code example.
|
||||
|
||||
pmpd1 = sparse2th(pmpd1)
|
||||
LG1 = G1.line_graph(backtracking=False)
|
||||
z = model(G1, LG1, pmpd1)
|
||||
_, pred = th.max(z, 1)
|
||||
visualize(pred, nx_G1)
|
||||
|
||||
#######################################################################################
|
||||
# Compared with the ground truth. Note that the color might be reversed for the
|
||||
# two communities because the model is for correctly predicting the partitioning.
|
||||
visualize(label1, nx_G1)
|
||||
|
||||
#########################################
|
||||
# Here is an animation to better understand the process. (40 epochs)
|
||||
#
|
||||
# .. figure:: https://i.imgur.com/KDUyE1S.gif
|
||||
# :alt: lgnn-anim
|
||||
#
|
||||
# Batching graphs for parallelism
|
||||
# --------------------------------
|
||||
#
|
||||
# LGNN takes a collection of different graphs.
|
||||
# You might consider whether batching can be used for parallelism.
|
||||
#
|
||||
# Batching has been into the data loader itself.
|
||||
# In the ``collate_fn`` for PyTorch data loader, graphs are batched using DGL's
|
||||
# batched_graph API. DGL batches graphs by merging them
|
||||
# into a large graph, with each smaller graph's adjacency matrix being a block
|
||||
# along the diagonal of the large graph's adjacency matrix. Concatenate
|
||||
# :math`\{Pm,Pd\}` as block diagonal matrix in correspondence to DGL batched
|
||||
# graph API.
|
||||
|
||||
|
||||
def collate_fn(batch):
|
||||
graphs, pmpds, labels = zip(*batch)
|
||||
batched_graphs = dgl.batch(graphs)
|
||||
batched_pmpds = sp.block_diag(pmpds)
|
||||
batched_labels = np.concatenate(labels, axis=0)
|
||||
return batched_graphs, batched_pmpds, batched_labels
|
||||
|
||||
|
||||
######################################################################################
|
||||
# You can find the complete code on Github at
|
||||
# `Community Detection with Graph Neural Networks (CDGNN) <https://github.com/dmlc/dgl/tree/master/examples/pytorch/line_graph>`_.
|
||||
@@ -0,0 +1,545 @@
|
||||
"""
|
||||
.. _model-gat:
|
||||
|
||||
Understand Graph Attention Network
|
||||
=======================================
|
||||
|
||||
**Authors:** `Hao Zhang <https://github.com/sufeidechabei/>`_, `Mufei Li
|
||||
<https://github.com/mufeili>`_, `Minjie Wang
|
||||
<https://jermainewang.github.io/>`_ `Zheng Zhang
|
||||
<https://shanghai.nyu.edu/academics/faculty/directory/zheng-zhang>`_
|
||||
|
||||
.. warning::
|
||||
|
||||
The tutorial aims at gaining insights into the paper, with code as a mean
|
||||
of explanation. The implementation thus is NOT optimized for running
|
||||
efficiency. For recommended implementation, please refer to the `official
|
||||
examples <https://github.com/dmlc/dgl/tree/master/examples>`_.
|
||||
|
||||
In this tutorial, you learn about a graph attention network (GAT) and how it can be
|
||||
implemented in PyTorch. You can also learn to visualize and understand what the attention
|
||||
mechanism has learned.
|
||||
|
||||
The research described in the paper `Graph Convolutional Network (GCN) <https://arxiv.org/abs/1609.02907>`_,
|
||||
indicates that combining local graph structure and node-level features yields
|
||||
good performance on node classification tasks. However, the way GCN aggregates
|
||||
is structure-dependent, which can hurt its generalizability.
|
||||
|
||||
One workaround is to simply average over all neighbor node features as described in
|
||||
the research paper `GraphSAGE
|
||||
<https://www-cs-faculty.stanford.edu/people/jure/pubs/graphsage-nips17.pdf>`_.
|
||||
However, `Graph Attention Network <https://arxiv.org/abs/1710.10903>`_ proposes a
|
||||
different type of aggregation. GAT uses weighting neighbor features with feature dependent and
|
||||
structure-free normalization, in the style of attention.
|
||||
"""
|
||||
###############################################################
|
||||
# Introducing attention to GCN
|
||||
# ----------------------------
|
||||
#
|
||||
# The key difference between GAT and GCN is how the information from the one-hop neighborhood is aggregated.
|
||||
#
|
||||
# For GCN, a graph convolution operation produces the normalized sum of the node features of neighbors.
|
||||
#
|
||||
#
|
||||
# .. math::
|
||||
#
|
||||
# h_i^{(l+1)}=\sigma\left(\sum_{j\in \mathcal{N}(i)} {\frac{1}{c_{ij}} W^{(l)}h^{(l)}_j}\right)
|
||||
#
|
||||
#
|
||||
# where :math:`\mathcal{N}(i)` is the set of its one-hop neighbors (to include
|
||||
# :math:`v_i` in the set, simply add a self-loop to each node),
|
||||
# :math:`c_{ij}=\sqrt{|\mathcal{N}(i)|}\sqrt{|\mathcal{N}(j)|}` is a
|
||||
# normalization constant based on graph structure, :math:`\sigma` is an
|
||||
# activation function (GCN uses ReLU), and :math:`W^{(l)}` is a shared
|
||||
# weight matrix for node-wise feature transformation. Another model proposed in
|
||||
# `GraphSAGE
|
||||
# <https://www-cs-faculty.stanford.edu/people/jure/pubs/graphsage-nips17.pdf>`_
|
||||
# employs the same update rule except that they set
|
||||
# :math:`c_{ij}=|\mathcal{N}(i)|`.
|
||||
#
|
||||
# GAT introduces the attention mechanism as a substitute for the statically
|
||||
# normalized convolution operation. Below are the equations to compute the node
|
||||
# embedding :math:`h_i^{(l+1)}` of layer :math:`l+1` from the embeddings of
|
||||
# layer :math:`l`.
|
||||
#
|
||||
# .. image:: https://data.dgl.ai/tutorial/gat/gat.png
|
||||
# :width: 450px
|
||||
# :align: center
|
||||
#
|
||||
# .. math::
|
||||
#
|
||||
# \begin{align}
|
||||
# z_i^{(l)}&=W^{(l)}h_i^{(l)},&(1) \\
|
||||
# e_{ij}^{(l)}&=\text{LeakyReLU}(\vec a^{(l)^T}(z_i^{(l)}||z_j^{(l)})),&(2)\\
|
||||
# \alpha_{ij}^{(l)}&=\frac{\exp(e_{ij}^{(l)})}{\sum_{k\in \mathcal{N}(i)}^{}\exp(e_{ik}^{(l)})},&(3)\\
|
||||
# h_i^{(l+1)}&=\sigma\left(\sum_{j\in \mathcal{N}(i)} {\alpha^{(l)}_{ij} z^{(l)}_j }\right),&(4)
|
||||
# \end{align}
|
||||
#
|
||||
#
|
||||
# Explanations:
|
||||
#
|
||||
#
|
||||
# * Equation (1) is a linear transformation of the lower layer embedding :math:`h_i^{(l)}`
|
||||
# and :math:`W^{(l)}` is its learnable weight matrix.
|
||||
# * Equation (2) computes a pair-wise *un-normalized* attention score between two neighbors.
|
||||
# Here, it first concatenates the :math:`z` embeddings of the two nodes, where :math:`||`
|
||||
# denotes concatenation, then takes a dot product of it and a learnable weight vector
|
||||
# :math:`\vec a^{(l)}`, and applies a LeakyReLU in the end. This form of attention is
|
||||
# usually called *additive attention*, contrast with the dot-product attention in the
|
||||
# Transformer model.
|
||||
# * Equation (3) applies a softmax to normalize the attention scores on each node's
|
||||
# incoming edges.
|
||||
# * Equation (4) is similar to GCN. The embeddings from neighbors are aggregated together,
|
||||
# scaled by the attention scores.
|
||||
#
|
||||
# There are other details from the paper, such as dropout and skip connections.
|
||||
# For the purpose of simplicity, those details are left out of this tutorial. To see more details,
|
||||
# download the `full example <https://github.com/dmlc/dgl/blob/master/examples/pytorch/gat/gat.py>`_.
|
||||
# In its essence, GAT is just a different aggregation function with attention
|
||||
# over features of neighbors, instead of a simple mean aggregation.
|
||||
#
|
||||
# GAT in DGL
|
||||
# ----------
|
||||
#
|
||||
# DGL provides an off-the-shelf implementation of the GAT layer under the ``dgl.nn.<backend>``
|
||||
# subpackage. Simply import the ``GATConv`` as the follows.
|
||||
|
||||
import os
|
||||
|
||||
os.environ["DGLBACKEND"] = "pytorch"
|
||||
###############################################################
|
||||
# Readers can skip the following step-by-step explanation of the implementation and
|
||||
# jump to the `Put everything together`_ for training and visualization results.
|
||||
#
|
||||
# To begin, you can get an overall impression about how a ``GATLayer`` module is
|
||||
# implemented in DGL. In this section, the four equations above are broken down
|
||||
# one at a time.
|
||||
#
|
||||
# .. note::
|
||||
#
|
||||
# This is showing how to implement a GAT from scratch. DGL provides a more
|
||||
# efficient :class:`builtin GAT layer module <dgl.nn.pytorch.conv.GATConv>`.
|
||||
#
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from dgl.nn.pytorch import GATConv
|
||||
|
||||
|
||||
class GATLayer(nn.Module):
|
||||
def __init__(self, g, in_dim, out_dim):
|
||||
super(GATLayer, self).__init__()
|
||||
self.g = g
|
||||
# equation (1)
|
||||
self.fc = nn.Linear(in_dim, out_dim, bias=False)
|
||||
# equation (2)
|
||||
self.attn_fc = nn.Linear(2 * out_dim, 1, bias=False)
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self):
|
||||
"""Reinitialize learnable parameters."""
|
||||
gain = nn.init.calculate_gain("relu")
|
||||
nn.init.xavier_normal_(self.fc.weight, gain=gain)
|
||||
nn.init.xavier_normal_(self.attn_fc.weight, gain=gain)
|
||||
|
||||
def edge_attention(self, edges):
|
||||
# edge UDF for equation (2)
|
||||
z2 = torch.cat([edges.src["z"], edges.dst["z"]], dim=1)
|
||||
a = self.attn_fc(z2)
|
||||
return {"e": F.leaky_relu(a)}
|
||||
|
||||
def message_func(self, edges):
|
||||
# message UDF for equation (3) & (4)
|
||||
return {"z": edges.src["z"], "e": edges.data["e"]}
|
||||
|
||||
def reduce_func(self, nodes):
|
||||
# reduce UDF for equation (3) & (4)
|
||||
# equation (3)
|
||||
alpha = F.softmax(nodes.mailbox["e"], dim=1)
|
||||
# equation (4)
|
||||
h = torch.sum(alpha * nodes.mailbox["z"], dim=1)
|
||||
return {"h": h}
|
||||
|
||||
def forward(self, h):
|
||||
# equation (1)
|
||||
z = self.fc(h)
|
||||
self.g.ndata["z"] = z
|
||||
# equation (2)
|
||||
self.g.apply_edges(self.edge_attention)
|
||||
# equation (3) & (4)
|
||||
self.g.update_all(self.message_func, self.reduce_func)
|
||||
return self.g.ndata.pop("h")
|
||||
|
||||
|
||||
##################################################################
|
||||
# Equation (1)
|
||||
# ^^^^^^^^^^^^
|
||||
#
|
||||
# .. math::
|
||||
#
|
||||
# z_i^{(l)}=W^{(l)}h_i^{(l)},(1)
|
||||
#
|
||||
# The first one shows linear transformation. It's common and can be
|
||||
# easily implemented in Pytorch using ``torch.nn.Linear``.
|
||||
#
|
||||
# Equation (2)
|
||||
# ^^^^^^^^^^^^
|
||||
#
|
||||
# .. math::
|
||||
#
|
||||
# e_{ij}^{(l)}=\text{LeakyReLU}(\vec a^{(l)^T}(z_i^{(l)}|z_j^{(l)})),(2)
|
||||
#
|
||||
# The un-normalized attention score :math:`e_{ij}` is calculated using the
|
||||
# embeddings of adjacent nodes :math:`i` and :math:`j`. This suggests that the
|
||||
# attention scores can be viewed as edge data, which can be calculated by the
|
||||
# ``apply_edges`` API. The argument to the ``apply_edges`` is an **Edge UDF**,
|
||||
# which is defined as below:
|
||||
|
||||
|
||||
def edge_attention(self, edges):
|
||||
# edge UDF for equation (2)
|
||||
z2 = torch.cat([edges.src["z"], edges.dst["z"]], dim=1)
|
||||
a = self.attn_fc(z2)
|
||||
return {"e": F.leaky_relu(a)}
|
||||
|
||||
|
||||
########################################################################3
|
||||
# Here, the dot product with the learnable weight vector :math:`\vec{a^{(l)}}`
|
||||
# is implemented again using PyTorch's linear transformation ``attn_fc``. Note
|
||||
# that ``apply_edges`` will **batch** all the edge data in one tensor, so the
|
||||
# ``cat``, ``attn_fc`` here are applied on all the edges in parallel.
|
||||
#
|
||||
# Equation (3) & (4)
|
||||
# ^^^^^^^^^^^^^^^^^^
|
||||
#
|
||||
# .. math::
|
||||
#
|
||||
# \begin{align}
|
||||
# \alpha_{ij}^{(l)}&=\frac{\exp(e_{ij}^{(l)})}{\sum_{k\in \mathcal{N}(i)}^{}\exp(e_{ik}^{(l)})},&(3)\\
|
||||
# h_i^{(l+1)}&=\sigma\left(\sum_{j\in \mathcal{N}(i)} {\alpha^{(l)}_{ij} z^{(l)}_j }\right),&(4)
|
||||
# \end{align}
|
||||
#
|
||||
# Similar to GCN, ``update_all`` API is used to trigger message passing on all
|
||||
# the nodes. The message function sends out two tensors: the transformed ``z``
|
||||
# embedding of the source node and the un-normalized attention score ``e`` on
|
||||
# each edge. The reduce function then performs two tasks:
|
||||
#
|
||||
#
|
||||
# * Normalize the attention scores using softmax (equation (3)).
|
||||
# * Aggregate neighbor embeddings weighted by the attention scores (equation(4)).
|
||||
#
|
||||
# Both tasks first fetch data from the mailbox and then manipulate it on the
|
||||
# second dimension (``dim=1``), on which the messages are batched.
|
||||
|
||||
|
||||
def reduce_func(self, nodes):
|
||||
# reduce UDF for equation (3) & (4)
|
||||
# equation (3)
|
||||
alpha = F.softmax(nodes.mailbox["e"], dim=1)
|
||||
# equation (4)
|
||||
h = torch.sum(alpha * nodes.mailbox["z"], dim=1)
|
||||
return {"h": h}
|
||||
|
||||
|
||||
#####################################################################
|
||||
# Multi-head attention
|
||||
# ^^^^^^^^^^^^^^^^^^^^
|
||||
#
|
||||
# Analogous to multiple channels in ConvNet, GAT introduces **multi-head
|
||||
# attention** to enrich the model capacity and to stabilize the learning
|
||||
# process. Each attention head has its own parameters and their outputs can be
|
||||
# merged in two ways:
|
||||
#
|
||||
# .. math:: \text{concatenation}: h^{(l+1)}_{i} =||_{k=1}^{K}\sigma\left(\sum_{j\in \mathcal{N}(i)}\alpha_{ij}^{k}W^{k}h^{(l)}_{j}\right)
|
||||
#
|
||||
# or
|
||||
#
|
||||
# .. math:: \text{average}: h_{i}^{(l+1)}=\sigma\left(\frac{1}{K}\sum_{k=1}^{K}\sum_{j\in\mathcal{N}(i)}\alpha_{ij}^{k}W^{k}h^{(l)}_{j}\right)
|
||||
#
|
||||
# where :math:`K` is the number of heads. You can use
|
||||
# concatenation for intermediary layers and average for the final layer.
|
||||
#
|
||||
# Use the above defined single-head ``GATLayer`` as the building block
|
||||
# for the ``MultiHeadGATLayer`` below:
|
||||
|
||||
|
||||
class MultiHeadGATLayer(nn.Module):
|
||||
def __init__(self, g, in_dim, out_dim, num_heads, merge="cat"):
|
||||
super(MultiHeadGATLayer, self).__init__()
|
||||
self.heads = nn.ModuleList()
|
||||
for i in range(num_heads):
|
||||
self.heads.append(GATLayer(g, in_dim, out_dim))
|
||||
self.merge = merge
|
||||
|
||||
def forward(self, h):
|
||||
head_outs = [attn_head(h) for attn_head in self.heads]
|
||||
if self.merge == "cat":
|
||||
# concat on the output feature dimension (dim=1)
|
||||
return torch.cat(head_outs, dim=1)
|
||||
else:
|
||||
# merge using average
|
||||
return torch.mean(torch.stack(head_outs))
|
||||
|
||||
|
||||
###########################################################################
|
||||
# Put everything together
|
||||
# ^^^^^^^^^^^^^^^^^^^^^^^
|
||||
#
|
||||
# Now, you can define a two-layer GAT model.
|
||||
|
||||
|
||||
class GAT(nn.Module):
|
||||
def __init__(self, g, in_dim, hidden_dim, out_dim, num_heads):
|
||||
super(GAT, self).__init__()
|
||||
self.layer1 = MultiHeadGATLayer(g, in_dim, hidden_dim, num_heads)
|
||||
# Be aware that the input dimension is hidden_dim*num_heads since
|
||||
# multiple head outputs are concatenated together. Also, only
|
||||
# one attention head in the output layer.
|
||||
self.layer2 = MultiHeadGATLayer(g, hidden_dim * num_heads, out_dim, 1)
|
||||
|
||||
def forward(self, h):
|
||||
h = self.layer1(h)
|
||||
h = F.elu(h)
|
||||
h = self.layer2(h)
|
||||
return h
|
||||
|
||||
|
||||
import networkx as nx
|
||||
|
||||
#############################################################################
|
||||
# We then load the Cora dataset using DGL's built-in data module.
|
||||
|
||||
from dgl import DGLGraph
|
||||
from dgl.data import citation_graph as citegrh
|
||||
|
||||
|
||||
def load_cora_data():
|
||||
data = citegrh.load_cora()
|
||||
g = data[0]
|
||||
mask = torch.BoolTensor(g.ndata["train_mask"])
|
||||
return g, g.ndata["feat"], g.ndata["label"], mask
|
||||
|
||||
|
||||
##############################################################################
|
||||
# The training loop is exactly the same as in the GCN tutorial.
|
||||
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
|
||||
g, features, labels, mask = load_cora_data()
|
||||
|
||||
# create the model, 2 heads, each head has hidden size 8
|
||||
net = GAT(g, in_dim=features.size()[1], hidden_dim=8, out_dim=7, num_heads=2)
|
||||
|
||||
# create optimizer
|
||||
optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)
|
||||
|
||||
# main loop
|
||||
dur = []
|
||||
for epoch in range(30):
|
||||
if epoch >= 3:
|
||||
t0 = time.time()
|
||||
|
||||
logits = net(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}".format(
|
||||
epoch, loss.item(), np.mean(dur)
|
||||
)
|
||||
)
|
||||
|
||||
#########################################################################
|
||||
# Visualizing and understanding attention learned
|
||||
# ----------------------------------------------
|
||||
#
|
||||
# Cora
|
||||
# ^^^^
|
||||
#
|
||||
# The following table summarizes the model performance on Cora that is reported in
|
||||
# `the GAT paper <https://arxiv.org/pdf/1710.10903.pdf>`_ and obtained with DGL
|
||||
# implementations.
|
||||
#
|
||||
# .. list-table::
|
||||
# :header-rows: 1
|
||||
#
|
||||
# * - Model
|
||||
# - Accuracy
|
||||
# * - GCN (paper)
|
||||
# - :math:`81.4\pm 0.5%`
|
||||
# * - GCN (dgl)
|
||||
# - :math:`82.05\pm 0.33%`
|
||||
# * - GAT (paper)
|
||||
# - :math:`83.0\pm 0.7%`
|
||||
# * - GAT (dgl)
|
||||
# - :math:`83.69\pm 0.529%`
|
||||
#
|
||||
# *What kind of attention distribution has our model learned?*
|
||||
#
|
||||
# Because the attention weight :math:`a_{ij}` is associated with edges, you can
|
||||
# visualize it by coloring edges. Below you can pick a subgraph of Cora and plot the
|
||||
# attention weights of the last ``GATLayer``. The nodes are colored according
|
||||
# to their labels, whereas the edges are colored according to the magnitude of
|
||||
# the attention weights, which can be referred with the colorbar on the right.
|
||||
#
|
||||
# .. image:: https://data.dgl.ai/tutorial/gat/cora-attention.png
|
||||
# :width: 600px
|
||||
# :align: center
|
||||
#
|
||||
# You can see that the model seems to learn different attention weights. To
|
||||
# understand the distribution more thoroughly, measure the `entropy
|
||||
# <https://en.wikipedia.org/wiki/Entropy_(information_theory>`_) of the
|
||||
# attention distribution. For any node :math:`i`,
|
||||
# :math:`\{\alpha_{ij}\}_{j\in\mathcal{N}(i)}` forms a discrete probability
|
||||
# distribution over all its neighbors with the entropy given by
|
||||
#
|
||||
# .. math:: H({\alpha_{ij}}_{j\in\mathcal{N}(i)})=-\sum_{j\in\mathcal{N}(i)} \alpha_{ij}\log\alpha_{ij}
|
||||
#
|
||||
# A low entropy means a high degree of concentration, and vice
|
||||
# versa. An entropy of 0 means all attention is on one source node. The uniform
|
||||
# distribution has the highest entropy of :math:`\log(\mathcal{N}(i))`.
|
||||
# Ideally, you want to see the model learns a distribution of lower entropy
|
||||
# (i.e, one or two neighbors are much more important than the others).
|
||||
#
|
||||
# Note that since nodes can have different degrees, the maximum entropy will
|
||||
# also be different. Therefore, you plot the aggregated histogram of entropy
|
||||
# values of all nodes in the entire graph. Below are the attention histogram of
|
||||
# learned by each attention head.
|
||||
#
|
||||
# |image2|
|
||||
#
|
||||
# As a reference, here is the histogram if all the nodes have uniform attention weight distribution.
|
||||
#
|
||||
# .. image:: https://data.dgl.ai/tutorial/gat/cora-attention-uniform-hist.png
|
||||
# :width: 250px
|
||||
# :align: center
|
||||
#
|
||||
# One can see that **the attention values learned is quite similar to uniform distribution**
|
||||
# (i.e, all neighbors are equally important). This partially
|
||||
# explains why the performance of GAT is close to that of GCN on Cora
|
||||
# (according to `author's reported result
|
||||
# <https://arxiv.org/pdf/1710.10903.pdf>`_, the accuracy difference averaged
|
||||
# over 100 runs is less than 2 percent). Attention does not matter
|
||||
# since it does not differentiate much.
|
||||
#
|
||||
# *Does that mean the attention mechanism is not useful?* No! A different
|
||||
# dataset exhibits an entirely different pattern, as you can see next.
|
||||
#
|
||||
# Protein-protein interaction (PPI) networks
|
||||
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
#
|
||||
# The PPI dataset used here consists of :math:`24` graphs corresponding to
|
||||
# different human tissues. Nodes can have up to :math:`121` kinds of labels, so
|
||||
# the label of node is represented as a binary tensor of size :math:`121`. The
|
||||
# task is to predict node label.
|
||||
#
|
||||
# Use :math:`20` graphs for training, :math:`2` for validation and :math:`2`
|
||||
# for test. The average number of nodes per graph is :math:`2372`. Each node
|
||||
# has :math:`50` features that are composed of positional gene sets, motif gene
|
||||
# sets, and immunological signatures. Critically, test graphs remain completely
|
||||
# unobserved during training, a setting called "inductive learning".
|
||||
#
|
||||
# Compare the performance of GAT and GCN for :math:`10` random runs on this
|
||||
# task and use hyperparameter search on the validation set to find the best
|
||||
# model.
|
||||
#
|
||||
# .. list-table::
|
||||
# :header-rows: 1
|
||||
#
|
||||
# * - Model
|
||||
# - F1 Score(micro)
|
||||
# * - GAT
|
||||
# - :math:`0.975 \pm 0.006`
|
||||
# * - GCN
|
||||
# - :math:`0.509 \pm 0.025`
|
||||
# * - Paper
|
||||
# - :math:`0.973 \pm 0.002`
|
||||
#
|
||||
# The table above is the result of this experiment, where you use micro `F1
|
||||
# score <https://en.wikipedia.org/wiki/F1_score>`_ to evaluate the model
|
||||
# performance.
|
||||
#
|
||||
# .. note::
|
||||
#
|
||||
# Below is the calculation process of F1 score:
|
||||
#
|
||||
# .. math::
|
||||
#
|
||||
# precision=\frac{\sum_{t=1}^{n}TP_{t}}{\sum_{t=1}^{n}(TP_{t} +FP_{t})}
|
||||
#
|
||||
# recall=\frac{\sum_{t=1}^{n}TP_{t}}{\sum_{t=1}^{n}(TP_{t} +FN_{t})}
|
||||
#
|
||||
# F1_{micro}=2\frac{precision*recall}{precision+recall}
|
||||
#
|
||||
# * :math:`TP_{t}` represents for number of nodes that both have and are predicted to have label :math:`t`
|
||||
# * :math:`FP_{t}` represents for number of nodes that do not have but are predicted to have label :math:`t`
|
||||
# * :math:`FN_{t}` represents for number of output classes labeled as :math:`t` but predicted as others.
|
||||
# * :math:`n` is the number of labels, i.e. :math:`121` in our case.
|
||||
#
|
||||
# During training, use ``BCEWithLogitsLoss`` as the loss function. The
|
||||
# learning curves of GAT and GCN are presented below; what is evident is the
|
||||
# dramatic performance adavantage of GAT over GCN.
|
||||
#
|
||||
# .. image:: https://data.dgl.ai/tutorial/gat/ppi-curve.png
|
||||
# :width: 300px
|
||||
# :align: center
|
||||
#
|
||||
# As before, you can have a statistical understanding of the attentions learned
|
||||
# by showing the histogram plot for the node-wise attention entropy. Below are
|
||||
# the attention histograms learned by different attention layers.
|
||||
#
|
||||
# *Attention learned in layer 1:*
|
||||
#
|
||||
# |image5|
|
||||
#
|
||||
# *Attention learned in layer 2:*
|
||||
#
|
||||
# |image6|
|
||||
#
|
||||
# *Attention learned in final layer:*
|
||||
#
|
||||
# |image7|
|
||||
#
|
||||
# Again, comparing with uniform distribution:
|
||||
#
|
||||
# .. image:: https://data.dgl.ai/tutorial/gat/ppi-uniform-hist.png
|
||||
# :width: 250px
|
||||
# :align: center
|
||||
#
|
||||
# Clearly, **GAT does learn sharp attention weights**! There is a clear pattern
|
||||
# over the layers as well: **the attention gets sharper with a higher
|
||||
# layer**.
|
||||
#
|
||||
# Unlike the Cora dataset where GAT's gain is minimal at best, for PPI there
|
||||
# is a significant performance gap between GAT and other GNN variants compared
|
||||
# in `the GAT paper <https://arxiv.org/pdf/1710.10903.pdf>`_ (at least 20 percent),
|
||||
# and the attention distributions between the two clearly differ. While this
|
||||
# deserves further research, one immediate conclusion is that GAT's advantage
|
||||
# lies perhaps more in its ability to handle a graph with more complex
|
||||
# neighborhood structure.
|
||||
#
|
||||
# What's next?
|
||||
# ------------
|
||||
#
|
||||
# So far, you have seen how to use DGL to implement GAT. There are some
|
||||
# missing details such as dropout, skip connections, and hyper-parameter tuning,
|
||||
# which are practices that do not involve DGL-related concepts. For more information
|
||||
# check out the full example.
|
||||
#
|
||||
# * See the optimized `full example <https://github.com/dmlc/dgl/blob/master/examples/pytorch/gat/gat.py>`_.
|
||||
# * The next tutorial describes how to speedup GAT models by parallelizing multiple attention heads and SPMV optimization.
|
||||
#
|
||||
# .. |image2| image:: https://data.dgl.ai/tutorial/gat/cora-attention-hist.png
|
||||
# .. |image5| image:: https://data.dgl.ai/tutorial/gat/ppi-first-layer-hist.png
|
||||
# .. |image6| image:: https://data.dgl.ai/tutorial/gat/ppi-second-layer-hist.png
|
||||
# .. |image7| image:: https://data.dgl.ai/tutorial/gat/ppi-final-layer-hist.png
|
||||
@@ -0,0 +1,37 @@
|
||||
.. _tutorials1-index:
|
||||
|
||||
Graph neural networks and its variants
|
||||
--------------------------------------------
|
||||
|
||||
* **Graph convolutional network (GCN)** `[research paper] <https://arxiv.org/abs/1609.02907>`__ `[tutorial]
|
||||
<1_gnn/1_gcn.html>`__ `[Pytorch code]
|
||||
<https://github.com/dmlc/dgl/blob/master/examples/pytorch/gcn>`__
|
||||
`[MXNet code]
|
||||
<https://github.com/dmlc/dgl/tree/master/examples/mxnet/gcn>`__:
|
||||
|
||||
* **Graph attention network (GAT)** `[research paper] <https://arxiv.org/abs/1710.10903>`__ `[tutorial]
|
||||
<1_gnn/9_gat.html>`__ `[Pytorch code]
|
||||
<https://github.com/dmlc/dgl/blob/master/examples/pytorch/gat>`__
|
||||
`[MXNet code]
|
||||
<https://github.com/dmlc/dgl/tree/master/examples/mxnet/gat>`__:
|
||||
GAT extends the GCN functionality by deploying multi-head attention
|
||||
among neighborhood of a node. This greatly enhances the capacity and
|
||||
expressiveness of the model.
|
||||
|
||||
* **Relational-GCN** `[research paper] <https://arxiv.org/abs/1703.06103>`__ `[tutorial]
|
||||
<1_gnn/4_rgcn.html>`__ `[Pytorch code]
|
||||
<https://github.com/dmlc/dgl/tree/master/examples/pytorch/rgcn>`__
|
||||
`[MXNet code]
|
||||
<https://github.com/dmlc/dgl/tree/master/examples/mxnet/rgcn>`__:
|
||||
Relational-GCN allows multiple edges among two entities of a
|
||||
graph. Edges with distinct relationships are encoded differently.
|
||||
|
||||
* **Line graph neural network (LGNN)** `[research paper] <https://openreview.net/pdf?id=H1g0Z3A9Fm>`__ `[tutorial]
|
||||
<1_gnn/6_line_graph.html>`__ `[Pytorch code]
|
||||
<https://github.com/dmlc/dgl/tree/master/examples/pytorch/line_graph>`__:
|
||||
This network focuses on community detection by inspecting graph structures. It
|
||||
uses representations of both the original graph and its line-graph
|
||||
companion. In addition to demonstrating how an algorithm can harness multiple
|
||||
graphs, this implementation shows how you can judiciously mix simple tensor
|
||||
operations and sparse-matrix tensor operations, along with message-passing with
|
||||
DGL.
|
||||
@@ -0,0 +1,422 @@
|
||||
"""
|
||||
.. _model-tree-lstm:
|
||||
|
||||
Tree-LSTM in DGL
|
||||
==========================
|
||||
|
||||
**Author**: Zihao Ye, Qipeng Guo, `Minjie Wang
|
||||
<https://jermainewang.github.io/>`_, `Jake Zhao
|
||||
<https://cs.nyu.edu/~jakezhao/>`_, Zheng Zhang
|
||||
|
||||
.. warning::
|
||||
|
||||
The tutorial aims at gaining insights into the paper, with code as a mean
|
||||
of explanation. The implementation thus is NOT optimized for running
|
||||
efficiency. For recommended implementation, please refer to the `official
|
||||
examples <https://github.com/dmlc/dgl/tree/master/examples>`_.
|
||||
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
##############################################################################
|
||||
#
|
||||
# In this tutorial, you learn to use Tree-LSTM networks for sentiment analysis.
|
||||
# The Tree-LSTM is a generalization of long short-term memory (LSTM) networks to tree-structured network topologies.
|
||||
#
|
||||
# The Tree-LSTM structure was first introduced by Kai et. al in an ACL 2015
|
||||
# paper: `Improved Semantic Representations From Tree-Structured Long
|
||||
# Short-Term Memory Networks <https://arxiv.org/pdf/1503.00075.pdf>`__.
|
||||
# The core idea is to introduce syntactic information for language tasks by
|
||||
# extending the chain-structured LSTM to a tree-structured LSTM. The dependency
|
||||
# tree and constituency tree techniques are leveraged to obtain a ''latent tree''.
|
||||
#
|
||||
# The challenge in training Tree-LSTMs is batching --- a standard
|
||||
# technique in machine learning to accelerate optimization. However, since trees
|
||||
# generally have different shapes by nature, parallization is non-trivial.
|
||||
# DGL offers an alternative. Pool all the trees into one single graph then
|
||||
# induce the message passing over them, guided by the structure of each tree.
|
||||
#
|
||||
# The task and the dataset
|
||||
# ------------------------
|
||||
#
|
||||
# The steps here use the
|
||||
# `Stanford Sentiment Treebank <https://nlp.stanford.edu/sentiment/>`__ in
|
||||
# ``dgl.data``. The dataset provides a fine-grained, tree-level sentiment
|
||||
# annotation. There are five classes: Very negative, negative, neutral, positive, and
|
||||
# very positive, which indicate the sentiment in the current subtree. Non-leaf
|
||||
# nodes in a constituency tree do not contain words, so use a special
|
||||
# ``PAD_WORD`` token to denote them. During training and inference
|
||||
# their embeddings would be masked to all-zero.
|
||||
#
|
||||
# .. figure:: https://i.loli.net/2018/11/08/5be3d4bfe031b.png
|
||||
# :alt:
|
||||
#
|
||||
# The figure displays one sample of the SST dataset, which is a
|
||||
# constituency parse tree with their nodes labeled with sentiment. To
|
||||
# speed up things, build a tiny set with five sentences and take a look
|
||||
# at the first one.
|
||||
#
|
||||
|
||||
from collections import namedtuple
|
||||
|
||||
os.environ["DGLBACKEND"] = "pytorch"
|
||||
import dgl
|
||||
from dgl.data.tree import SSTDataset
|
||||
|
||||
|
||||
SSTBatch = namedtuple("SSTBatch", ["graph", "mask", "wordid", "label"])
|
||||
|
||||
# Each sample in the dataset is a constituency tree. The leaf nodes
|
||||
# represent words. The word is an int value stored in the "x" field.
|
||||
# The non-leaf nodes have a special word PAD_WORD. The sentiment
|
||||
# label is stored in the "y" feature field.
|
||||
trainset = SSTDataset(mode="tiny") # the "tiny" set has only five trees
|
||||
tiny_sst = [tr for tr in trainset]
|
||||
num_vocabs = trainset.vocab_size
|
||||
num_classes = trainset.num_classes
|
||||
|
||||
vocab = trainset.vocab # vocabulary dict: key -> id
|
||||
inv_vocab = {
|
||||
v: k for k, v in vocab.items()
|
||||
} # inverted vocabulary dict: id -> word
|
||||
|
||||
a_tree = tiny_sst[0]
|
||||
for token in a_tree.ndata["x"].tolist():
|
||||
if token != trainset.PAD_WORD:
|
||||
print(inv_vocab[token], end=" ")
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
##############################################################################
|
||||
# Step 1: Batching
|
||||
# ----------------
|
||||
#
|
||||
# Add all the trees to one graph, using
|
||||
# the :func:`~dgl.batched_graph.batch` API.
|
||||
#
|
||||
|
||||
import networkx as nx
|
||||
|
||||
graph = dgl.batch(tiny_sst)
|
||||
|
||||
|
||||
def plot_tree(g):
|
||||
# this plot requires pygraphviz package
|
||||
pos = nx.nx_agraph.graphviz_layout(g, prog="dot")
|
||||
nx.draw(
|
||||
g,
|
||||
pos,
|
||||
with_labels=False,
|
||||
node_size=10,
|
||||
node_color=[[0.5, 0.5, 0.5]],
|
||||
arrowsize=4,
|
||||
)
|
||||
plt.show()
|
||||
|
||||
|
||||
plot_tree(graph.to_networkx())
|
||||
|
||||
#################################################################################
|
||||
# You can read more about the definition of :func:`~dgl.batch`, or
|
||||
# skip ahead to the next step:
|
||||
# .. note::
|
||||
#
|
||||
# **Definition**: :func:`~dgl.batch` unions a list of :math:`B`
|
||||
# :class:`~dgl.DGLGraph`\ s and returns a :class:`~dgl.DGLGraph` of batch
|
||||
# size :math:`B`.
|
||||
#
|
||||
# - The union includes all the nodes,
|
||||
# edges, and their features. The order of nodes, edges, and features are
|
||||
# preserved.
|
||||
#
|
||||
# - Given that you have :math:`V_i` nodes for graph
|
||||
# :math:`\mathcal{G}_i`, the node ID :math:`j` in graph
|
||||
# :math:`\mathcal{G}_i` correspond to node ID
|
||||
# :math:`j + \sum_{k=1}^{i-1} V_k` in the batched graph.
|
||||
#
|
||||
# - Therefore, performing feature transformation and message passing on
|
||||
# the batched graph is equivalent to doing those
|
||||
# on all ``DGLGraph`` constituents in parallel.
|
||||
#
|
||||
# - Duplicate references to the same graph are
|
||||
# treated as deep copies; the nodes, edges, and features are duplicated,
|
||||
# and mutation on one reference does not affect the other.
|
||||
# - The batched graph keeps track of the meta
|
||||
# information of the constituents so it can be
|
||||
# :func:`~dgl.batched_graph.unbatch`\ ed to list of ``DGLGraph``\ s.
|
||||
#
|
||||
# Step 2: Tree-LSTM cell with message-passing APIs
|
||||
# ------------------------------------------------
|
||||
#
|
||||
# Researchers have proposed two types of Tree-LSTMs: Child-Sum
|
||||
# Tree-LSTMs, and :math:`N`-ary Tree-LSTMs. In this tutorial you focus
|
||||
# on applying *Binary* Tree-LSTM to binarized constituency trees. This
|
||||
# application is also known as *Constituency Tree-LSTM*. Use PyTorch
|
||||
# as a backend framework to set up the network.
|
||||
#
|
||||
# In `N`-ary Tree-LSTM, each unit at node :math:`j` maintains a hidden
|
||||
# representation :math:`h_j` and a memory cell :math:`c_j`. The unit
|
||||
# :math:`j` takes the input vector :math:`x_j` and the hidden
|
||||
# representations of the child units: :math:`h_{jl}, 1\leq l\leq N` as
|
||||
# input, then update its new hidden representation :math:`h_j` and memory
|
||||
# cell :math:`c_j` by:
|
||||
#
|
||||
# .. math::
|
||||
#
|
||||
# i_j & = & \sigma\left(W^{(i)}x_j + \sum_{l=1}^{N}U^{(i)}_l h_{jl} + b^{(i)}\right), & (1)\\
|
||||
# f_{jk} & = & \sigma\left(W^{(f)}x_j + \sum_{l=1}^{N}U_{kl}^{(f)} h_{jl} + b^{(f)} \right), & (2)\\
|
||||
# o_j & = & \sigma\left(W^{(o)}x_j + \sum_{l=1}^{N}U_{l}^{(o)} h_{jl} + b^{(o)} \right), & (3) \\
|
||||
# u_j & = & \textrm{tanh}\left(W^{(u)}x_j + \sum_{l=1}^{N} U_l^{(u)}h_{jl} + b^{(u)} \right), & (4)\\
|
||||
# c_j & = & i_j \odot u_j + \sum_{l=1}^{N} f_{jl} \odot c_{jl}, &(5) \\
|
||||
# h_j & = & o_j \cdot \textrm{tanh}(c_j), &(6) \\
|
||||
#
|
||||
# It can be decomposed into three phases: ``message_func``,
|
||||
# ``reduce_func`` and ``apply_node_func``.
|
||||
#
|
||||
# .. note::
|
||||
# ``apply_node_func`` is a new node UDF that has not been introduced before. In
|
||||
# ``apply_node_func``, a user specifies what to do with node features,
|
||||
# without considering edge features and messages. In a Tree-LSTM case,
|
||||
# ``apply_node_func`` is a must, since there exists (leaf) nodes with
|
||||
# :math:`0` incoming edges, which would not be updated with
|
||||
# ``reduce_func``.
|
||||
#
|
||||
|
||||
import torch as th
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class TreeLSTMCell(nn.Module):
|
||||
def __init__(self, x_size, h_size):
|
||||
super(TreeLSTMCell, self).__init__()
|
||||
self.W_iou = nn.Linear(x_size, 3 * h_size, bias=False)
|
||||
self.U_iou = nn.Linear(2 * h_size, 3 * h_size, bias=False)
|
||||
self.b_iou = nn.Parameter(th.zeros(1, 3 * h_size))
|
||||
self.U_f = nn.Linear(2 * h_size, 2 * h_size)
|
||||
|
||||
def message_func(self, edges):
|
||||
return {"h": edges.src["h"], "c": edges.src["c"]}
|
||||
|
||||
def reduce_func(self, nodes):
|
||||
# concatenate h_jl for equation (1), (2), (3), (4)
|
||||
h_cat = nodes.mailbox["h"].view(nodes.mailbox["h"].size(0), -1)
|
||||
# equation (2)
|
||||
f = th.sigmoid(self.U_f(h_cat)).view(*nodes.mailbox["h"].size())
|
||||
# second term of equation (5)
|
||||
c = th.sum(f * nodes.mailbox["c"], 1)
|
||||
return {"iou": self.U_iou(h_cat), "c": c}
|
||||
|
||||
def apply_node_func(self, nodes):
|
||||
# equation (1), (3), (4)
|
||||
iou = nodes.data["iou"] + self.b_iou
|
||||
i, o, u = th.chunk(iou, 3, 1)
|
||||
i, o, u = th.sigmoid(i), th.sigmoid(o), th.tanh(u)
|
||||
# equation (5)
|
||||
c = i * u + nodes.data["c"]
|
||||
# equation (6)
|
||||
h = o * th.tanh(c)
|
||||
return {"h": h, "c": c}
|
||||
|
||||
|
||||
##############################################################################
|
||||
# Step 3: Define traversal
|
||||
# ------------------------
|
||||
#
|
||||
# After you define the message-passing functions, induce the
|
||||
# right order to trigger them. This is a significant departure from models
|
||||
# such as GCN, where all nodes are pulling messages from upstream ones
|
||||
# *simultaneously*.
|
||||
#
|
||||
# In the case of Tree-LSTM, messages start from leaves of the tree, and
|
||||
# propagate/processed upwards until they reach the roots. A visualization
|
||||
# is as follows:
|
||||
#
|
||||
# .. figure:: https://i.loli.net/2018/11/09/5be4b5d2df54d.gif
|
||||
# :alt:
|
||||
#
|
||||
# DGL defines a generator to perform the topological sort, each item is a
|
||||
# tensor recording the nodes from bottom level to the roots. One can
|
||||
# appreciate the degree of parallelism by inspecting the difference of the
|
||||
# followings:
|
||||
#
|
||||
|
||||
# to heterogenous graph
|
||||
trv_a_tree = dgl.graph(a_tree.edges())
|
||||
print("Traversing one tree:")
|
||||
print(dgl.topological_nodes_generator(trv_a_tree))
|
||||
|
||||
# to heterogenous graph
|
||||
trv_graph = dgl.graph(graph.edges())
|
||||
print("Traversing many trees at the same time:")
|
||||
print(dgl.topological_nodes_generator(trv_graph))
|
||||
|
||||
##############################################################################
|
||||
# Call :meth:`~dgl.DGLGraph.prop_nodes` to trigger the message passing:
|
||||
|
||||
import dgl.function as fn
|
||||
import torch as th
|
||||
|
||||
trv_graph.ndata["a"] = th.ones(graph.num_nodes(), 1)
|
||||
traversal_order = dgl.topological_nodes_generator(trv_graph)
|
||||
trv_graph.prop_nodes(
|
||||
traversal_order,
|
||||
message_func=fn.copy_u("a", "a"),
|
||||
reduce_func=fn.sum("a", "a"),
|
||||
)
|
||||
|
||||
# the following is a syntax sugar that does the same
|
||||
# dgl.prop_nodes_topo(graph)
|
||||
|
||||
##############################################################################
|
||||
# .. note::
|
||||
#
|
||||
# Before you call :meth:`~dgl.DGLGraph.prop_nodes`, specify a
|
||||
# `message_func` and `reduce_func` in advance. In the example, you can see built-in
|
||||
# copy-from-source and sum functions as message functions, and a reduce
|
||||
# function for demonstration.
|
||||
#
|
||||
# Putting it together
|
||||
# -------------------
|
||||
#
|
||||
# Here is the complete code that specifies the ``Tree-LSTM`` class.
|
||||
#
|
||||
|
||||
|
||||
class TreeLSTM(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
num_vocabs,
|
||||
x_size,
|
||||
h_size,
|
||||
num_classes,
|
||||
dropout,
|
||||
pretrained_emb=None,
|
||||
):
|
||||
super(TreeLSTM, self).__init__()
|
||||
self.x_size = x_size
|
||||
self.embedding = nn.Embedding(num_vocabs, x_size)
|
||||
if pretrained_emb is not None:
|
||||
print("Using glove")
|
||||
self.embedding.weight.data.copy_(pretrained_emb)
|
||||
self.embedding.weight.requires_grad = True
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
self.linear = nn.Linear(h_size, num_classes)
|
||||
self.cell = TreeLSTMCell(x_size, h_size)
|
||||
|
||||
def forward(self, batch, h, c):
|
||||
"""Compute tree-lstm prediction given a batch.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
batch : dgl.data.SSTBatch
|
||||
The data batch.
|
||||
h : Tensor
|
||||
Initial hidden state.
|
||||
c : Tensor
|
||||
Initial cell state.
|
||||
|
||||
Returns
|
||||
-------
|
||||
logits : Tensor
|
||||
The prediction of each node.
|
||||
"""
|
||||
g = batch.graph
|
||||
# to heterogenous graph
|
||||
g = dgl.graph(g.edges())
|
||||
# feed embedding
|
||||
embeds = self.embedding(batch.wordid * batch.mask)
|
||||
g.ndata["iou"] = self.cell.W_iou(
|
||||
self.dropout(embeds)
|
||||
) * batch.mask.float().unsqueeze(-1)
|
||||
g.ndata["h"] = h
|
||||
g.ndata["c"] = c
|
||||
# propagate
|
||||
dgl.prop_nodes_topo(
|
||||
g,
|
||||
message_func=self.cell.message_func,
|
||||
reduce_func=self.cell.reduce_func,
|
||||
apply_node_func=self.cell.apply_node_func,
|
||||
)
|
||||
# compute logits
|
||||
h = self.dropout(g.ndata.pop("h"))
|
||||
logits = self.linear(h)
|
||||
return logits
|
||||
|
||||
|
||||
import torch.nn.functional as F
|
||||
|
||||
##############################################################################
|
||||
# Main Loop
|
||||
# ---------
|
||||
#
|
||||
# Finally, you could write a training paradigm in PyTorch.
|
||||
#
|
||||
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
device = th.device("cpu")
|
||||
# hyper parameters
|
||||
x_size = 256
|
||||
h_size = 256
|
||||
dropout = 0.5
|
||||
lr = 0.05
|
||||
weight_decay = 1e-4
|
||||
epochs = 10
|
||||
|
||||
# create the model
|
||||
model = TreeLSTM(
|
||||
trainset.vocab_size, x_size, h_size, trainset.num_classes, dropout
|
||||
)
|
||||
print(model)
|
||||
|
||||
# create the optimizer
|
||||
optimizer = th.optim.Adagrad(
|
||||
model.parameters(), lr=lr, weight_decay=weight_decay
|
||||
)
|
||||
|
||||
|
||||
def batcher(dev):
|
||||
def batcher_dev(batch):
|
||||
batch_trees = dgl.batch(batch)
|
||||
return SSTBatch(
|
||||
graph=batch_trees,
|
||||
mask=batch_trees.ndata["mask"].to(device),
|
||||
wordid=batch_trees.ndata["x"].to(device),
|
||||
label=batch_trees.ndata["y"].to(device),
|
||||
)
|
||||
|
||||
return batcher_dev
|
||||
|
||||
|
||||
train_loader = DataLoader(
|
||||
dataset=tiny_sst,
|
||||
batch_size=5,
|
||||
collate_fn=batcher(device),
|
||||
shuffle=False,
|
||||
num_workers=0,
|
||||
)
|
||||
|
||||
# training loop
|
||||
for epoch in range(epochs):
|
||||
for step, batch in enumerate(train_loader):
|
||||
g = batch.graph
|
||||
n = g.num_nodes()
|
||||
h = th.zeros((n, h_size))
|
||||
c = th.zeros((n, h_size))
|
||||
logits = model(batch, h, c)
|
||||
logp = F.log_softmax(logits, 1)
|
||||
loss = F.nll_loss(logp, batch.label, reduction="sum")
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
pred = th.argmax(logits, 1)
|
||||
acc = float(th.sum(th.eq(batch.label, pred))) / len(batch.label)
|
||||
print(
|
||||
"Epoch {:05d} | Step {:05d} | Loss {:.4f} | Acc {:.4f} |".format(
|
||||
epoch, step, loss.item(), acc
|
||||
)
|
||||
)
|
||||
##############################################################################
|
||||
# To train the model on a full dataset with different settings (such as CPU or GPU),
|
||||
# refer to the `PyTorch example <https://github.com/dmlc/dgl/tree/master/examples/pytorch/tree_lstm>`__.
|
||||
# There is also an implementation of the Child-Sum Tree-LSTM.
|
||||
@@ -0,0 +1,16 @@
|
||||
.. _tutorials2-index:
|
||||
|
||||
Batching many small graphs
|
||||
-------------------------------
|
||||
|
||||
* **Tree-LSTM** `[paper] <https://arxiv.org/abs/1503.00075>`__ `[tutorial]
|
||||
<2_small_graph/3_tree-lstm.html>`__ `[PyTorch code]
|
||||
<https://github.com/dmlc/dgl/blob/master/examples/pytorch/tree_lstm>`__:
|
||||
Sentences have inherent structures that are thrown
|
||||
away by treating them simply as sequences. Tree-LSTM is a powerful model
|
||||
that learns the representation by using prior syntactic structures such as a parse-tree.
|
||||
The challenge in training is that simply by padding
|
||||
a sentence to the maximum length no longer works. Trees of different
|
||||
sentences have different sizes and topologies. DGL solves this problem by
|
||||
adding the trees to a bigger container graph, and then using message-passing
|
||||
to explore maximum parallelism. Batching is a key API for this.
|
||||
@@ -0,0 +1,793 @@
|
||||
"""
|
||||
.. _model-dgmg:
|
||||
|
||||
Generative Models of Graphs
|
||||
===========================================
|
||||
|
||||
**Author**: `Mufei Li <https://github.com/mufeili>`_,
|
||||
`Lingfan Yu <https://github.com/ylfdq1118>`_, Zheng Zhang
|
||||
|
||||
.. warning::
|
||||
|
||||
The tutorial aims at gaining insights into the paper, with code as a mean
|
||||
of explanation. The implementation thus is NOT optimized for running
|
||||
efficiency. For recommended implementation, please refer to the `official
|
||||
examples <https://github.com/dmlc/dgl/tree/master/examples>`_.
|
||||
|
||||
"""
|
||||
|
||||
##############################################################################
|
||||
#
|
||||
# In this tutorial, you learn how to train and generate one graph at
|
||||
# a time. You also explore parallelism within the graph embedding operation, which is an
|
||||
# essential building block. The tutorial ends with a simple optimization that
|
||||
# delivers double the speed by batching across graphs.
|
||||
#
|
||||
# Earlier tutorials showed how embedding a graph or
|
||||
# a node enables you to work on tasks such as `semi-supervised classification for nodes
|
||||
# <http://docs.dgl.ai/tutorials/models/1_gcn.html#sphx-glr-tutorials-models-1-gcn-py>`__
|
||||
# or `sentiment analysis
|
||||
# <http://docs.dgl.ai/tutorials/models/3_tree-lstm.html#sphx-glr-tutorials-models-3-tree-lstm-py>`__.
|
||||
# Wouldn't it be interesting to predict the future evolution of the graph and
|
||||
# perform the analysis iteratively?
|
||||
#
|
||||
# To address the evolution of the graphs, you generate a variety of graph samples. In other words, you need
|
||||
# **generative models** of graphs. In-addition to learning
|
||||
# node and edge features, you would need to model the distribution of arbitrary graphs.
|
||||
# While general generative models can model the density function explicitly and
|
||||
# implicitly and generate samples at once or sequentially, you only focus
|
||||
# on explicit generative models for sequential generation here. Typical applications
|
||||
# include drug or materials discovery, chemical processes, or proteomics.
|
||||
#
|
||||
# Introduction
|
||||
# --------------------
|
||||
# The primitive actions of mutating a graph in Deep Graph Library (DGL) are nothing more than ``add_nodes``
|
||||
# and ``add_edges``. That is, if you were to draw a circle of three nodes,
|
||||
#
|
||||
# .. figure:: https://user-images.githubusercontent.com/19576924/48313438-78baf000-e5f7-11e8-931e-cd00ab34fa50.gif
|
||||
# :alt:
|
||||
#
|
||||
# you can write the code as follows.
|
||||
#
|
||||
|
||||
import os
|
||||
|
||||
os.environ["DGLBACKEND"] = "pytorch"
|
||||
import dgl
|
||||
|
||||
g = dgl.DGLGraph()
|
||||
g.add_nodes(1) # Add node 0
|
||||
g.add_nodes(1) # Add node 1
|
||||
|
||||
# Edges in DGLGraph are directed by default.
|
||||
# For undirected edges, add edges for both directions.
|
||||
g.add_edges([1, 0], [0, 1]) # Add edges (1, 0), (0, 1)
|
||||
g.add_nodes(1) # Add node 2
|
||||
g.add_edges([2, 1], [1, 2]) # Add edges (2, 1), (1, 2)
|
||||
g.add_edges([2, 0], [0, 2]) # Add edges (2, 0), (0, 2)
|
||||
|
||||
#######################################################################################
|
||||
# Real-world graphs are much more complex. There are many families of graphs,
|
||||
# with different sizes, topologies, node types, edge types, and the possibility
|
||||
# of multigraphs. Besides, a same graph can be generated in many different
|
||||
# orders. Regardless, the generative process entails a few steps.
|
||||
#
|
||||
# - Encode a changing graph.
|
||||
# - Perform actions stochastically.
|
||||
# - If you are training, collect error signals and optimize the model parameters.
|
||||
#
|
||||
# When it comes to implementation, another important aspect is speed. How do you
|
||||
# parallelize the computation, given that generating a graph is fundamentally a
|
||||
# sequential process?
|
||||
#
|
||||
# .. note::
|
||||
#
|
||||
# To be sure, this is not necessarily a hard constraint. Subgraphs can be
|
||||
# built in parallel and then get assembled. But we
|
||||
# will restrict ourselves to the sequential processes for this tutorial.
|
||||
#
|
||||
#
|
||||
# DGMG: The main flow
|
||||
# --------------------
|
||||
# For this tutorial, you use
|
||||
# `Deep Generative Models of Graphs <https://arxiv.org/abs/1803.03324>`__
|
||||
# ) (DGMG) to implement a graph generative model using DGL. Its algorithmic
|
||||
# framework is general but also challenging to parallelize.
|
||||
#
|
||||
# .. note::
|
||||
#
|
||||
# While it's possible for DGMG to handle complex graphs with typed nodes,
|
||||
# typed edges, and multigraphs, here you use a simplified version of it
|
||||
# for generating graph topologies.
|
||||
#
|
||||
# DGMG generates a graph by following a state machine, which is basically a
|
||||
# two-level loop. Generate one node at a time and connect it to a subset of
|
||||
# the existing nodes, one at a time. This is similar to language modeling. The
|
||||
# generative process is an iterative one that emits one word or character or sentence
|
||||
# at a time, conditioned on the sequence generated so far.
|
||||
#
|
||||
# At each time step, you either:
|
||||
# - Add a new node to the graph
|
||||
# - Select two existing nodes and add an edge between them
|
||||
#
|
||||
# .. figure:: https://user-images.githubusercontent.com/19576924/48605003-7f11e900-e9b6-11e8-8880-87362348e154.png
|
||||
# :alt:
|
||||
#
|
||||
# The Python code will look as follows. In fact, this is *exactly* how inference
|
||||
# with DGMG is implemented in DGL.
|
||||
#
|
||||
|
||||
|
||||
def forward_inference(self):
|
||||
stop = self.add_node_and_update()
|
||||
while (not stop) and (self.g.num_nodes() < self.v_max + 1):
|
||||
num_trials = 0
|
||||
to_add_edge = self.add_edge_or_not()
|
||||
while to_add_edge and (num_trials < self.g.num_nodes() - 1):
|
||||
self.choose_dest_and_update()
|
||||
num_trials += 1
|
||||
to_add_edge = self.add_edge_or_not()
|
||||
stop = self.add_node_and_update()
|
||||
return self.g
|
||||
|
||||
|
||||
#######################################################################################
|
||||
# Assume you have a pre-trained model for generating cycles of nodes 10-20.
|
||||
# How does it generate a cycle on-the-fly during inference? Use the code below
|
||||
# to create an animation with your own model.
|
||||
#
|
||||
# ::
|
||||
#
|
||||
# import torch
|
||||
# import matplotlib.animation as animation
|
||||
# import matplotlib.pyplot as plt
|
||||
# import networkx as nx
|
||||
# from copy import deepcopy
|
||||
#
|
||||
# if __name__ == '__main__':
|
||||
# # pre-trained model saved with path ./model.pth
|
||||
# model = torch.load('./model.pth')
|
||||
# model.eval()
|
||||
# g = model()
|
||||
#
|
||||
# src_list = g.edges()[1]
|
||||
# dest_list = g.edges()[0]
|
||||
#
|
||||
# evolution = []
|
||||
#
|
||||
# nx_g = nx.Graph()
|
||||
# evolution.append(deepcopy(nx_g))
|
||||
#
|
||||
# for i in range(0, len(src_list), 2):
|
||||
# src = src_list[i].item()
|
||||
# dest = dest_list[i].item()
|
||||
# if src not in nx_g.nodes():
|
||||
# nx_g.add_node(src)
|
||||
# evolution.append(deepcopy(nx_g))
|
||||
# if dest not in nx_g.nodes():
|
||||
# nx_g.add_node(dest)
|
||||
# evolution.append(deepcopy(nx_g))
|
||||
# nx_g.add_edges_from([(src, dest), (dest, src)])
|
||||
# evolution.append(deepcopy(nx_g))
|
||||
#
|
||||
# def animate(i):
|
||||
# ax.cla()
|
||||
# g_t = evolution[i]
|
||||
# nx.draw_circular(g_t, with_labels=True, ax=ax,
|
||||
# node_color=['#FEBD69'] * g_t.num_nodes())
|
||||
#
|
||||
# fig, ax = plt.subplots()
|
||||
# ani = animation.FuncAnimation(fig, animate,
|
||||
# frames=len(evolution),
|
||||
# interval=600)
|
||||
#
|
||||
# .. figure:: https://user-images.githubusercontent.com/19576924/48928548-2644d200-ef1b-11e8-8591-da93345382ad.gif
|
||||
# :alt:
|
||||
#
|
||||
# DGMG: Optimization objective
|
||||
# ------------------------------
|
||||
# Similar to language modeling, DGMG trains the model with *behavior cloning*,
|
||||
# or *teacher forcing*. Assume for each graph there exists a sequence of
|
||||
# *oracle actions* :math:`a_{1},\cdots,a_{T}` that generates it. What the model
|
||||
# does is to follow these actions, compute the joint probabilities of such
|
||||
# action sequences, and maximize them.
|
||||
#
|
||||
# By chain rule, the probability of taking :math:`a_{1},\cdots,a_{T}` is:
|
||||
#
|
||||
# .. math::
|
||||
#
|
||||
# p(a_{1},\cdots, a_{T}) = p(a_{1})p(a_{2}|a_{1})\cdots p(a_{T}|a_{1},\cdots,a_{T-1}).\\
|
||||
#
|
||||
# The optimization objective is then simply the typical MLE loss:
|
||||
#
|
||||
# .. math::
|
||||
#
|
||||
# -\log p(a_{1},\cdots,a_{T})=-\sum_{t=1}^{T}\log p(a_{t}|a_{1},\cdots, a_{t-1}).\\
|
||||
#
|
||||
|
||||
|
||||
def forward_train(self, actions):
|
||||
"""
|
||||
- actions: list
|
||||
- Contains a_1, ..., a_T described above
|
||||
- self.prepare_for_train()
|
||||
- Initializes self.action_step to be 0, which will get
|
||||
incremented by 1 every time it is called.
|
||||
- Initializes objects recording log p(a_t|a_1,...a_{t-1})
|
||||
|
||||
Returns
|
||||
-------
|
||||
- self.get_log_prob(): log p(a_1, ..., a_T)
|
||||
"""
|
||||
self.prepare_for_train()
|
||||
|
||||
stop = self.add_node_and_update(a=actions[self.action_step])
|
||||
while not stop:
|
||||
to_add_edge = self.add_edge_or_not(a=actions[self.action_step])
|
||||
while to_add_edge:
|
||||
self.choose_dest_and_update(a=actions[self.action_step])
|
||||
to_add_edge = self.add_edge_or_not(a=actions[self.action_step])
|
||||
stop = self.add_node_and_update(a=actions[self.action_step])
|
||||
return self.get_log_prob()
|
||||
|
||||
|
||||
#######################################################################################
|
||||
# The key difference between ``forward_train`` and ``forward_inference`` is
|
||||
# that the training process takes oracle actions as input and returns log
|
||||
# probabilities for evaluating the loss.
|
||||
#
|
||||
# DGMG: The implementation
|
||||
# --------------------------
|
||||
# The ``DGMG`` class
|
||||
# ``````````````````````````
|
||||
# Below you can find the skeleton code for the model. You gradually
|
||||
# fill in the details for each function.
|
||||
#
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class DGMGSkeleton(nn.Module):
|
||||
def __init__(self, v_max):
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
v_max: int
|
||||
Max number of nodes considered
|
||||
"""
|
||||
super(DGMGSkeleton, self).__init__()
|
||||
|
||||
# Graph configuration
|
||||
self.v_max = v_max
|
||||
|
||||
def add_node_and_update(self, a=None):
|
||||
"""Decide if to add a new node.
|
||||
If a new node should be added, update the graph."""
|
||||
return NotImplementedError
|
||||
|
||||
def add_edge_or_not(self, a=None):
|
||||
"""Decide if a new edge should be added."""
|
||||
return NotImplementedError
|
||||
|
||||
def choose_dest_and_update(self, a=None):
|
||||
"""Choose destination and connect it to the latest node.
|
||||
Add edges for both directions and update the graph."""
|
||||
return NotImplementedError
|
||||
|
||||
def forward_train(self, actions):
|
||||
"""Forward at training time. It records the probability
|
||||
of generating a ground truth graph following the actions."""
|
||||
return NotImplementedError
|
||||
|
||||
def forward_inference(self):
|
||||
"""Forward at inference time.
|
||||
It generates graphs on the fly."""
|
||||
return NotImplementedError
|
||||
|
||||
def forward(self, actions=None):
|
||||
# The graph you will work on
|
||||
self.g = dgl.DGLGraph()
|
||||
|
||||
# If there are some features for nodes and edges,
|
||||
# zero tensors will be set for those of new nodes and edges.
|
||||
self.g.set_n_initializer(dgl.frame.zero_initializer)
|
||||
self.g.set_e_initializer(dgl.frame.zero_initializer)
|
||||
|
||||
if self.training:
|
||||
return self.forward_train(actions=actions)
|
||||
else:
|
||||
return self.forward_inference()
|
||||
|
||||
|
||||
#######################################################################################
|
||||
# Encoding a dynamic graph
|
||||
# ``````````````````````````
|
||||
# All the actions generating a graph are sampled from probability
|
||||
# distributions. In order to do that, you project the structured data,
|
||||
# namely the graph, onto an Euclidean space. The challenge is that such
|
||||
# process, called *embedding*, needs to be repeated as the graphs mutate.
|
||||
#
|
||||
# Graph embedding
|
||||
# ''''''''''''''''''''''''''
|
||||
# Let :math:`G=(V,E)` be an arbitrary graph. Each node :math:`v` has an
|
||||
# embedding vector :math:`\textbf{h}_{v} \in \mathbb{R}^{n}`. Similarly,
|
||||
# the graph has an embedding vector :math:`\textbf{h}_{G} \in \mathbb{R}^{k}`.
|
||||
# Typically, :math:`k > n` since a graph contains more information than
|
||||
# an individual node.
|
||||
#
|
||||
# The graph embedding is a weighted sum of node embeddings under a linear
|
||||
# transformation:
|
||||
#
|
||||
# .. math::
|
||||
#
|
||||
# \textbf{h}_{G} =\sum_{v\in V}\text{Sigmoid}(g_m(\textbf{h}_{v}))f_{m}(\textbf{h}_{v}),\\
|
||||
#
|
||||
# The first term, :math:`\text{Sigmoid}(g_m(\textbf{h}_{v}))`, computes a
|
||||
# gating function and can be thought of as how much the overall graph embedding
|
||||
# attends on each node. The second term :math:`f_{m}:\mathbb{R}^{n}\rightarrow\mathbb{R}^{k}`
|
||||
# maps the node embeddings to the space of graph embeddings.
|
||||
#
|
||||
# Implement graph embedding as a ``GraphEmbed`` class.
|
||||
#
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class GraphEmbed(nn.Module):
|
||||
def __init__(self, node_hidden_size):
|
||||
super(GraphEmbed, self).__init__()
|
||||
|
||||
# Setting from the paper
|
||||
self.graph_hidden_size = 2 * node_hidden_size
|
||||
|
||||
# Embed graphs
|
||||
self.node_gating = nn.Sequential(
|
||||
nn.Linear(node_hidden_size, 1), nn.Sigmoid()
|
||||
)
|
||||
self.node_to_graph = nn.Linear(node_hidden_size, self.graph_hidden_size)
|
||||
|
||||
def forward(self, g):
|
||||
if g.num_nodes() == 0:
|
||||
return torch.zeros(1, self.graph_hidden_size)
|
||||
else:
|
||||
# Node features are stored as hv in ndata.
|
||||
hvs = g.ndata["hv"]
|
||||
return (self.node_gating(hvs) * self.node_to_graph(hvs)).sum(
|
||||
0, keepdim=True
|
||||
)
|
||||
|
||||
|
||||
#######################################################################################
|
||||
# Update node embeddings via graph propagation
|
||||
# '''''''''''''''''''''''''''''''''''''''''''''
|
||||
#
|
||||
# The mechanism of updating node embeddings in DGMG is similar to that for
|
||||
# graph convolutional networks. For a node :math:`v` in the graph, its
|
||||
# neighbor :math:`u` sends a message to it with
|
||||
#
|
||||
# .. math::
|
||||
#
|
||||
# \textbf{m}_{u\rightarrow v}=\textbf{W}_{m}\text{concat}([\textbf{h}_{v}, \textbf{h}_{u}, \textbf{x}_{u, v}]) + \textbf{b}_{m},\\
|
||||
#
|
||||
# where :math:`\textbf{x}_{u,v}` is the embedding of the edge between
|
||||
# :math:`u` and :math:`v`.
|
||||
#
|
||||
# After receiving messages from all its neighbors, :math:`v` summarizes them
|
||||
# with a node activation vector
|
||||
#
|
||||
# .. math::
|
||||
#
|
||||
# \textbf{a}_{v} = \sum_{u: (u, v)\in E}\textbf{m}_{u\rightarrow v}\\
|
||||
#
|
||||
# and use this information to update its own feature:
|
||||
#
|
||||
# .. math::
|
||||
#
|
||||
# \textbf{h}'_{v} = \textbf{GRU}(\textbf{h}_{v}, \textbf{a}_{v}).\\
|
||||
#
|
||||
# Performing all the operations above once for all nodes synchronously is
|
||||
# called one round of graph propagation. The more rounds of graph propagation
|
||||
# you perform, the longer distance messages travel throughout the graph.
|
||||
#
|
||||
# With DGL, you implement graph propagation with ``g.update_all``.
|
||||
# The message notation here can be a bit confusing. Researchers can refer
|
||||
# to :math:`\textbf{m}_{u\rightarrow v}` as messages, however the message function
|
||||
# below only passes :math:`\text{concat}([\textbf{h}_{u}, \textbf{x}_{u, v}])`.
|
||||
# The operation :math:`\textbf{W}_{m}\text{concat}([\textbf{h}_{v}, \textbf{h}_{u}, \textbf{x}_{u, v}]) + \textbf{b}_{m}`
|
||||
# is then performed across all edges at once for efficiency consideration.
|
||||
#
|
||||
|
||||
from functools import partial
|
||||
|
||||
|
||||
class GraphProp(nn.Module):
|
||||
def __init__(self, num_prop_rounds, node_hidden_size):
|
||||
super(GraphProp, self).__init__()
|
||||
|
||||
self.num_prop_rounds = num_prop_rounds
|
||||
|
||||
# Setting from the paper
|
||||
self.node_activation_hidden_size = 2 * node_hidden_size
|
||||
|
||||
message_funcs = []
|
||||
node_update_funcs = []
|
||||
self.reduce_funcs = []
|
||||
|
||||
for t in range(num_prop_rounds):
|
||||
# input being [hv, hu, xuv]
|
||||
message_funcs.append(
|
||||
nn.Linear(
|
||||
2 * node_hidden_size + 1, self.node_activation_hidden_size
|
||||
)
|
||||
)
|
||||
|
||||
self.reduce_funcs.append(partial(self.dgmg_reduce, round=t))
|
||||
node_update_funcs.append(
|
||||
nn.GRUCell(self.node_activation_hidden_size, node_hidden_size)
|
||||
)
|
||||
self.message_funcs = nn.ModuleList(message_funcs)
|
||||
self.node_update_funcs = nn.ModuleList(node_update_funcs)
|
||||
|
||||
def dgmg_msg(self, edges):
|
||||
"""For an edge u->v, return concat([h_u, x_uv])"""
|
||||
return {"m": torch.cat([edges.src["hv"], edges.data["he"]], dim=1)}
|
||||
|
||||
def dgmg_reduce(self, nodes, round):
|
||||
hv_old = nodes.data["hv"]
|
||||
m = nodes.mailbox["m"]
|
||||
message = torch.cat(
|
||||
[hv_old.unsqueeze(1).expand(-1, m.size(1), -1), m], dim=2
|
||||
)
|
||||
node_activation = (self.message_funcs[round](message)).sum(1)
|
||||
|
||||
return {"a": node_activation}
|
||||
|
||||
def forward(self, g):
|
||||
if g.num_edges() > 0:
|
||||
for t in range(self.num_prop_rounds):
|
||||
g.update_all(
|
||||
message_func=self.dgmg_msg, reduce_func=self.reduce_funcs[t]
|
||||
)
|
||||
g.ndata["hv"] = self.node_update_funcs[t](
|
||||
g.ndata["a"], g.ndata["hv"]
|
||||
)
|
||||
|
||||
|
||||
#######################################################################################
|
||||
# Actions
|
||||
# ``````````````````````````
|
||||
# All actions are sampled from distributions parameterized using neural networks
|
||||
# and here they are in turn.
|
||||
#
|
||||
# Action 1: Add nodes
|
||||
# ''''''''''''''''''''''''''
|
||||
#
|
||||
# Given the graph embedding vector :math:`\textbf{h}_{G}`, evaluate
|
||||
#
|
||||
# .. math::
|
||||
#
|
||||
# \text{Sigmoid}(\textbf{W}_{\text{add node}}\textbf{h}_{G}+b_{\text{add node}}),\\
|
||||
#
|
||||
# which is then used to parametrize a Bernoulli distribution for deciding whether
|
||||
# to add a new node.
|
||||
#
|
||||
# If a new node is to be added, initialize its feature with
|
||||
#
|
||||
# .. math::
|
||||
#
|
||||
# \textbf{W}_{\text{init}}\text{concat}([\textbf{h}_{\text{init}} , \textbf{h}_{G}])+\textbf{b}_{\text{init}},\\
|
||||
#
|
||||
# where :math:`\textbf{h}_{\text{init}}` is a learnable embedding module for
|
||||
# untyped nodes.
|
||||
#
|
||||
|
||||
import torch.nn.functional as F
|
||||
from torch.distributions import Bernoulli
|
||||
|
||||
|
||||
def bernoulli_action_log_prob(logit, action):
|
||||
"""Calculate the log p of an action with respect to a Bernoulli
|
||||
distribution. Use logit rather than prob for numerical stability."""
|
||||
if action == 0:
|
||||
return F.logsigmoid(-logit)
|
||||
else:
|
||||
return F.logsigmoid(logit)
|
||||
|
||||
|
||||
class AddNode(nn.Module):
|
||||
def __init__(self, graph_embed_func, node_hidden_size):
|
||||
super(AddNode, self).__init__()
|
||||
|
||||
self.graph_op = {"embed": graph_embed_func}
|
||||
|
||||
self.stop = 1
|
||||
self.add_node = nn.Linear(graph_embed_func.graph_hidden_size, 1)
|
||||
|
||||
# If to add a node, initialize its hv
|
||||
self.node_type_embed = nn.Embedding(1, node_hidden_size)
|
||||
self.initialize_hv = nn.Linear(
|
||||
node_hidden_size + graph_embed_func.graph_hidden_size,
|
||||
node_hidden_size,
|
||||
)
|
||||
|
||||
self.init_node_activation = torch.zeros(1, 2 * node_hidden_size)
|
||||
|
||||
def _initialize_node_repr(self, g, node_type, graph_embed):
|
||||
"""Whenver a node is added, initialize its representation."""
|
||||
num_nodes = g.num_nodes()
|
||||
hv_init = self.initialize_hv(
|
||||
torch.cat(
|
||||
[
|
||||
self.node_type_embed(torch.LongTensor([node_type])),
|
||||
graph_embed,
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
)
|
||||
g.nodes[num_nodes - 1].data["hv"] = hv_init
|
||||
g.nodes[num_nodes - 1].data["a"] = self.init_node_activation
|
||||
|
||||
def prepare_training(self):
|
||||
self.log_prob = []
|
||||
|
||||
def forward(self, g, action=None):
|
||||
graph_embed = self.graph_op["embed"](g)
|
||||
|
||||
logit = self.add_node(graph_embed)
|
||||
prob = torch.sigmoid(logit)
|
||||
|
||||
if not self.training:
|
||||
action = Bernoulli(prob).sample().item()
|
||||
stop = bool(action == self.stop)
|
||||
|
||||
if not stop:
|
||||
g.add_nodes(1)
|
||||
self._initialize_node_repr(g, action, graph_embed)
|
||||
if self.training:
|
||||
sample_log_prob = bernoulli_action_log_prob(logit, action)
|
||||
|
||||
self.log_prob.append(sample_log_prob)
|
||||
return stop
|
||||
|
||||
|
||||
#######################################################################################
|
||||
# Action 2: Add edges
|
||||
# ''''''''''''''''''''''''''
|
||||
#
|
||||
# Given the graph embedding vector :math:`\textbf{h}_{G}` and the node
|
||||
# embedding vector :math:`\textbf{h}_{v}` for the latest node :math:`v`,
|
||||
# you evaluate
|
||||
#
|
||||
# .. math::
|
||||
#
|
||||
# \text{Sigmoid}(\textbf{W}_{\text{add edge}}\text{concat}([\textbf{h}_{G}, \textbf{h}_{v}])+b_{\text{add edge}}),\\
|
||||
#
|
||||
# which is then used to parametrize a Bernoulli distribution for deciding
|
||||
# whether to add a new edge starting from :math:`v`.
|
||||
#
|
||||
|
||||
|
||||
class AddEdge(nn.Module):
|
||||
def __init__(self, graph_embed_func, node_hidden_size):
|
||||
super(AddEdge, self).__init__()
|
||||
|
||||
self.graph_op = {"embed": graph_embed_func}
|
||||
self.add_edge = nn.Linear(
|
||||
graph_embed_func.graph_hidden_size + node_hidden_size, 1
|
||||
)
|
||||
|
||||
def prepare_training(self):
|
||||
self.log_prob = []
|
||||
|
||||
def forward(self, g, action=None):
|
||||
graph_embed = self.graph_op["embed"](g)
|
||||
src_embed = g.nodes[g.num_nodes() - 1].data["hv"]
|
||||
|
||||
logit = self.add_edge(torch.cat([graph_embed, src_embed], dim=1))
|
||||
prob = torch.sigmoid(logit)
|
||||
|
||||
if self.training:
|
||||
sample_log_prob = bernoulli_action_log_prob(logit, action)
|
||||
self.log_prob.append(sample_log_prob)
|
||||
else:
|
||||
action = Bernoulli(prob).sample().item()
|
||||
to_add_edge = bool(action == 0)
|
||||
return to_add_edge
|
||||
|
||||
|
||||
#######################################################################################
|
||||
# Action 3: Choose a destination
|
||||
# '''''''''''''''''''''''''''''''''
|
||||
#
|
||||
# When action 2 returns `True`, choose a destination for the
|
||||
# latest node :math:`v`.
|
||||
#
|
||||
# For each possible destination :math:`u\in\{0, \cdots, v-1\}`, the
|
||||
# probability of choosing it is given by
|
||||
#
|
||||
# .. math::
|
||||
#
|
||||
# \frac{\text{exp}(\textbf{W}_{\text{dest}}\text{concat}([\textbf{h}_{u}, \textbf{h}_{v}])+\textbf{b}_{\text{dest}})}{\sum_{i=0}^{v-1}\text{exp}(\textbf{W}_{\text{dest}}\text{concat}([\textbf{h}_{i}, \textbf{h}_{v}])+\textbf{b}_{\text{dest}})}\\
|
||||
#
|
||||
|
||||
from torch.distributions import Categorical
|
||||
|
||||
|
||||
class ChooseDestAndUpdate(nn.Module):
|
||||
def __init__(self, graph_prop_func, node_hidden_size):
|
||||
super(ChooseDestAndUpdate, self).__init__()
|
||||
|
||||
self.graph_op = {"prop": graph_prop_func}
|
||||
self.choose_dest = nn.Linear(2 * node_hidden_size, 1)
|
||||
|
||||
def _initialize_edge_repr(self, g, src_list, dest_list):
|
||||
# For untyped edges, only add 1 to indicate its existence.
|
||||
# For multiple edge types, use a one-hot representation
|
||||
# or an embedding module.
|
||||
edge_repr = torch.ones(len(src_list), 1)
|
||||
g.edges[src_list, dest_list].data["he"] = edge_repr
|
||||
|
||||
def prepare_training(self):
|
||||
self.log_prob = []
|
||||
|
||||
def forward(self, g, dest):
|
||||
src = g.num_nodes() - 1
|
||||
possible_dests = range(src)
|
||||
|
||||
src_embed_expand = g.nodes[src].data["hv"].expand(src, -1)
|
||||
possible_dests_embed = g.nodes[possible_dests].data["hv"]
|
||||
|
||||
dests_scores = self.choose_dest(
|
||||
torch.cat([possible_dests_embed, src_embed_expand], dim=1)
|
||||
).view(1, -1)
|
||||
dests_probs = F.softmax(dests_scores, dim=1)
|
||||
|
||||
if not self.training:
|
||||
dest = Categorical(dests_probs).sample().item()
|
||||
if not g.has_edges_between(src, dest):
|
||||
# For undirected graphs, add edges for both directions
|
||||
# so that you can perform graph propagation.
|
||||
src_list = [src, dest]
|
||||
dest_list = [dest, src]
|
||||
|
||||
g.add_edges(src_list, dest_list)
|
||||
self._initialize_edge_repr(g, src_list, dest_list)
|
||||
|
||||
self.graph_op["prop"](g)
|
||||
if self.training:
|
||||
if dests_probs.nelement() > 1:
|
||||
self.log_prob.append(
|
||||
F.log_softmax(dests_scores, dim=1)[:, dest : dest + 1]
|
||||
)
|
||||
|
||||
|
||||
#######################################################################################
|
||||
# Putting it together
|
||||
# ``````````````````````````
|
||||
#
|
||||
# You are now ready to have a complete implementation of the model class.
|
||||
#
|
||||
|
||||
|
||||
class DGMG(DGMGSkeleton):
|
||||
def __init__(self, v_max, node_hidden_size, num_prop_rounds):
|
||||
super(DGMG, self).__init__(v_max)
|
||||
|
||||
# Graph embedding module
|
||||
self.graph_embed = GraphEmbed(node_hidden_size)
|
||||
|
||||
# Graph propagation module
|
||||
self.graph_prop = GraphProp(num_prop_rounds, node_hidden_size)
|
||||
|
||||
# Actions
|
||||
self.add_node_agent = AddNode(self.graph_embed, node_hidden_size)
|
||||
self.add_edge_agent = AddEdge(self.graph_embed, node_hidden_size)
|
||||
self.choose_dest_agent = ChooseDestAndUpdate(
|
||||
self.graph_prop, node_hidden_size
|
||||
)
|
||||
|
||||
# Forward functions
|
||||
self.forward_train = partial(forward_train, self=self)
|
||||
self.forward_inference = partial(forward_inference, self=self)
|
||||
|
||||
@property
|
||||
def action_step(self):
|
||||
old_step_count = self.step_count
|
||||
self.step_count += 1
|
||||
|
||||
return old_step_count
|
||||
|
||||
def prepare_for_train(self):
|
||||
self.step_count = 0
|
||||
|
||||
self.add_node_agent.prepare_training()
|
||||
self.add_edge_agent.prepare_training()
|
||||
self.choose_dest_agent.prepare_training()
|
||||
|
||||
def add_node_and_update(self, a=None):
|
||||
"""Decide if to add a new node.
|
||||
If a new node should be added, update the graph."""
|
||||
|
||||
return self.add_node_agent(self.g, a)
|
||||
|
||||
def add_edge_or_not(self, a=None):
|
||||
"""Decide if a new edge should be added."""
|
||||
|
||||
return self.add_edge_agent(self.g, a)
|
||||
|
||||
def choose_dest_and_update(self, a=None):
|
||||
"""Choose destination and connect it to the latest node.
|
||||
Add edges for both directions and update the graph."""
|
||||
|
||||
self.choose_dest_agent(self.g, a)
|
||||
|
||||
def get_log_prob(self):
|
||||
add_node_log_p = torch.cat(self.add_node_agent.log_prob).sum()
|
||||
add_edge_log_p = torch.cat(self.add_edge_agent.log_prob).sum()
|
||||
choose_dest_log_p = torch.cat(self.choose_dest_agent.log_prob).sum()
|
||||
return add_node_log_p + add_edge_log_p + choose_dest_log_p
|
||||
|
||||
|
||||
#######################################################################################
|
||||
# Below is an animation where a graph is generated on the fly
|
||||
# after every 10 batches of training for the first 400 batches. You
|
||||
# can see how the model improves over time and begins generating cycles.
|
||||
#
|
||||
# .. figure:: https://user-images.githubusercontent.com/19576924/48929291-60fe3880-ef22-11e8-832a-fbe56656559a.gif
|
||||
# :alt:
|
||||
#
|
||||
# For generative models, you can evaluate performance by checking the percentage
|
||||
# of valid graphs among the graphs it generates on the fly.
|
||||
|
||||
import torch.utils.model_zoo as model_zoo
|
||||
|
||||
# Download a pre-trained model state dict for generating cycles with 10-20 nodes.
|
||||
state_dict = model_zoo.load_url(
|
||||
"https://data.dgl.ai/model/dgmg_cycles-5a0c40be.pth"
|
||||
)
|
||||
model = DGMG(v_max=20, node_hidden_size=16, num_prop_rounds=2)
|
||||
model.load_state_dict(state_dict)
|
||||
model.eval()
|
||||
|
||||
|
||||
def is_valid(g):
|
||||
# Check if g is a cycle having 10-20 nodes.
|
||||
def _get_previous(i, v_max):
|
||||
if i == 0:
|
||||
return v_max
|
||||
else:
|
||||
return i - 1
|
||||
|
||||
def _get_next(i, v_max):
|
||||
if i == v_max:
|
||||
return 0
|
||||
else:
|
||||
return i + 1
|
||||
|
||||
size = g.num_nodes()
|
||||
|
||||
if size < 10 or size > 20:
|
||||
return False
|
||||
for node in range(size):
|
||||
neighbors = g.successors(node)
|
||||
|
||||
if len(neighbors) != 2:
|
||||
return False
|
||||
if _get_previous(node, size - 1) not in neighbors:
|
||||
return False
|
||||
if _get_next(node, size - 1) not in neighbors:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
num_valid = 0
|
||||
for i in range(100):
|
||||
g = model()
|
||||
num_valid += is_valid(g)
|
||||
del model
|
||||
print("Among 100 graphs generated, {}% are valid.".format(num_valid))
|
||||
|
||||
#######################################################################################
|
||||
# For the complete implementation, see the `DGL DGMG example
|
||||
# <https://github.com/dmlc/dgl/tree/master/examples/pytorch/dgmg>`__.
|
||||
#
|
||||
@@ -0,0 +1,14 @@
|
||||
.. _tutorials3-index:
|
||||
|
||||
Generative models
|
||||
--------------------
|
||||
|
||||
* **DGMG** `[paper] <https://arxiv.org/abs/1803.03324>`__ `[tutorial]
|
||||
<3_generative_model/5_dgmg.html>`__ `[PyTorch code]
|
||||
<https://github.com/dmlc/dgl/tree/master/examples/pytorch/dgmg>`__:
|
||||
This model belongs to the family that deals with structural
|
||||
generation. Deep generative models of graphs (DGMG) uses a state-machine approach.
|
||||
It is also very challenging because, unlike Tree-LSTM, every
|
||||
sample has a dynamic, probability-driven structure that is not available
|
||||
before training. You can progressively leverage intra- and
|
||||
inter-graph parallelism to steadily improve the performance.
|
||||
@@ -0,0 +1,291 @@
|
||||
"""
|
||||
.. _model-capsule:
|
||||
|
||||
Capsule Network
|
||||
===========================
|
||||
|
||||
**Author**: Jinjing Zhou, `Jake Zhao <https://cs.nyu.edu/~jakezhao/>`_, Zheng Zhang, Jinyang Li
|
||||
|
||||
In this tutorial, you learn how to describe one of the more classical models in terms of graphs. The approach
|
||||
offers a different perspective. The tutorial describes how to implement a Capsule model for the
|
||||
`capsule network <http://arxiv.org/abs/1710.09829>`__.
|
||||
|
||||
.. warning::
|
||||
|
||||
The tutorial aims at gaining insights into the paper, with code as a mean
|
||||
of explanation. The implementation thus is NOT optimized for running
|
||||
efficiency. For recommended implementation, please refer to the `official
|
||||
examples <https://github.com/dmlc/dgl/tree/master/examples>`_.
|
||||
|
||||
"""
|
||||
#######################################################################################
|
||||
# Key ideas of Capsule
|
||||
# --------------------
|
||||
#
|
||||
# The Capsule model offers two key ideas: Richer representation and dynamic routing.
|
||||
#
|
||||
# **Richer representation** -- In classic convolutional networks, a scalar
|
||||
# value represents the activation of a given feature. By contrast, a
|
||||
# capsule outputs a vector. The vector's length represents the probability
|
||||
# of a feature being present. The vector's orientation represents the
|
||||
# various properties of the feature (such as pose, deformation, texture
|
||||
# etc.).
|
||||
#
|
||||
# |image0|
|
||||
#
|
||||
# **Dynamic routing** -- The output of a capsule is sent to
|
||||
# certain parents in the layer above based on how well the capsule's
|
||||
# prediction agrees with that of a parent. Such dynamic
|
||||
# routing-by-agreement generalizes the static routing of max-pooling.
|
||||
#
|
||||
# During training, routing is accomplished iteratively. Each iteration adjusts
|
||||
# routing weights between capsules based on their observed agreements.
|
||||
# It's a manner similar to a k-means algorithm or `competitive
|
||||
# learning <https://en.wikipedia.org/wiki/Competitive_learning>`__.
|
||||
#
|
||||
# In this tutorial, you see how a capsule's dynamic routing algorithm can be
|
||||
# naturally expressed as a graph algorithm. The implementation is adapted
|
||||
# from `Cedric
|
||||
# Chee <https://github.com/cedrickchee/capsule-net-pytorch>`__, replacing
|
||||
# only the routing layer. This version achieves similar speed and accuracy.
|
||||
#
|
||||
# Model implementation
|
||||
# ----------------------
|
||||
# Step 1: Setup and graph initialization
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
#
|
||||
# The connectivity between two layers of capsules form a directed,
|
||||
# bipartite graph, as shown in the Figure below.
|
||||
#
|
||||
# |image1|
|
||||
#
|
||||
# Each node :math:`j` is associated with feature :math:`v_j`,
|
||||
# representing its capsule’s output. Each edge is associated with
|
||||
# features :math:`b_{ij}` and :math:`\hat{u}_{j|i}`. :math:`b_{ij}`
|
||||
# determines routing weights, and :math:`\hat{u}_{j|i}` represents the
|
||||
# prediction of capsule :math:`i` for :math:`j`.
|
||||
#
|
||||
# Here's how we set up the graph and initialize node and edge features.
|
||||
|
||||
import os
|
||||
|
||||
os.environ["DGLBACKEND"] = "pytorch"
|
||||
import dgl
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import torch as th
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def init_graph(in_nodes, out_nodes, f_size):
|
||||
u = np.repeat(np.arange(in_nodes), out_nodes)
|
||||
v = np.tile(np.arange(in_nodes, in_nodes + out_nodes), in_nodes)
|
||||
g = dgl.DGLGraph((u, v))
|
||||
# init states
|
||||
g.ndata["v"] = th.zeros(in_nodes + out_nodes, f_size)
|
||||
g.edata["b"] = th.zeros(in_nodes * out_nodes, 1)
|
||||
return g
|
||||
|
||||
|
||||
#########################################################################################
|
||||
# Step 2: Define message passing functions
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
#
|
||||
# This is the pseudocode for Capsule's routing algorithm.
|
||||
#
|
||||
# |image2|
|
||||
# Implement pseudocode lines 4-7 in the class `DGLRoutingLayer` as the following steps:
|
||||
#
|
||||
# 1. Calculate coupling coefficients.
|
||||
#
|
||||
# - Coefficients are the softmax over all out-edge of in-capsules.
|
||||
# :math:`\textbf{c}_{i,j} = \text{softmax}(\textbf{b}_{i,j})`.
|
||||
#
|
||||
# 2. Calculate weighted sum over all in-capsules.
|
||||
#
|
||||
# - Output of a capsule is equal to the weighted sum of its in-capsules
|
||||
# :math:`s_j=\sum_i c_{ij}\hat{u}_{j|i}`
|
||||
#
|
||||
# 3. Squash outputs.
|
||||
#
|
||||
# - Squash the length of a Capsule's output vector to range (0,1), so it can represent the probability (of some feature being present).
|
||||
# - :math:`v_j=\text{squash}(s_j)=\frac{||s_j||^2}{1+||s_j||^2}\frac{s_j}{||s_j||}`
|
||||
#
|
||||
# 4. Update weights by the amount of agreement.
|
||||
#
|
||||
# - The scalar product :math:`\hat{u}_{j|i}\cdot v_j` can be considered as how well capsule :math:`i` agrees with :math:`j`. It is used to update
|
||||
# :math:`b_{ij}=b_{ij}+\hat{u}_{j|i}\cdot v_j`
|
||||
|
||||
import dgl.function as fn
|
||||
|
||||
|
||||
class DGLRoutingLayer(nn.Module):
|
||||
def __init__(self, in_nodes, out_nodes, f_size):
|
||||
super(DGLRoutingLayer, self).__init__()
|
||||
self.g = init_graph(in_nodes, out_nodes, f_size)
|
||||
self.in_nodes = in_nodes
|
||||
self.out_nodes = out_nodes
|
||||
self.in_indx = list(range(in_nodes))
|
||||
self.out_indx = list(range(in_nodes, in_nodes + out_nodes))
|
||||
|
||||
def forward(self, u_hat, routing_num=1):
|
||||
self.g.edata["u_hat"] = u_hat
|
||||
|
||||
for r in range(routing_num):
|
||||
# step 1 (line 4): normalize over out edges
|
||||
edges_b = self.g.edata["b"].view(self.in_nodes, self.out_nodes)
|
||||
self.g.edata["c"] = F.softmax(edges_b, dim=1).view(-1, 1)
|
||||
self.g.edata["c u_hat"] = self.g.edata["c"] * self.g.edata["u_hat"]
|
||||
|
||||
# Execute step 1 & 2
|
||||
self.g.update_all(fn.copy_e("c u_hat", "m"), fn.sum("m", "s"))
|
||||
|
||||
# step 3 (line 6)
|
||||
self.g.nodes[self.out_indx].data["v"] = self.squash(
|
||||
self.g.nodes[self.out_indx].data["s"], dim=1
|
||||
)
|
||||
|
||||
# step 4 (line 7)
|
||||
v = th.cat(
|
||||
[self.g.nodes[self.out_indx].data["v"]] * self.in_nodes, dim=0
|
||||
)
|
||||
self.g.edata["b"] = self.g.edata["b"] + (
|
||||
self.g.edata["u_hat"] * v
|
||||
).sum(dim=1, keepdim=True)
|
||||
|
||||
@staticmethod
|
||||
def squash(s, dim=1):
|
||||
sq = th.sum(s**2, dim=dim, keepdim=True)
|
||||
s_norm = th.sqrt(sq)
|
||||
s = (sq / (1.0 + sq)) * (s / s_norm)
|
||||
return s
|
||||
|
||||
|
||||
############################################################################################################
|
||||
# Step 3: Testing
|
||||
# ~~~~~~~~~~~~~~~
|
||||
#
|
||||
# Make a simple 20x10 capsule layer.
|
||||
in_nodes = 20
|
||||
out_nodes = 10
|
||||
f_size = 4
|
||||
u_hat = th.randn(in_nodes * out_nodes, f_size)
|
||||
routing = DGLRoutingLayer(in_nodes, out_nodes, f_size)
|
||||
|
||||
############################################################################################################
|
||||
# You can visualize a Capsule network's behavior by monitoring the entropy
|
||||
# of coupling coefficients. They should start high and then drop, as the
|
||||
# weights gradually concentrate on fewer edges.
|
||||
entropy_list = []
|
||||
dist_list = []
|
||||
|
||||
for i in range(10):
|
||||
routing(u_hat)
|
||||
dist_matrix = routing.g.edata["c"].view(in_nodes, out_nodes)
|
||||
entropy = (-dist_matrix * th.log(dist_matrix)).sum(dim=1)
|
||||
entropy_list.append(entropy.data.numpy())
|
||||
dist_list.append(dist_matrix.data.numpy())
|
||||
stds = np.std(entropy_list, axis=1)
|
||||
means = np.mean(entropy_list, axis=1)
|
||||
plt.errorbar(np.arange(len(entropy_list)), means, stds, marker="o")
|
||||
plt.ylabel("Entropy of Weight Distribution")
|
||||
plt.xlabel("Number of Routing")
|
||||
plt.xticks(np.arange(len(entropy_list)))
|
||||
plt.close()
|
||||
############################################################################################################
|
||||
# |image3|
|
||||
#
|
||||
# Alternatively, we can also watch the evolution of histograms.
|
||||
|
||||
import matplotlib.animation as animation
|
||||
import seaborn as sns
|
||||
|
||||
fig = plt.figure(dpi=150)
|
||||
fig.clf()
|
||||
ax = fig.subplots()
|
||||
|
||||
|
||||
def dist_animate(i):
|
||||
ax.cla()
|
||||
sns.distplot(dist_list[i].reshape(-1), kde=False, ax=ax)
|
||||
ax.set_xlabel("Weight Distribution Histogram")
|
||||
ax.set_title("Routing: %d" % (i))
|
||||
|
||||
|
||||
ani = animation.FuncAnimation(
|
||||
fig, dist_animate, frames=len(entropy_list), interval=500
|
||||
)
|
||||
plt.close()
|
||||
|
||||
############################################################################################################
|
||||
# |image4|
|
||||
#
|
||||
# You can monitor the how lower-level Capsules gradually attach to one of the
|
||||
# higher level ones.
|
||||
import networkx as nx
|
||||
from networkx.algorithms import bipartite
|
||||
|
||||
g = routing.g.to_networkx()
|
||||
X, Y = bipartite.sets(g)
|
||||
height_in = 10
|
||||
height_out = height_in * 0.8
|
||||
height_in_y = np.linspace(0, height_in, in_nodes)
|
||||
height_out_y = np.linspace((height_in - height_out) / 2, height_out, out_nodes)
|
||||
pos = dict()
|
||||
|
||||
fig2 = plt.figure(figsize=(8, 3), dpi=150)
|
||||
fig2.clf()
|
||||
ax = fig2.subplots()
|
||||
pos.update(
|
||||
(n, (i, 1)) for i, n in zip(height_in_y, X)
|
||||
) # put nodes from X at x=1
|
||||
pos.update(
|
||||
(n, (i, 2)) for i, n in zip(height_out_y, Y)
|
||||
) # put nodes from Y at x=2
|
||||
|
||||
|
||||
def weight_animate(i):
|
||||
ax.cla()
|
||||
ax.axis("off")
|
||||
ax.set_title("Routing: %d " % i)
|
||||
dm = dist_list[i]
|
||||
nx.draw_networkx_nodes(
|
||||
g, pos, nodelist=range(in_nodes), node_color="r", node_size=100, ax=ax
|
||||
)
|
||||
nx.draw_networkx_nodes(
|
||||
g,
|
||||
pos,
|
||||
nodelist=range(in_nodes, in_nodes + out_nodes),
|
||||
node_color="b",
|
||||
node_size=100,
|
||||
ax=ax,
|
||||
)
|
||||
for edge in g.edges():
|
||||
nx.draw_networkx_edges(
|
||||
g,
|
||||
pos,
|
||||
edgelist=[edge],
|
||||
width=dm[edge[0], edge[1] - in_nodes] * 1.5,
|
||||
ax=ax,
|
||||
)
|
||||
|
||||
|
||||
ani2 = animation.FuncAnimation(
|
||||
fig2, weight_animate, frames=len(dist_list), interval=500
|
||||
)
|
||||
plt.close()
|
||||
|
||||
############################################################################################################
|
||||
# |image5|
|
||||
#
|
||||
# The full code of this visualization is provided on
|
||||
# `GitHub <https://github.com/dmlc/dgl/blob/master/examples/pytorch/capsule/simple_routing.py>`__. The complete
|
||||
# code that trains on MNIST is also on `GitHub <https://github.com/dmlc/dgl/tree/tutorial/examples/pytorch/capsule>`__.
|
||||
#
|
||||
# .. |image0| image:: https://i.imgur.com/55Ovkdh.png
|
||||
# .. |image1| image:: https://i.imgur.com/9tc6GLl.png
|
||||
# .. |image2| image:: https://i.imgur.com/mv1W9Rv.png
|
||||
# .. |image3| image:: https://i.imgur.com/dMvu7p3.png
|
||||
# .. |image4| image:: https://github.com/VoVAllen/DGL_Capsule/raw/master/routing_dist.gif
|
||||
# .. |image5| image:: https://github.com/VoVAllen/DGL_Capsule/raw/master/routing_vis.gif
|
||||
@@ -0,0 +1,888 @@
|
||||
"""
|
||||
.. _model-transformer:
|
||||
|
||||
Transformer as a Graph Neural Network
|
||||
======================================
|
||||
|
||||
**Author**: Zihao Ye, Jinjing Zhou, Qipeng Guo, Quan Gan, Zheng Zhang
|
||||
|
||||
.. warning::
|
||||
|
||||
The tutorial aims at gaining insights into the paper, with code as a mean
|
||||
of explanation. The implementation thus is NOT optimized for running
|
||||
efficiency. For recommended implementation, please refer to the `official
|
||||
examples <https://github.com/dmlc/dgl/tree/master/examples>`_.
|
||||
|
||||
"""
|
||||
################################################################################################
|
||||
# In this tutorial, you learn about a simplified implementation of the Transformer model.
|
||||
# You can see highlights of the most important design points. For instance, there is
|
||||
# only single-head attention. The complete code can be found
|
||||
# `here <https://github.com/dmlc/dgl/tree/master/examples/pytorch/transformer>`__.
|
||||
#
|
||||
# The overall structure is similar to the one from the research papaer `Annotated
|
||||
# Transformer <http://nlp.seas.harvard.edu/2018/04/03/attention.html>`__.
|
||||
#
|
||||
# The Transformer model, as a replacement of CNN/RNN architecture for
|
||||
# sequence modeling, was introduced in the research paper: `Attention is All
|
||||
# You Need <https://arxiv.org/pdf/1706.03762.pdf>`__. It improved the
|
||||
# state of the art for machine translation as well as natural language
|
||||
# inference task
|
||||
# (`GPT <https://s3-us-west-2.amazonaws.com/openai-assets/research-covers/language-unsupervised/language_understanding_paper.pdf>`__).
|
||||
# Recent work on pre-training Transformer with large scale corpus
|
||||
# (`BERT <https://arxiv.org/pdf/1810.04805.pdf>`__) supports that it is
|
||||
# capable of learning high-quality semantic representation.
|
||||
#
|
||||
# The interesting part of Transformer is its extensive employment of
|
||||
# attention. The classic use of attention comes from machine translation
|
||||
# model, where the output token attends to all input tokens.
|
||||
#
|
||||
# Transformer additionally applies *self-attention* in both decoder and
|
||||
# encoder. This process forces words relate to each other to combine
|
||||
# together, irrespective of their positions in the sequence. This is
|
||||
# different from RNN-based model, where words (in the source sentence) are
|
||||
# combined along the chain, which is thought to be too constrained.
|
||||
#
|
||||
# Attention layer of Transformer
|
||||
# ------------------------------
|
||||
#
|
||||
# In the attention layer of Transformer, for each node the module learns to
|
||||
# assign weights on its in-coming edges. For node pair :math:`(i, j)`
|
||||
# (from :math:`i` to :math:`j`) with node
|
||||
# :math:`x_i, x_j \in \mathbb{R}^n`, the score of their connection is
|
||||
# defined as follows:
|
||||
#
|
||||
# .. math::
|
||||
#
|
||||
#
|
||||
# q_j = W_q\cdot x_j \\
|
||||
# k_i = W_k\cdot x_i\\
|
||||
# v_i = W_v\cdot x_i\\
|
||||
# \textrm{score} = q_j^T k_i
|
||||
#
|
||||
# where :math:`W_q, W_k, W_v \in \mathbb{R}^{n\times d_k}` map the
|
||||
# representations :math:`x` to “query”, “key”, and “value” space
|
||||
# respectively.
|
||||
#
|
||||
# There are other possibilities to implement the score function. The dot
|
||||
# product measures the similarity of a given query :math:`q_j` and a key
|
||||
# :math:`k_i`: if :math:`j` needs the information stored in :math:`i`, the
|
||||
# query vector at position :math:`j` (:math:`q_j`) is supposed to be close
|
||||
# to key vector at position :math:`i` (:math:`k_i`).
|
||||
#
|
||||
# The score is then used to compute the sum of the incoming values,
|
||||
# normalized over the weights of edges, stored in :math:`\textrm{wv}`.
|
||||
# Then apply an affine layer to :math:`\textrm{wv}` to get the output
|
||||
# :math:`o`:
|
||||
#
|
||||
# .. math::
|
||||
#
|
||||
#
|
||||
# w_{ji} = \frac{\exp\{\textrm{score}_{ji} \}}{\sum\limits_{(k, i)\in E}\exp\{\textrm{score}_{ki} \}} \\
|
||||
# \textrm{wv}_i = \sum_{(k, i)\in E} w_{ki} v_k \\
|
||||
# o = W_o\cdot \textrm{wv} \\
|
||||
#
|
||||
# Multi-head attention layer
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
#
|
||||
# In Transformer, attention is *multi-headed*. A head is very much like a
|
||||
# channel in a convolutional network. The multi-head attention consists of
|
||||
# multiple attention heads, in which each head refers to a single
|
||||
# attention module. :math:`\textrm{wv}^{(i)}` for all the heads are
|
||||
# concatenated and mapped to output :math:`o` with an affine layer:
|
||||
#
|
||||
# .. math::
|
||||
#
|
||||
#
|
||||
# o = W_o \cdot \textrm{concat}\left(\left[\textrm{wv}^{(0)}, \textrm{wv}^{(1)}, \cdots, \textrm{wv}^{(h)}\right]\right)
|
||||
#
|
||||
# The code below wraps necessary components for multi-head attention, and
|
||||
# provides two interfaces.
|
||||
#
|
||||
# - ``get`` maps state ‘x’, to query, key and value, which is required by
|
||||
# following steps(\ ``propagate_attention``).
|
||||
# - ``get_o`` maps the updated value after attention to the output
|
||||
# :math:`o` for post-processing.
|
||||
#
|
||||
# .. code::
|
||||
#
|
||||
# class MultiHeadAttention(nn.Module):
|
||||
# "Multi-Head Attention"
|
||||
# def __init__(self, h, dim_model):
|
||||
# "h: number of heads; dim_model: hidden dimension"
|
||||
# super(MultiHeadAttention, self).__init__()
|
||||
# self.d_k = dim_model // h
|
||||
# self.h = h
|
||||
# # W_q, W_k, W_v, W_o
|
||||
# self.linears = clones(nn.Linear(dim_model, dim_model), 4)
|
||||
#
|
||||
# def get(self, x, fields='qkv'):
|
||||
# "Return a dict of queries / keys / values."
|
||||
# batch_size = x.shape[0]
|
||||
# ret = {}
|
||||
# if 'q' in fields:
|
||||
# ret['q'] = self.linears[0](x).view(batch_size, self.h, self.d_k)
|
||||
# if 'k' in fields:
|
||||
# ret['k'] = self.linears[1](x).view(batch_size, self.h, self.d_k)
|
||||
# if 'v' in fields:
|
||||
# ret['v'] = self.linears[2](x).view(batch_size, self.h, self.d_k)
|
||||
# return ret
|
||||
#
|
||||
# def get_o(self, x):
|
||||
# "get output of the multi-head attention"
|
||||
# batch_size = x.shape[0]
|
||||
# return self.linears[3](x.view(batch_size, -1))
|
||||
#
|
||||
#
|
||||
# How DGL implements Transformer with a graph neural network
|
||||
# ----------------------------------------------------------
|
||||
#
|
||||
# You get a different perspective of Transformer by treating the
|
||||
# attention as edges in a graph and adopt message passing on the edges to
|
||||
# induce the appropriate processing.
|
||||
#
|
||||
# Graph structure
|
||||
# ~~~~~~~~~~~~~~~
|
||||
#
|
||||
# Construct the graph by mapping tokens of the source and target
|
||||
# sentence to nodes. The complete Transformer graph is made up of three
|
||||
# subgraphs:
|
||||
#
|
||||
# **Source language graph**. This is a complete graph, each
|
||||
# token :math:`s_i` can attend to any other token :math:`s_j` (including
|
||||
# self-loops). |image0|
|
||||
# **Target language graph**. The graph is
|
||||
# half-complete, in that :math:`t_i` attends only to :math:`t_j` if
|
||||
# :math:`i > j` (an output token can not depend on future words). |image1|
|
||||
# **Cross-language graph**. This is a bi-partitie graph, where there is
|
||||
# an edge from every source token :math:`s_i` to every target token
|
||||
# :math:`t_j`, meaning every target token can attend on source tokens.
|
||||
# |image2|
|
||||
#
|
||||
# The full picture looks like this: |image3|
|
||||
#
|
||||
# Pre-build the graphs in dataset preparation stage.
|
||||
#
|
||||
# Message passing
|
||||
# ~~~~~~~~~~~~~~~
|
||||
#
|
||||
# Once you define the graph structure, move on to defining the
|
||||
# computation for message passing.
|
||||
#
|
||||
# Assuming that you have already computed all the queries :math:`q_i`, keys
|
||||
# :math:`k_i` and values :math:`v_i`. For each node :math:`i` (no matter
|
||||
# whether it is a source token or target token), you can decompose the
|
||||
# attention computation into two steps:
|
||||
#
|
||||
# 1. **Message computation:** Compute attention score
|
||||
# :math:`\mathrm{score}_{ij}` between :math:`i` and all nodes :math:`j`
|
||||
# to be attended over, by taking the scaled-dot product between
|
||||
# :math:`q_i` and :math:`k_j`. The message sent from :math:`j` to
|
||||
# :math:`i` will consist of the score :math:`\mathrm{score}_{ij}` and
|
||||
# the value :math:`v_j`.
|
||||
# 2. **Message aggregation:** Aggregate the values :math:`v_j` from all
|
||||
# :math:`j` according to the scores :math:`\mathrm{score}_{ij}`.
|
||||
#
|
||||
# Simple implementation
|
||||
# ^^^^^^^^^^^^^^^^^^^^
|
||||
#
|
||||
# Message computation
|
||||
# '''''''''''''''''''
|
||||
#
|
||||
# Compute ``score`` and send source node’s ``v`` to destination’s mailbox
|
||||
#
|
||||
# .. code::
|
||||
#
|
||||
# def message_func(edges):
|
||||
# return {'score': ((edges.src['k'] * edges.dst['q'])
|
||||
# .sum(-1, keepdim=True)),
|
||||
# 'v': edges.src['v']}
|
||||
#
|
||||
# Message aggregation
|
||||
# '''''''''''''''''''
|
||||
#
|
||||
# Normalize over all in-edges and weighted sum to get output
|
||||
#
|
||||
# .. code::
|
||||
#
|
||||
# import torch as th
|
||||
# import torch.nn.functional as F
|
||||
#
|
||||
# def reduce_func(nodes, d_k=64):
|
||||
# v = nodes.mailbox['v']
|
||||
# att = F.softmax(nodes.mailbox['score'] / th.sqrt(d_k), 1)
|
||||
# return {'dx': (att * v).sum(1)}
|
||||
#
|
||||
# Execute on specific edges
|
||||
# '''''''''''''''''''''''''
|
||||
#
|
||||
# .. code::
|
||||
#
|
||||
# import functools.partial as partial
|
||||
# def naive_propagate_attention(self, g, eids):
|
||||
# g.send_and_recv(eids, message_func, partial(reduce_func, d_k=self.d_k))
|
||||
#
|
||||
# Speeding up with built-in functions
|
||||
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
#
|
||||
# To speed up the message passing process, use DGL’s built-in
|
||||
# functions, including:
|
||||
#
|
||||
# - ``fn.src_mul_egdes(src_field, edges_field, out_field)`` multiplies
|
||||
# source’s attribute and edges attribute, and send the result to the
|
||||
# destination node’s mailbox keyed by ``out_field``.
|
||||
# - ``fn.copy_e(edges_field, out_field)`` copies edge’s attribute to
|
||||
# destination node’s mailbox.
|
||||
# - ``fn.sum(edges_field, out_field)`` sums up
|
||||
# edge’s attribute and sends aggregation to destination node’s mailbox.
|
||||
#
|
||||
# Here, you assemble those built-in functions into ``propagate_attention``,
|
||||
# which is also the main graph operation function in the final
|
||||
# implementation. To accelerate it, break the ``softmax`` operation into
|
||||
# the following steps. Recall that for each head there are two phases.
|
||||
#
|
||||
# 1. Compute attention score by multiply src node’s ``k`` and dst node’s
|
||||
# ``q``
|
||||
#
|
||||
# - ``g.apply_edges(src_dot_dst('k', 'q', 'score'), eids)``
|
||||
#
|
||||
# 2. Scaled Softmax over all dst nodes’ in-coming edges
|
||||
#
|
||||
# - Step 1: Exponentialize score with scale normalize constant
|
||||
#
|
||||
# - ``g.apply_edges(scaled_exp('score', np.sqrt(self.d_k)))``
|
||||
#
|
||||
# .. math:: \textrm{score}_{ij}\leftarrow\exp{\left(\frac{\textrm{score}_{ij}}{ \sqrt{d_k}}\right)}
|
||||
#
|
||||
# - Step 2: Get the “values” on associated nodes weighted by “scores”
|
||||
# on in-coming edges of each node; get the sum of “scores” on
|
||||
# in-coming edges of each node for normalization. Note that here
|
||||
# :math:`\textrm{wv}` is not normalized.
|
||||
#
|
||||
# - ``msg: fn.u_mul_e('v', 'score', 'v'), reduce: fn.sum('v', 'wv')``
|
||||
#
|
||||
# .. math:: \textrm{wv}_j=\sum_{i=1}^{N} \textrm{score}_{ij} \cdot v_i
|
||||
#
|
||||
# - ``msg: fn.copy_e('score', 'score'), reduce: fn.sum('score', 'z')``
|
||||
#
|
||||
# .. math:: \textrm{z}_j=\sum_{i=1}^{N} \textrm{score}_{ij}
|
||||
#
|
||||
# The normalization of :math:`\textrm{wv}` is left to post processing.
|
||||
#
|
||||
# .. code::
|
||||
#
|
||||
# def src_dot_dst(src_field, dst_field, out_field):
|
||||
# def func(edges):
|
||||
# return {out_field: (edges.src[src_field] * edges.dst[dst_field]).sum(-1, keepdim=True)}
|
||||
#
|
||||
# return func
|
||||
#
|
||||
# def scaled_exp(field, scale_constant):
|
||||
# def func(edges):
|
||||
# # clamp for softmax numerical stability
|
||||
# return {field: th.exp((edges.data[field] / scale_constant).clamp(-5, 5))}
|
||||
#
|
||||
# return func
|
||||
#
|
||||
#
|
||||
# def propagate_attention(self, g, eids):
|
||||
# # Compute attention score
|
||||
# g.apply_edges(src_dot_dst('k', 'q', 'score'), eids)
|
||||
# g.apply_edges(scaled_exp('score', np.sqrt(self.d_k)))
|
||||
# # Update node state
|
||||
# g.send_and_recv(eids,
|
||||
# [fn.u_mul_e('v', 'score', 'v'), fn.copy_e('score', 'score')],
|
||||
# [fn.sum('v', 'wv'), fn.sum('score', 'z')])
|
||||
#
|
||||
# Preprocessing and postprocessing
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
#
|
||||
# In Transformer, data needs to be pre- and post-processed before and
|
||||
# after the ``propagate_attention`` function.
|
||||
#
|
||||
# **Preprocessing** The preprocessing function ``pre_func`` first
|
||||
# normalizes the node representations and then map them to a set of
|
||||
# queries, keys and values, using self-attention as an example:
|
||||
#
|
||||
# .. math::
|
||||
#
|
||||
#
|
||||
# x \leftarrow \textrm{LayerNorm}(x) \\
|
||||
# [q, k, v] \leftarrow [W_q, W_k, W_v ]\cdot x
|
||||
#
|
||||
# **Postprocessing** The postprocessing function ``post_funcs`` completes
|
||||
# the whole computation correspond to one layer of the transformer: 1.
|
||||
# Normalize :math:`\textrm{wv}` and get the output of Multi-Head Attention
|
||||
# Layer :math:`o`.
|
||||
#
|
||||
# .. math::
|
||||
#
|
||||
#
|
||||
# \textrm{wv} \leftarrow \frac{\textrm{wv}}{z} \\
|
||||
# o \leftarrow W_o\cdot \textrm{wv} + b_o
|
||||
#
|
||||
# add residual connection:
|
||||
#
|
||||
# .. math::
|
||||
#
|
||||
#
|
||||
# x \leftarrow x + o
|
||||
#
|
||||
# 2. Applying a two layer position-wise feed forward layer on :math:`x`
|
||||
# then add residual connection:
|
||||
#
|
||||
# .. math::
|
||||
#
|
||||
#
|
||||
# x \leftarrow x + \textrm{LayerNorm}(\textrm{FFN}(x))
|
||||
#
|
||||
# where :math:`\textrm{FFN}` refers to the feed forward function.
|
||||
#
|
||||
# .. code::
|
||||
#
|
||||
# class Encoder(nn.Module):
|
||||
# def __init__(self, layer, N):
|
||||
# super(Encoder, self).__init__()
|
||||
# self.N = N
|
||||
# self.layers = clones(layer, N)
|
||||
# self.norm = LayerNorm(layer.size)
|
||||
#
|
||||
# def pre_func(self, i, fields='qkv'):
|
||||
# layer = self.layers[i]
|
||||
# def func(nodes):
|
||||
# x = nodes.data['x']
|
||||
# norm_x = layer.sublayer[0].norm(x)
|
||||
# return layer.self_attn.get(norm_x, fields=fields)
|
||||
# return func
|
||||
#
|
||||
# def post_func(self, i):
|
||||
# layer = self.layers[i]
|
||||
# def func(nodes):
|
||||
# x, wv, z = nodes.data['x'], nodes.data['wv'], nodes.data['z']
|
||||
# o = layer.self_attn.get_o(wv / z)
|
||||
# x = x + layer.sublayer[0].dropout(o)
|
||||
# x = layer.sublayer[1](x, layer.feed_forward)
|
||||
# return {'x': x if i < self.N - 1 else self.norm(x)}
|
||||
# return func
|
||||
#
|
||||
# class Decoder(nn.Module):
|
||||
# def __init__(self, layer, N):
|
||||
# super(Decoder, self).__init__()
|
||||
# self.N = N
|
||||
# self.layers = clones(layer, N)
|
||||
# self.norm = LayerNorm(layer.size)
|
||||
#
|
||||
# def pre_func(self, i, fields='qkv', l=0):
|
||||
# layer = self.layers[i]
|
||||
# def func(nodes):
|
||||
# x = nodes.data['x']
|
||||
# if fields == 'kv':
|
||||
# norm_x = x # In enc-dec attention, x has already been normalized.
|
||||
# else:
|
||||
# norm_x = layer.sublayer[l].norm(x)
|
||||
# return layer.self_attn.get(norm_x, fields)
|
||||
# return func
|
||||
#
|
||||
# def post_func(self, i, l=0):
|
||||
# layer = self.layers[i]
|
||||
# def func(nodes):
|
||||
# x, wv, z = nodes.data['x'], nodes.data['wv'], nodes.data['z']
|
||||
# o = layer.self_attn.get_o(wv / z)
|
||||
# x = x + layer.sublayer[l].dropout(o)
|
||||
# if l == 1:
|
||||
# x = layer.sublayer[2](x, layer.feed_forward)
|
||||
# return {'x': x if i < self.N - 1 else self.norm(x)}
|
||||
# return func
|
||||
#
|
||||
# This completes all procedures of one layer of encoder and decoder in
|
||||
# Transformer.
|
||||
#
|
||||
# .. note::
|
||||
#
|
||||
# The sublayer connection part is little bit different from the
|
||||
# original paper. However, this implementation is the same as `The Annotated
|
||||
# Transformer <http://nlp.seas.harvard.edu/2018/04/03/attention.html>`__
|
||||
# and
|
||||
# `OpenNMT <https://github.com/OpenNMT/OpenNMT-py/blob/cd29c1dbfb35f4a2701ff52a1bf4e5bdcf02802e/onmt/encoders/transformer.py>`__.
|
||||
#
|
||||
# Main class of Transformer graph
|
||||
# -------------------------------
|
||||
#
|
||||
# The processing flow of Transformer can be seen as a 2-stage
|
||||
# message-passing within the complete graph (adding pre- and post-
|
||||
# processing appropriately): 1) self-attention in encoder, 2)
|
||||
# self-attention in decoder followed by cross-attention between encoder
|
||||
# and decoder, as shown below. |image4|
|
||||
#
|
||||
# .. code:: python
|
||||
#
|
||||
# class Transformer(nn.Module):
|
||||
# def __init__(self, encoder, decoder, src_embed, tgt_embed, pos_enc, generator, h, d_k):
|
||||
# super(Transformer, self).__init__()
|
||||
# self.encoder, self.decoder = encoder, decoder
|
||||
# self.src_embed, self.tgt_embed = src_embed, tgt_embed
|
||||
# self.pos_enc = pos_enc
|
||||
# self.generator = generator
|
||||
# self.h, self.d_k = h, d_k
|
||||
#
|
||||
# def propagate_attention(self, g, eids):
|
||||
# # Compute attention score
|
||||
# g.apply_edges(src_dot_dst('k', 'q', 'score'), eids)
|
||||
# g.apply_edges(scaled_exp('score', np.sqrt(self.d_k)))
|
||||
# # Send weighted values to target nodes
|
||||
# g.send_and_recv(eids,
|
||||
# [fn.u_mul_e('v', 'score', 'v'), fn.copy_e('score', 'score')],
|
||||
# [fn.sum('v', 'wv'), fn.sum('score', 'z')])
|
||||
#
|
||||
# def update_graph(self, g, eids, pre_pairs, post_pairs):
|
||||
# "Update the node states and edge states of the graph."
|
||||
#
|
||||
# # Pre-compute queries and key-value pairs.
|
||||
# for pre_func, nids in pre_pairs:
|
||||
# g.apply_nodes(pre_func, nids)
|
||||
# self.propagate_attention(g, eids)
|
||||
# # Further calculation after attention mechanism
|
||||
# for post_func, nids in post_pairs:
|
||||
# g.apply_nodes(post_func, nids)
|
||||
#
|
||||
# def forward(self, graph):
|
||||
# g = graph.g
|
||||
# nids, eids = graph.nids, graph.eids
|
||||
#
|
||||
# # Word Embedding and Position Embedding
|
||||
# src_embed, src_pos = self.src_embed(graph.src[0]), self.pos_enc(graph.src[1])
|
||||
# tgt_embed, tgt_pos = self.tgt_embed(graph.tgt[0]), self.pos_enc(graph.tgt[1])
|
||||
# g.nodes[nids['enc']].data['x'] = self.pos_enc.dropout(src_embed + src_pos)
|
||||
# g.nodes[nids['dec']].data['x'] = self.pos_enc.dropout(tgt_embed + tgt_pos)
|
||||
#
|
||||
# for i in range(self.encoder.N):
|
||||
# # Step 1: Encoder Self-attention
|
||||
# pre_func = self.encoder.pre_func(i, 'qkv')
|
||||
# post_func = self.encoder.post_func(i)
|
||||
# nodes, edges = nids['enc'], eids['ee']
|
||||
# self.update_graph(g, edges, [(pre_func, nodes)], [(post_func, nodes)])
|
||||
#
|
||||
# for i in range(self.decoder.N):
|
||||
# # Step 2: Dncoder Self-attention
|
||||
# pre_func = self.decoder.pre_func(i, 'qkv')
|
||||
# post_func = self.decoder.post_func(i)
|
||||
# nodes, edges = nids['dec'], eids['dd']
|
||||
# self.update_graph(g, edges, [(pre_func, nodes)], [(post_func, nodes)])
|
||||
# # Step 3: Encoder-Decoder attention
|
||||
# pre_q = self.decoder.pre_func(i, 'q', 1)
|
||||
# pre_kv = self.decoder.pre_func(i, 'kv', 1)
|
||||
# post_func = self.decoder.post_func(i, 1)
|
||||
# nodes_e, nodes_d, edges = nids['enc'], nids['dec'], eids['ed']
|
||||
# self.update_graph(g, edges, [(pre_q, nodes_d), (pre_kv, nodes_e)], [(post_func, nodes_d)])
|
||||
#
|
||||
# return self.generator(g.ndata['x'][nids['dec']])
|
||||
#
|
||||
#
|
||||
# .. note::
|
||||
#
|
||||
# By calling ``update_graph`` function, you can create your own
|
||||
# Transformer on any subgraphs with nearly the same code. This
|
||||
# flexibility enables us to discover new, sparse structures (c.f. local attention
|
||||
# mentioned `here <https://arxiv.org/pdf/1508.04025.pdf>`__). Note in this
|
||||
# implementation you don't use mask or padding, which makes the logic
|
||||
# more clear and saves memory. The trade-off is that the implementation is
|
||||
# slower.
|
||||
#
|
||||
# Training
|
||||
# --------
|
||||
#
|
||||
# This tutorial does not cover several other techniques such as Label
|
||||
# Smoothing and Noam Optimizations mentioned in the original paper. For
|
||||
# detailed description about these modules, read `The
|
||||
# Annotated
|
||||
# Transformer <http://nlp.seas.harvard.edu/2018/04/03/attention.html>`__
|
||||
# written by Harvard NLP team.
|
||||
#
|
||||
# Task and the dataset
|
||||
# ~~~~~~~~~~~~~~~~~~~~
|
||||
#
|
||||
# The Transformer is a general framework for a variety of NLP tasks. This tutorial focuses
|
||||
# on the sequence to sequence learning: it’s a typical case to illustrate how it works.
|
||||
#
|
||||
# As for the dataset, there are two example tasks: copy and sort, together
|
||||
# with two real-world translation tasks: multi30k en-de task and wmt14
|
||||
# en-de task.
|
||||
#
|
||||
# - **copy dataset**: copy input sequences to output. (train/valid/test:
|
||||
# 9000, 1000, 1000)
|
||||
# - **sort dataset**: sort input sequences as output. (train/valid/test:
|
||||
# 9000, 1000, 1000)
|
||||
# - **Multi30k en-de**, translate sentences from En to De.
|
||||
# (train/valid/test: 29000, 1000, 1000)
|
||||
# - **WMT14 en-de**, translate sentences from En to De.
|
||||
# (Train/Valid/Test: 4500966/3000/3003)
|
||||
#
|
||||
# .. note::
|
||||
# Training with wmt14 requires multi-GPU support and is not available. Contributions are welcome!
|
||||
#
|
||||
# Graph building
|
||||
# ~~~~~~~~~~~~~~
|
||||
#
|
||||
# **Batching** This is similar to the way you handle Tree-LSTM. Build a graph pool in
|
||||
# advance, including all possible combination of input lengths and output
|
||||
# lengths. Then for each sample in a batch, call ``dgl.batch`` to batch
|
||||
# graphs of their sizes together in to a single large graph.
|
||||
#
|
||||
# You can wrap the process of creating graph pool and building
|
||||
# BatchedGraph in ``dataset.GraphPool`` and
|
||||
# ``dataset.TranslationDataset``.
|
||||
#
|
||||
# .. code:: python
|
||||
#
|
||||
# graph_pool = GraphPool()
|
||||
#
|
||||
# data_iter = dataset(graph_pool, mode='train', batch_size=1, devices=devices)
|
||||
# for graph in data_iter:
|
||||
# print(graph.nids['enc']) # encoder node ids
|
||||
# print(graph.nids['dec']) # decoder node ids
|
||||
# print(graph.eids['ee']) # encoder-encoder edge ids
|
||||
# print(graph.eids['ed']) # encoder-decoder edge ids
|
||||
# print(graph.eids['dd']) # decoder-decoder edge ids
|
||||
# print(graph.src[0]) # Input word index list
|
||||
# print(graph.src[1]) # Input positions
|
||||
# print(graph.tgt[0]) # Output word index list
|
||||
# print(graph.tgt[1]) # Ouptut positions
|
||||
# break
|
||||
#
|
||||
# Output:
|
||||
#
|
||||
# .. code::
|
||||
#
|
||||
# tensor([0, 1, 2, 3, 4, 5, 6, 7, 8], device='cuda:0')
|
||||
# tensor([ 9, 10, 11, 12, 13, 14, 15, 16, 17, 18], device='cuda:0')
|
||||
# tensor([ 0, 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, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71,
|
||||
# 72, 73, 74, 75, 76, 77, 78, 79, 80], device='cuda:0')
|
||||
# tensor([ 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94,
|
||||
# 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108,
|
||||
# 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122,
|
||||
# 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136,
|
||||
# 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150,
|
||||
# 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164,
|
||||
# 165, 166, 167, 168, 169, 170], device='cuda:0')
|
||||
# tensor([171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181, 182, 183, 184,
|
||||
# 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198,
|
||||
# 199, 200, 201, 202, 203, 204, 205, 206, 207, 208, 209, 210, 211, 212,
|
||||
# 213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 224, 225],
|
||||
# device='cuda:0')
|
||||
# tensor([28, 25, 7, 26, 6, 4, 5, 9, 18], device='cuda:0')
|
||||
# tensor([0, 1, 2, 3, 4, 5, 6, 7, 8], device='cuda:0')
|
||||
# tensor([ 0, 28, 25, 7, 26, 6, 4, 5, 9, 18], device='cuda:0')
|
||||
# tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], device='cuda:0')
|
||||
#
|
||||
# Put it all together
|
||||
# -------------------
|
||||
#
|
||||
# Train a one-head transformer with one layer, 128 dimension on copy
|
||||
# task. Set other parameters to the default.
|
||||
#
|
||||
# Inference module is not included in this tutorial. It
|
||||
# requires beam search. For a full implementation, see the `GitHub
|
||||
# repo <https://github.com/dmlc/dgl/tree/master/examples/pytorch/transformer>`__.
|
||||
#
|
||||
# .. code:: python
|
||||
#
|
||||
# from tqdm.auto import tqdm
|
||||
# import torch as th
|
||||
# import numpy as np
|
||||
#
|
||||
# from loss import LabelSmoothing, SimpleLossCompute
|
||||
# from modules import make_model
|
||||
# from optims import NoamOpt
|
||||
# from dgl.contrib.transformer import get_dataset, GraphPool
|
||||
#
|
||||
# def run_epoch(data_iter, model, loss_compute, is_train=True):
|
||||
# for i, g in tqdm(enumerate(data_iter)):
|
||||
# with th.set_grad_enabled(is_train):
|
||||
# output = model(g)
|
||||
# loss = loss_compute(output, g.tgt_y, g.n_tokens)
|
||||
# print('average loss: {}'.format(loss_compute.avg_loss))
|
||||
# print('accuracy: {}'.format(loss_compute.accuracy))
|
||||
#
|
||||
# N = 1
|
||||
# batch_size = 128
|
||||
# devices = ['cuda' if th.cuda.is_available() else 'cpu']
|
||||
#
|
||||
# dataset = get_dataset("copy")
|
||||
# V = dataset.vocab_size
|
||||
# criterion = LabelSmoothing(V, padding_idx=dataset.pad_id, smoothing=0.1)
|
||||
# dim_model = 128
|
||||
#
|
||||
# # Create model
|
||||
# model = make_model(V, V, N=N, dim_model=128, dim_ff=128, h=1)
|
||||
#
|
||||
# # Sharing weights between Encoder & Decoder
|
||||
# model.src_embed.lut.weight = model.tgt_embed.lut.weight
|
||||
# model.generator.proj.weight = model.tgt_embed.lut.weight
|
||||
#
|
||||
# model, criterion = model.to(devices[0]), criterion.to(devices[0])
|
||||
# model_opt = NoamOpt(dim_model, 1, 400,
|
||||
# th.optim.Adam(model.parameters(), lr=1e-3, betas=(0.9, 0.98), eps=1e-9))
|
||||
# loss_compute = SimpleLossCompute
|
||||
#
|
||||
# att_maps = []
|
||||
# for epoch in range(4):
|
||||
# train_iter = dataset(graph_pool, mode='train', batch_size=batch_size, devices=devices)
|
||||
# valid_iter = dataset(graph_pool, mode='valid', batch_size=batch_size, devices=devices)
|
||||
# print('Epoch: {} Training...'.format(epoch))
|
||||
# model.train(True)
|
||||
# run_epoch(train_iter, model,
|
||||
# loss_compute(criterion, model_opt), is_train=True)
|
||||
# print('Epoch: {} Evaluating...'.format(epoch))
|
||||
# model.att_weight_map = None
|
||||
# model.eval()
|
||||
# run_epoch(valid_iter, model,
|
||||
# loss_compute(criterion, None), is_train=False)
|
||||
# att_maps.append(model.att_weight_map)
|
||||
#
|
||||
# Visualization
|
||||
# -------------
|
||||
#
|
||||
# After training, you can visualize the attention that the Transformer generates
|
||||
# on copy task.
|
||||
#
|
||||
# .. code:: python
|
||||
#
|
||||
# src_seq = dataset.get_seq_by_id(VIZ_IDX, mode='valid', field='src')
|
||||
# tgt_seq = dataset.get_seq_by_id(VIZ_IDX, mode='valid', field='tgt')[:-1]
|
||||
# # visualize head 0 of encoder-decoder attention
|
||||
# att_animation(att_maps, 'e2d', src_seq, tgt_seq, 0)
|
||||
#
|
||||
# |image5| from the figure you see the decoder nodes gradually learns to
|
||||
# attend to corresponding nodes in input sequence, which is the expected
|
||||
# behavior.
|
||||
#
|
||||
# Multi-head attention
|
||||
# ~~~~~~~~~~~~~~~~~~~~
|
||||
#
|
||||
# Besides the attention of a one-head attention trained on toy task. We
|
||||
# also visualize the attention scores of Encoder’s Self Attention,
|
||||
# Decoder’s Self Attention and the Encoder-Decoder attention of an
|
||||
# one-Layer Transformer network trained on multi-30k dataset.
|
||||
#
|
||||
# From the visualization you see the diversity of different heads, which is what you would
|
||||
# expect. Different heads learn different relations between word pairs.
|
||||
#
|
||||
# - **Encoder Self-Attention** |image6|
|
||||
#
|
||||
# - **Encoder-Decoder Attention** Most words in target sequence attend on
|
||||
# their related words in source sequence, for example: when generating
|
||||
# “See” (in De), several heads attend on “lake”; when generating
|
||||
# “Eisfischerhütte”, several heads attend on “ice”. |image7|
|
||||
#
|
||||
# - **Decoder Self-Attention** Most words attend on their previous few
|
||||
# words. |image8|
|
||||
#
|
||||
# Adaptive Universal Transformer
|
||||
# ------------------------------
|
||||
#
|
||||
# A recent research paper by Google, `Universal
|
||||
# Transformer <https://arxiv.org/pdf/1807.03819.pdf>`__, is an example to
|
||||
# show how ``update_graph`` adapts to more complex updating rules.
|
||||
#
|
||||
# The Universal Transformer was proposed to address the problem that
|
||||
# vanilla Transformer is not computationally universal by introducing
|
||||
# recurrence in Transformer:
|
||||
#
|
||||
# - The basic idea of Universal Transformer is to repeatedly revise its
|
||||
# representations of all symbols in the sequence with each recurrent
|
||||
# step by applying a Transformer layer on the representations.
|
||||
# - Compared to vanilla Transformer, Universal Transformer shares weights
|
||||
# among its layers, and it does not fix the recurrence time (which
|
||||
# means the number of layers in Transformer).
|
||||
#
|
||||
# A further optimization employs an `adaptive computation time
|
||||
# (ACT) <https://arxiv.org/pdf/1603.08983.pdf>`__ mechanism to allow the
|
||||
# model to dynamically adjust the number of times the representation of
|
||||
# each position in a sequence is revised (refereed to as **step**
|
||||
# hereafter). This model is also known as the Adaptive Universal
|
||||
# Transformer (AUT).
|
||||
#
|
||||
# In AUT, you maintain an active nodes list. In each step :math:`t`, we
|
||||
# compute a halting probability: :math:`h (0<h<1)` for all nodes in this
|
||||
# list by:
|
||||
#
|
||||
# .. math:: h^t_i = \sigma(W_h x^t_i + b_h)
|
||||
#
|
||||
# then dynamically decide which nodes are still active. A node is halted
|
||||
# at time :math:`T` if and only if
|
||||
# :math:`\sum_{t=1}^{T-1} h_t < 1 - \varepsilon \leq \sum_{t=1}^{T}h_t`.
|
||||
# Halted nodes are removed from the list. The procedure proceeds until the
|
||||
# list is empty or a pre-defined maximum step is reached. From DGL’s
|
||||
# perspective, this means that the “active” graph becomes sparser over
|
||||
# time.
|
||||
#
|
||||
# The final state of a node :math:`s_i` is a weighted average of
|
||||
# :math:`x_i^t` by :math:`h_i^t`:
|
||||
#
|
||||
# .. math:: s_i = \sum_{t=1}^{T} h_i^t\cdot x_i^t
|
||||
#
|
||||
# In DGL, implement an algorithm by calling
|
||||
# ``update_graph`` on nodes that are still active and edges associated
|
||||
# with this nodes. The following code shows the Universal Transformer
|
||||
# class in DGL:
|
||||
#
|
||||
# .. code::
|
||||
#
|
||||
# class UTransformer(nn.Module):
|
||||
# "Universal Transformer(https://arxiv.org/pdf/1807.03819.pdf) with ACT(https://arxiv.org/pdf/1603.08983.pdf)."
|
||||
# MAX_DEPTH = 8
|
||||
# thres = 0.99
|
||||
# act_loss_weight = 0.01
|
||||
# def __init__(self, encoder, decoder, src_embed, tgt_embed, pos_enc, time_enc, generator, h, d_k):
|
||||
# super(UTransformer, self).__init__()
|
||||
# self.encoder, self.decoder = encoder, decoder
|
||||
# self.src_embed, self.tgt_embed = src_embed, tgt_embed
|
||||
# self.pos_enc, self.time_enc = pos_enc, time_enc
|
||||
# self.halt_enc = HaltingUnit(h * d_k)
|
||||
# self.halt_dec = HaltingUnit(h * d_k)
|
||||
# self.generator = generator
|
||||
# self.h, self.d_k = h, d_k
|
||||
#
|
||||
# def step_forward(self, nodes):
|
||||
# # add positional encoding and time encoding, increment step by one
|
||||
# x = nodes.data['x']
|
||||
# step = nodes.data['step']
|
||||
# pos = nodes.data['pos']
|
||||
# return {'x': self.pos_enc.dropout(x + self.pos_enc(pos.view(-1)) + self.time_enc(step.view(-1))),
|
||||
# 'step': step + 1}
|
||||
#
|
||||
# def halt_and_accum(self, name, end=False):
|
||||
# "field: 'enc' or 'dec'"
|
||||
# halt = self.halt_enc if name == 'enc' else self.halt_dec
|
||||
# thres = self.thres
|
||||
# def func(nodes):
|
||||
# p = halt(nodes.data['x'])
|
||||
# sum_p = nodes.data['sum_p'] + p
|
||||
# active = (sum_p < thres) & (1 - end)
|
||||
# _continue = active.float()
|
||||
# r = nodes.data['r'] * (1 - _continue) + (1 - sum_p) * _continue
|
||||
# s = nodes.data['s'] + ((1 - _continue) * r + _continue * p) * nodes.data['x']
|
||||
# return {'p': p, 'sum_p': sum_p, 'r': r, 's': s, 'active': active}
|
||||
# return func
|
||||
#
|
||||
# def propagate_attention(self, g, eids):
|
||||
# # Compute attention score
|
||||
# g.apply_edges(src_dot_dst('k', 'q', 'score'), eids)
|
||||
# g.apply_edges(scaled_exp('score', np.sqrt(self.d_k)), eids)
|
||||
# # Send weighted values to target nodes
|
||||
# g.send_and_recv(eids,
|
||||
# [fn.u_mul_e('v', 'score', 'v'), fn.copy_e('score', 'score')],
|
||||
# [fn.sum('v', 'wv'), fn.sum('score', 'z')])
|
||||
#
|
||||
# def update_graph(self, g, eids, pre_pairs, post_pairs):
|
||||
# "Update the node states and edge states of the graph."
|
||||
# # Pre-compute queries and key-value pairs.
|
||||
# for pre_func, nids in pre_pairs:
|
||||
# g.apply_nodes(pre_func, nids)
|
||||
# self.propagate_attention(g, eids)
|
||||
# # Further calculation after attention mechanism
|
||||
# for post_func, nids in post_pairs:
|
||||
# g.apply_nodes(post_func, nids)
|
||||
#
|
||||
# def forward(self, graph):
|
||||
# g = graph.g
|
||||
# N, E = graph.n_nodes, graph.n_edges
|
||||
# nids, eids = graph.nids, graph.eids
|
||||
#
|
||||
# # embed & pos
|
||||
# g.nodes[nids['enc']].data['x'] = self.src_embed(graph.src[0])
|
||||
# g.nodes[nids['dec']].data['x'] = self.tgt_embed(graph.tgt[0])
|
||||
# g.nodes[nids['enc']].data['pos'] = graph.src[1]
|
||||
# g.nodes[nids['dec']].data['pos'] = graph.tgt[1]
|
||||
#
|
||||
# # init step
|
||||
# device = next(self.parameters()).device
|
||||
# g.ndata['s'] = th.zeros(N, self.h * self.d_k, dtype=th.float, device=device) # accumulated state
|
||||
# g.ndata['p'] = th.zeros(N, 1, dtype=th.float, device=device) # halting prob
|
||||
# g.ndata['r'] = th.ones(N, 1, dtype=th.float, device=device) # remainder
|
||||
# g.ndata['sum_p'] = th.zeros(N, 1, dtype=th.float, device=device) # sum of pondering values
|
||||
# g.ndata['step'] = th.zeros(N, 1, dtype=th.long, device=device) # step
|
||||
# g.ndata['active'] = th.ones(N, 1, dtype=th.uint8, device=device) # active
|
||||
#
|
||||
# for step in range(self.MAX_DEPTH):
|
||||
# pre_func = self.encoder.pre_func('qkv')
|
||||
# post_func = self.encoder.post_func()
|
||||
# nodes = g.filter_nodes(lambda v: v.data['active'].view(-1), nids['enc'])
|
||||
# if len(nodes) == 0: break
|
||||
# edges = g.filter_edges(lambda e: e.dst['active'].view(-1), eids['ee'])
|
||||
# end = step == self.MAX_DEPTH - 1
|
||||
# self.update_graph(g, edges,
|
||||
# [(self.step_forward, nodes), (pre_func, nodes)],
|
||||
# [(post_func, nodes), (self.halt_and_accum('enc', end), nodes)])
|
||||
#
|
||||
# g.nodes[nids['enc']].data['x'] = self.encoder.norm(g.nodes[nids['enc']].data['s'])
|
||||
#
|
||||
# for step in range(self.MAX_DEPTH):
|
||||
# pre_func = self.decoder.pre_func('qkv')
|
||||
# post_func = self.decoder.post_func()
|
||||
# nodes = g.filter_nodes(lambda v: v.data['active'].view(-1), nids['dec'])
|
||||
# if len(nodes) == 0: break
|
||||
# edges = g.filter_edges(lambda e: e.dst['active'].view(-1), eids['dd'])
|
||||
# self.update_graph(g, edges,
|
||||
# [(self.step_forward, nodes), (pre_func, nodes)],
|
||||
# [(post_func, nodes)])
|
||||
#
|
||||
# pre_q = self.decoder.pre_func('q', 1)
|
||||
# pre_kv = self.decoder.pre_func('kv', 1)
|
||||
# post_func = self.decoder.post_func(1)
|
||||
# nodes_e = nids['enc']
|
||||
# edges = g.filter_edges(lambda e: e.dst['active'].view(-1), eids['ed'])
|
||||
# end = step == self.MAX_DEPTH - 1
|
||||
# self.update_graph(g, edges,
|
||||
# [(pre_q, nodes), (pre_kv, nodes_e)],
|
||||
# [(post_func, nodes), (self.halt_and_accum('dec', end), nodes)])
|
||||
#
|
||||
# g.nodes[nids['dec']].data['x'] = self.decoder.norm(g.nodes[nids['dec']].data['s'])
|
||||
# act_loss = th.mean(g.ndata['r']) # ACT loss
|
||||
#
|
||||
# return self.generator(g.ndata['x'][nids['dec']]), act_loss * self.act_loss_weight
|
||||
#
|
||||
# Call ``filter_nodes`` and ``filter_edge`` to find nodes/edges
|
||||
# that are still active:
|
||||
#
|
||||
# .. note::
|
||||
#
|
||||
# - :func:`~dgl.DGLGraph.filter_nodes` takes a predicate and a node
|
||||
# ID list/tensor as input, then returns a tensor of node IDs that satisfy
|
||||
# the given predicate.
|
||||
# - :func:`~dgl.DGLGraph.filter_edges` takes a predicate
|
||||
# and an edge ID list/tensor as input, then returns a tensor of edge IDs
|
||||
# that satisfy the given predicate.
|
||||
#
|
||||
# For the full implementation, see the `GitHub
|
||||
# repo <https://github.com/dmlc/dgl/tree/master/examples/pytorch/transformer/modules/act.py>`__.
|
||||
#
|
||||
# The figure below shows the effect of Adaptive Computational
|
||||
# Time. Different positions of a sentence were revised different times.
|
||||
#
|
||||
# |image9|
|
||||
#
|
||||
# You can also visualize the dynamics of step distribution on nodes during the
|
||||
# training of AUT on sort task(reach 99.7% accuracy), which demonstrates
|
||||
# how AUT learns to reduce recurrence steps during training. |image10|
|
||||
#
|
||||
# .. |image0| image:: https://i.imgur.com/zV5LmTX.png
|
||||
# .. |image1| image:: https://i.imgur.com/dETQMMx.png
|
||||
# .. |image2| image:: https://i.imgur.com/hnGP229.png
|
||||
# .. |image3| image:: https://i.imgur.com/Hj2rRGT.png
|
||||
# .. |image4| image:: https://i.imgur.com/zlUpJ41.png
|
||||
# .. |image5| image:: https://s1.ax1x.com/2018/12/06/F126xI.gif
|
||||
# .. |image6| image:: https://i.imgur.com/HjYb7F2.png
|
||||
# .. |image7| image:: https://i.imgur.com/383J5O5.png
|
||||
# .. |image8| image:: https://i.imgur.com/c0UWB1V.png
|
||||
# .. |image9| image:: https://s1.ax1x.com/2018/12/06/F1sGod.png
|
||||
# .. |image10| image:: https://s1.ax1x.com/2018/12/06/F1r8Cq.gif
|
||||
#
|
||||
# .. note::
|
||||
# The notebook itself is not executable due to many dependencies.
|
||||
# Download `7_transformer.py <https://data.dgl.ai/tutorial/7_transformer.py>`__,
|
||||
# and copy the python script to directory ``examples/pytorch/transformer``
|
||||
# then run ``python 7_transformer.py`` to see how it works.
|
||||
@@ -0,0 +1,24 @@
|
||||
.. _tutorials4-index:
|
||||
|
||||
|
||||
Revisit classic models from a graph perspective
|
||||
-------------------------------------------------------
|
||||
|
||||
* **Capsule** `[paper] <https://arxiv.org/abs/1710.09829>`__ `[tutorial]
|
||||
<4_old_wines/2_capsule.html>`__ `[PyTorch code]
|
||||
<https://github.com/dmlc/dgl/tree/master/examples/pytorch/capsule>`__:
|
||||
This new computer vision model has two key ideas. First, enhancing the feature
|
||||
representation in a vector form (instead of a scalar) called *capsule*. Second,
|
||||
replacing max-pooling with dynamic routing. The idea of dynamic routing is to
|
||||
integrate a lower level capsule to one or several higher level capsules
|
||||
with non-parametric message-passing. A tutorial shows how the latter can be
|
||||
implemented with DGL APIs.
|
||||
|
||||
|
||||
* **Transformer** `[paper] <https://arxiv.org/abs/1706.03762>`__ `[tutorial] <4_old_wines/7_transformer.html>`__
|
||||
`[PyTorch code] <https://github.com/dmlc/dgl/tree/master/examples/pytorch/transformer>`__ and **Universal Transformer**
|
||||
`[paper] <https://arxiv.org/abs/1807.03819>`__ `[tutorial] <4_old_wines/7_transformer.html>`__
|
||||
`[PyTorch code] <https://github.com/dmlc/dgl/tree/master/examples/pytorch/transformer/modules/act.py>`__:
|
||||
These two models replace recurrent neural networks (RNNs) with several layers of multi-head attention to
|
||||
encode and discover structures among tokens of a sentence. These attention
|
||||
mechanisms are similarly formulated as graph operations with message-passing.
|
||||
@@ -0,0 +1,2 @@
|
||||
Paper Study with DGL
|
||||
=========================================
|
||||
Reference in New Issue
Block a user