chore: import upstream snapshot with attribution

tutorials/blitz/.gitignore

@@ -0,0 +1,2 @@
+*.dgl
+*.csv

tutorials/blitz/1_introduction.py

@@ -0,0 +1,228 @@
+"""
+Node Classification with DGL
+============================
+
+GNNs are powerful tools for many machine learning tasks on graphs. In
+this introductory tutorial, you will learn the basic workflow of using
+GNNs for node classification, i.e. predicting the category of a node in
+a graph.
+
+By completing this tutorial, you will be able to
+
+-  Load a DGL-provided dataset.
+-  Build a GNN model with DGL-provided neural network modules.
+-  Train and evaluate a GNN model for node classification on either CPU
+   or GPU.
+
+This tutorial assumes that you have experience in building neural
+networks with PyTorch.
+
+(Time estimate: 13 minutes)
+
+"""
+
+import os
+
+os.environ["DGLBACKEND"] = "pytorch"
+import dgl
+import dgl.data
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+######################################################################
+# Overview of Node Classification with GNN
+# ----------------------------------------
+#
+# One of the most popular and widely adopted tasks on graph data is node
+# classification, where a model needs to predict the ground truth category
+# of each node. Before graph neural networks, many proposed methods are
+# using either connectivity alone (such as DeepWalk or node2vec), or simple
+# combinations of connectivity and the node's own features.  GNNs, by
+# contrast, offers an opportunity to obtain node representations by
+# combining the connectivity and features of a *local neighborhood*.
+#
+# `Kipf et
+# al., <https://arxiv.org/abs/1609.02907>`__ is an example that formulates
+# the node classification problem as a semi-supervised node classification
+# task. With the help of only a small portion of labeled nodes, a graph
+# neural network (GNN) can accurately predict the node category of the
+# others.
+#
+# This tutorial will show how to build such a GNN for semi-supervised node
+# classification with only a small number of labels on the Cora
+# dataset,
+# a citation network with papers as nodes and citations as edges. The task
+# is to predict the category of a given paper. Each paper node contains a
+# word count vector as its features, normalized so that they sum up to one,
+# as described in Section 5.2 of
+# `the paper <https://arxiv.org/abs/1609.02907>`__.
+#
+# Loading Cora Dataset
+# --------------------
+#
+
+
+dataset = dgl.data.CoraGraphDataset()
+print(f"Number of categories: {dataset.num_classes}")
+
+
+######################################################################
+# A DGL Dataset object may contain one or multiple graphs. The Cora
+# dataset used in this tutorial only consists of one single graph.
+#
+
+g = dataset[0]
+
+
+######################################################################
+# A DGL graph can store node features and edge features in two
+# dictionary-like attributes called ``ndata`` and ``edata``.
+# In the DGL Cora dataset, the graph contains the following node features:
+#
+# - ``train_mask``: A boolean tensor indicating whether the node is in the
+#   training set.
+#
+# - ``val_mask``: A boolean tensor indicating whether the node is in the
+#   validation set.
+#
+# - ``test_mask``: A boolean tensor indicating whether the node is in the
+#   test set.
+#
+# - ``label``: The ground truth node category.
+#
+# -  ``feat``: The node features.
+#
+
+print("Node features")
+print(g.ndata)
+print("Edge features")
+print(g.edata)
+
+
+######################################################################
+# Defining a Graph Convolutional Network (GCN)
+# --------------------------------------------
+#
+# This tutorial will build a two-layer `Graph Convolutional Network
+# (GCN) <http://tkipf.github.io/graph-convolutional-networks/>`__. Each
+# layer computes new node representations by aggregating neighbor
+# information.
+#
+# To build a multi-layer GCN you can simply stack ``dgl.nn.GraphConv``
+# modules, which inherit ``torch.nn.Module``.
+#
+
+from dgl.nn import GraphConv
+
+
+class GCN(nn.Module):
+    def __init__(self, in_feats, h_feats, num_classes):
+        super(GCN, self).__init__()
+        self.conv1 = GraphConv(in_feats, h_feats)
+        self.conv2 = GraphConv(h_feats, num_classes)
+
+    def forward(self, g, in_feat):
+        h = self.conv1(g, in_feat)
+        h = F.relu(h)
+        h = self.conv2(g, h)
+        return h
+
+
+# Create the model with given dimensions
+model = GCN(g.ndata["feat"].shape[1], 16, dataset.num_classes)
+
+
+######################################################################
+# DGL provides implementation of many popular neighbor aggregation
+# modules. You can easily invoke them with one line of code.
+#
+
+
+######################################################################
+# Training the GCN
+# ----------------
+#
+# Training this GCN is similar to training other PyTorch neural networks.
+#
+
+
+def train(g, model):
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+    best_val_acc = 0
+    best_test_acc = 0
+
+    features = g.ndata["feat"]
+    labels = g.ndata["label"]
+    train_mask = g.ndata["train_mask"]
+    val_mask = g.ndata["val_mask"]
+    test_mask = g.ndata["test_mask"]
+    for e in range(100):
+        # Forward
+        logits = model(g, features)
+
+        # Compute prediction
+        pred = logits.argmax(1)
+
+        # Compute loss
+        # Note that you should only compute the losses of the nodes in the training set.
+        loss = F.cross_entropy(logits[train_mask], labels[train_mask])
+
+        # Compute accuracy on training/validation/test
+        train_acc = (pred[train_mask] == labels[train_mask]).float().mean()
+        val_acc = (pred[val_mask] == labels[val_mask]).float().mean()
+        test_acc = (pred[test_mask] == labels[test_mask]).float().mean()
+
+        # Save the best validation accuracy and the corresponding test accuracy.
+        if best_val_acc < val_acc:
+            best_val_acc = val_acc
+            best_test_acc = test_acc
+
+        # Backward
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        if e % 5 == 0:
+            print(
+                f"In epoch {e}, loss: {loss:.3f}, val acc: {val_acc:.3f} (best {best_val_acc:.3f}), test acc: {test_acc:.3f} (best {best_test_acc:.3f})"
+            )
+
+
+model = GCN(g.ndata["feat"].shape[1], 16, dataset.num_classes)
+train(g, model)
+
+
+######################################################################
+# Training on GPU
+# ---------------
+#
+# Training on GPU requires to put both the model and the graph onto GPU
+# with the ``to`` method, similar to what you will do in PyTorch.
+#
+# .. code:: python
+#
+#    g = g.to('cuda')
+#    model = GCN(g.ndata['feat'].shape[1], 16, dataset.num_classes).to('cuda')
+#    train(g, model)
+#
+
+
+######################################################################
+# What’s next?
+# ------------
+#
+# -  :doc:`How does DGL represent a graph <2_dglgraph>`?
+# -  :doc:`Write your own GNN module <3_message_passing>`.
+# -  :doc:`Link prediction (predicting existence of edges) on full
+#    graph <4_link_predict>`.
+# -  :doc:`Graph classification <5_graph_classification>`.
+# -  :doc:`Make your own dataset <6_load_data>`.
+# -  :ref:`The list of supported graph convolution
+#    modules <apinn-pytorch>`.
+# -  :ref:`The list of datasets provided by DGL <apidata>`.
+#
+
+
+# Thumbnail credits: Stanford CS224W Notes
+# sphinx_gallery_thumbnail_path = '_static/blitz_1_introduction.png'

tutorials/blitz/2_dglgraph.py

@@ -0,0 +1,236 @@
+"""
+How Does DGL Represent A Graph?
+===============================
+
+By the end of this tutorial you will be able to:
+
+-  Construct a graph in DGL from scratch.
+-  Assign node and edge features to a graph.
+-  Query properties of a DGL graph such as node degrees and
+   connectivity.
+-  Transform a DGL graph into another graph.
+-  Load and save DGL graphs.
+
+(Time estimate: 16 minutes)
+
+"""
+
+
+######################################################################
+# DGL Graph Construction
+# ----------------------
+#
+# DGL represents a directed graph as a ``DGLGraph`` object. You can
+# construct a graph by specifying the number of nodes in the graph as well
+# as the list of source and destination nodes.  Nodes in the graph have
+# consecutive IDs starting from 0.
+#
+# For instance, the following code constructs a directed star graph with 5
+# leaves. The center node's ID is 0. The edges go from the
+# center node to the leaves.
+#
+
+import os
+
+os.environ["DGLBACKEND"] = "pytorch"
+import dgl
+import numpy as np
+import torch
+
+g = dgl.graph(([0, 0, 0, 0, 0], [1, 2, 3, 4, 5]), num_nodes=6)
+# Equivalently, PyTorch LongTensors also work.
+g = dgl.graph(
+    (torch.LongTensor([0, 0, 0, 0, 0]), torch.LongTensor([1, 2, 3, 4, 5])),
+    num_nodes=6,
+)
+
+# You can omit the number of nodes argument if you can tell the number of nodes from the edge list alone.
+g = dgl.graph(([0, 0, 0, 0, 0], [1, 2, 3, 4, 5]))
+
+
+######################################################################
+# Edges in the graph have consecutive IDs starting from 0, and are
+# in the same order as the list of source and destination nodes during
+# creation.
+#
+
+# Print the source and destination nodes of every edge.
+print(g.edges())
+
+
+######################################################################
+# .. note::
+#
+#    ``DGLGraph``'s are always directed to best fit the computation
+#    pattern of graph neural networks, where the messages sent
+#    from one node to the other are often different between both
+#    directions. If you want to handle undirected graphs, you may consider
+#    treating it as a bidirectional graph. See `Graph
+#    Transformations`_ for an example of making
+#    a bidirectional graph.
+#
+
+
+######################################################################
+# Assigning Node and Edge Features to Graph
+# -----------------------------------------
+#
+# Many graph data contain attributes on nodes and edges.
+# Although the types of node and edge attributes can be arbitrary in real
+# world, ``DGLGraph`` only accepts attributes stored in tensors (with
+# numerical contents). Consequently, an attribute of all the nodes or
+# edges must have the same shape. In the context of deep learning, those
+# attributes are often called *features*.
+#
+# You can assign and retrieve node and edge features via ``ndata`` and
+# ``edata`` interface.
+#
+
+# Assign a 3-dimensional node feature vector for each node.
+g.ndata["x"] = torch.randn(6, 3)
+# Assign a 4-dimensional edge feature vector for each edge.
+g.edata["a"] = torch.randn(5, 4)
+# Assign a 5x4 node feature matrix for each node.  Node and edge features in DGL can be multi-dimensional.
+g.ndata["y"] = torch.randn(6, 5, 4)
+
+print(g.edata["a"])
+
+
+######################################################################
+# .. note::
+#
+#    The vast development of deep learning has provided us many
+#    ways to encode various types of attributes into numerical features.
+#    Here are some general suggestions:
+#
+#    -  For categorical attributes (e.g. gender, occupation), consider
+#       converting them to integers or one-hot encoding.
+#    -  For variable length string contents (e.g. news article, quote),
+#       consider applying a language model.
+#    -  For images, consider applying a vision model such as CNNs.
+#
+#    You can find plenty of materials on how to encode such attributes
+#    into a tensor in the `PyTorch Deep Learning
+#    Tutorials <https://pytorch.org/tutorials/>`__.
+#
+
+
+######################################################################
+# Querying Graph Structures
+# -------------------------
+#
+# ``DGLGraph`` object provides various methods to query a graph structure.
+#
+
+print(g.num_nodes())
+print(g.num_edges())
+# Out degrees of the center node
+print(g.out_degrees(0))
+# In degrees of the center node - note that the graph is directed so the in degree should be 0.
+print(g.in_degrees(0))
+
+
+######################################################################
+# Graph Transformations
+# ---------------------
+#
+
+
+######################################################################
+# DGL provides many APIs to transform a graph to another such as
+# extracting a subgraph:
+#
+
+# Induce a subgraph from node 0, node 1 and node 3 from the original graph.
+sg1 = g.subgraph([0, 1, 3])
+# Induce a subgraph from edge 0, edge 1 and edge 3 from the original graph.
+sg2 = g.edge_subgraph([0, 1, 3])
+
+
+######################################################################
+# You can obtain the node/edge mapping from the subgraph to the original
+# graph by looking into the node feature ``dgl.NID`` or edge feature
+# ``dgl.EID`` in the new graph.
+#
+
+# The original IDs of each node in sg1
+print(sg1.ndata[dgl.NID])
+# The original IDs of each edge in sg1
+print(sg1.edata[dgl.EID])
+# The original IDs of each node in sg2
+print(sg2.ndata[dgl.NID])
+# The original IDs of each edge in sg2
+print(sg2.edata[dgl.EID])
+
+
+######################################################################
+# ``subgraph`` and ``edge_subgraph`` also copies the original features
+# to the subgraph:
+#
+
+# The original node feature of each node in sg1
+print(sg1.ndata["x"])
+# The original edge feature of each node in sg1
+print(sg1.edata["a"])
+# The original node feature of each node in sg2
+print(sg2.ndata["x"])
+# The original edge feature of each node in sg2
+print(sg2.edata["a"])
+
+
+######################################################################
+# Another common transformation is to add a reverse edge for each edge in
+# the original graph with ``dgl.add_reverse_edges``.
+#
+# .. note::
+#
+#    If you have an undirected graph, it is better to convert it
+#    into a bidirectional graph first via adding reverse edges.
+#
+
+newg = dgl.add_reverse_edges(g)
+print(newg.edges())
+
+
+######################################################################
+# Loading and Saving Graphs
+# -------------------------
+#
+# You can save a graph or a list of graphs via ``dgl.save_graphs`` and
+# load them back with ``dgl.load_graphs``.
+#
+
+# Save graphs
+dgl.save_graphs("graph.dgl", g)
+dgl.save_graphs("graphs.dgl", [g, sg1, sg2])
+
+# Load graphs
+(g,), _ = dgl.load_graphs("graph.dgl")
+print(g)
+(g, sg1, sg2), _ = dgl.load_graphs("graphs.dgl")
+print(g)
+print(sg1)
+print(sg2)
+
+
+######################################################################
+# What’s next?
+# ------------
+#
+# -  See
+#    :ref:`here <apigraph-querying-graph-structure>`
+#    for a list of graph structure query APIs.
+# -  See
+#    :ref:`here <api-subgraph-extraction>`
+#    for a list of subgraph extraction routines.
+# -  See
+#    :ref:`here <api-transform>`
+#    for a list of graph transformation routines.
+# -  API reference of :func:`dgl.save_graphs`
+#    and
+#    :func:`dgl.load_graphs`
+#
+
+
+# Thumbnail credits: Wikipedia
+# sphinx_gallery_thumbnail_path = '_static/blitz_2_dglgraph.png'

tutorials/blitz/3_message_passing.py

@@ -0,0 +1,379 @@
+"""
+Write your own GNN module
+=========================
+
+Sometimes, your model goes beyond simply stacking existing GNN modules.
+For example, you would like to invent a new way of aggregating neighbor
+information by considering node importance or edge weights.
+
+By the end of this tutorial you will be able to
+
+-  Understand DGL’s message passing APIs.
+-  Implement GraphSAGE convolution module by your own.
+
+This tutorial assumes that you already know :doc:`the basics of training a
+GNN for node classification <1_introduction>`.
+
+(Time estimate: 10 minutes)
+
+"""
+
+import os
+
+os.environ["DGLBACKEND"] = "pytorch"
+import dgl
+import dgl.function as fn
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+######################################################################
+# Message passing and GNNs
+# ------------------------
+#
+# DGL follows the *message passing paradigm* inspired by the Message
+# Passing Neural Network proposed by `Gilmer et
+# al. <https://arxiv.org/abs/1704.01212>`__ Essentially, they found many
+# GNN models can fit into the following framework:
+#
+# .. math::
+#
+#
+#    m_{u\to v}^{(l)} = M^{(l)}\left(h_v^{(l-1)}, h_u^{(l-1)}, e_{u\to v}^{(l-1)}\right)
+#
+# .. math::
+#
+#
+#    m_{v}^{(l)} = \sum_{u\in\mathcal{N}(v)}m_{u\to v}^{(l)}
+#
+# .. math::
+#
+#
+#    h_v^{(l)} = U^{(l)}\left(h_v^{(l-1)}, m_v^{(l)}\right)
+#
+# where DGL calls :math:`M^{(l)}` the *message function*, :math:`\sum` the
+# *reduce function* and :math:`U^{(l)}` the *update function*. Note that
+# :math:`\sum` here can represent any function and is not necessarily a
+# summation.
+#
+
+
+######################################################################
+# For example, the `GraphSAGE convolution (Hamilton et al.,
+# 2017) <https://cs.stanford.edu/people/jure/pubs/graphsage-nips17.pdf>`__
+# takes the following mathematical form:
+#
+# .. math::
+#
+#
+#    h_{\mathcal{N}(v)}^k\leftarrow \text{Average}\{h_u^{k-1},\forall u\in\mathcal{N}(v)\}
+#
+# .. math::
+#
+#
+#    h_v^k\leftarrow \text{ReLU}\left(W^k\cdot \text{CONCAT}(h_v^{k-1}, h_{\mathcal{N}(v)}^k) \right)
+#
+# You can see that message passing is directional: the message sent from
+# one node :math:`u` to other node :math:`v` is not necessarily the same
+# as the other message sent from node :math:`v` to node :math:`u` in the
+# opposite direction.
+#
+# Although DGL has builtin support of GraphSAGE via
+# :class:`dgl.nn.SAGEConv <dgl.nn.pytorch.SAGEConv>`,
+# here is how you can implement GraphSAGE convolution in DGL by your own.
+#
+
+
+class SAGEConv(nn.Module):
+    """Graph convolution module used by the GraphSAGE model.
+
+    Parameters
+    ----------
+    in_feat : int
+        Input feature size.
+    out_feat : int
+        Output feature size.
+    """
+
+    def __init__(self, in_feat, out_feat):
+        super(SAGEConv, self).__init__()
+        # A linear submodule for projecting the input and neighbor feature to the output.
+        self.linear = nn.Linear(in_feat * 2, out_feat)
+
+    def forward(self, g, h):
+        """Forward computation
+
+        Parameters
+        ----------
+        g : Graph
+            The input graph.
+        h : Tensor
+            The input node feature.
+        """
+        with g.local_scope():
+            g.ndata["h"] = h
+            # update_all is a message passing API.
+            g.update_all(
+                message_func=fn.copy_u("h", "m"),
+                reduce_func=fn.mean("m", "h_N"),
+            )
+            h_N = g.ndata["h_N"]
+            h_total = torch.cat([h, h_N], dim=1)
+            return self.linear(h_total)
+
+
+######################################################################
+# The central piece in this code is the
+# :func:`g.update_all <dgl.DGLGraph.update_all>`
+# function, which gathers and averages the neighbor features. There are
+# three concepts here:
+#
+# * Message function ``fn.copy_u('h', 'm')`` that
+#   copies the node feature under name ``'h'`` as *messages* with name
+#   ``'m'`` sent to neighbors.
+#
+# * Reduce function ``fn.mean('m', 'h_N')`` that averages
+#   all the received messages under name ``'m'`` and saves the result as a
+#   new node feature ``'h_N'``.
+#
+# * ``update_all`` tells DGL to trigger the
+#   message and reduce functions for all the nodes and edges.
+#
+
+
+######################################################################
+# Afterwards, you can stack your own GraphSAGE convolution layers to form
+# a multi-layer GraphSAGE network.
+#
+
+
+class Model(nn.Module):
+    def __init__(self, in_feats, h_feats, num_classes):
+        super(Model, self).__init__()
+        self.conv1 = SAGEConv(in_feats, h_feats)
+        self.conv2 = SAGEConv(h_feats, num_classes)
+
+    def forward(self, g, in_feat):
+        h = self.conv1(g, in_feat)
+        h = F.relu(h)
+        h = self.conv2(g, h)
+        return h
+
+
+######################################################################
+# Training loop
+# ~~~~~~~~~~~~~
+# The following code for data loading and training loop is directly copied
+# from the introduction tutorial.
+#
+
+import dgl.data
+
+dataset = dgl.data.CoraGraphDataset()
+g = dataset[0]
+
+
+def train(g, model):
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+    all_logits = []
+    best_val_acc = 0
+    best_test_acc = 0
+
+    features = g.ndata["feat"]
+    labels = g.ndata["label"]
+    train_mask = g.ndata["train_mask"]
+    val_mask = g.ndata["val_mask"]
+    test_mask = g.ndata["test_mask"]
+    for e in range(200):
+        # Forward
+        logits = model(g, features)
+
+        # Compute prediction
+        pred = logits.argmax(1)
+
+        # Compute loss
+        # Note that we should only compute the losses of the nodes in the training set,
+        # i.e. with train_mask 1.
+        loss = F.cross_entropy(logits[train_mask], labels[train_mask])
+
+        # Compute accuracy on training/validation/test
+        train_acc = (pred[train_mask] == labels[train_mask]).float().mean()
+        val_acc = (pred[val_mask] == labels[val_mask]).float().mean()
+        test_acc = (pred[test_mask] == labels[test_mask]).float().mean()
+
+        # Save the best validation accuracy and the corresponding test accuracy.
+        if best_val_acc < val_acc:
+            best_val_acc = val_acc
+            best_test_acc = test_acc
+
+        # Backward
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        all_logits.append(logits.detach())
+
+        if e % 5 == 0:
+            print(
+                "In epoch {}, loss: {:.3f}, val acc: {:.3f} (best {:.3f}), test acc: {:.3f} (best {:.3f})".format(
+                    e, loss, val_acc, best_val_acc, test_acc, best_test_acc
+                )
+            )
+
+
+model = Model(g.ndata["feat"].shape[1], 16, dataset.num_classes)
+train(g, model)
+
+
+######################################################################
+# More customization
+# ------------------
+#
+# In DGL, we provide many built-in message and reduce functions under the
+# ``dgl.function`` package. You can find more details in :ref:`the API
+# doc <apifunction>`.
+#
+
+
+######################################################################
+# These APIs allow one to quickly implement new graph convolution modules.
+# For example, the following implements a new ``SAGEConv`` that aggregates
+# neighbor representations using a weighted average. Note that ``edata``
+# member can hold edge features which can also take part in message
+# passing.
+#
+
+
+class WeightedSAGEConv(nn.Module):
+    """Graph convolution module used by the GraphSAGE model with edge weights.
+
+    Parameters
+    ----------
+    in_feat : int
+        Input feature size.
+    out_feat : int
+        Output feature size.
+    """
+
+    def __init__(self, in_feat, out_feat):
+        super(WeightedSAGEConv, self).__init__()
+        # A linear submodule for projecting the input and neighbor feature to the output.
+        self.linear = nn.Linear(in_feat * 2, out_feat)
+
+    def forward(self, g, h, w):
+        """Forward computation
+
+        Parameters
+        ----------
+        g : Graph
+            The input graph.
+        h : Tensor
+            The input node feature.
+        w : Tensor
+            The edge weight.
+        """
+        with g.local_scope():
+            g.ndata["h"] = h
+            g.edata["w"] = w
+            g.update_all(
+                message_func=fn.u_mul_e("h", "w", "m"),
+                reduce_func=fn.mean("m", "h_N"),
+            )
+            h_N = g.ndata["h_N"]
+            h_total = torch.cat([h, h_N], dim=1)
+            return self.linear(h_total)
+
+
+######################################################################
+# Because the graph in this dataset does not have edge weights, we
+# manually assign all edge weights to one in the ``forward()`` function of
+# the model. You can replace it with your own edge weights.
+#
+
+
+class Model(nn.Module):
+    def __init__(self, in_feats, h_feats, num_classes):
+        super(Model, self).__init__()
+        self.conv1 = WeightedSAGEConv(in_feats, h_feats)
+        self.conv2 = WeightedSAGEConv(h_feats, num_classes)
+
+    def forward(self, g, in_feat):
+        h = self.conv1(g, in_feat, torch.ones(g.num_edges(), 1).to(g.device))
+        h = F.relu(h)
+        h = self.conv2(g, h, torch.ones(g.num_edges(), 1).to(g.device))
+        return h
+
+
+model = Model(g.ndata["feat"].shape[1], 16, dataset.num_classes)
+train(g, model)
+
+
+######################################################################
+# Even more customization by user-defined function
+# ------------------------------------------------
+#
+# DGL allows user-defined message and reduce function for the maximal
+# expressiveness. Here is a user-defined message function that is
+# equivalent to ``fn.u_mul_e('h', 'w', 'm')``.
+#
+
+
+def u_mul_e_udf(edges):
+    return {"m": edges.src["h"] * edges.data["w"]}
+
+
+######################################################################
+# ``edges`` has three members: ``src``, ``data`` and ``dst``, representing
+# the source node feature, edge feature, and destination node feature for
+# all edges.
+#
+
+
+######################################################################
+# You can also write your own reduce function. For example, the following
+# is equivalent to the builtin ``fn.mean('m', 'h_N')`` function that averages
+# the incoming messages:
+#
+
+
+def mean_udf(nodes):
+    return {"h_N": nodes.mailbox["m"].mean(1)}
+
+
+######################################################################
+# In short, DGL will group the nodes by their in-degrees, and for each
+# group DGL stacks the incoming messages along the second dimension. You
+# can then perform a reduction along the second dimension to aggregate
+# messages.
+#
+# For more details on customizing message and reduce function with
+# user-defined function, please refer to the :ref:`API
+# reference <apiudf>`.
+#
+
+
+######################################################################
+# Best practice of writing custom GNN modules
+# -------------------------------------------
+#
+# DGL recommends the following practice ranked by preference:
+#
+# -  Use ``dgl.nn`` modules.
+# -  Use ``dgl.nn.functional`` functions which contain lower-level complex
+#    operations such as computing a softmax for each node over incoming
+#    edges.
+# -  Use ``update_all`` with builtin message and reduce functions.
+# -  Use user-defined message or reduce functions.
+#
+
+
+######################################################################
+# What’s next?
+# ------------
+#
+# -  :ref:`Writing Efficient Message Passing
+#    Code <guide-message-passing-efficient>`.
+#
+
+
+# Thumbnail credits: Representation Learning on Networks, Jure Leskovec, WWW 2018
+# sphinx_gallery_thumbnail_path = '_static/blitz_3_message_passing.png'

tutorials/blitz/4_link_predict.py

@@ -0,0 +1,364 @@
+"""
+Link Prediction using Graph Neural Networks
+===========================================
+
+In the :doc:`introduction <1_introduction>`, you have already learned
+the basic workflow of using GNNs for node classification,
+i.e. predicting the category of a node in a graph. This tutorial will
+teach you how to train a GNN for link prediction, i.e. predicting the
+existence of an edge between two arbitrary nodes in a graph.
+
+By the end of this tutorial you will be able to
+
+-  Build a GNN-based link prediction model.
+-  Train and evaluate the model on a small DGL-provided dataset.
+
+(Time estimate: 28 minutes)
+
+"""
+
+import itertools
+import os
+
+os.environ["DGLBACKEND"] = "pytorch"
+
+import dgl
+import dgl.data
+import numpy as np
+import scipy.sparse as sp
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+######################################################################
+# Overview of Link Prediction with GNN
+# ------------------------------------
+#
+# Many applications such as social recommendation, item recommendation,
+# knowledge graph completion, etc., can be formulated as link prediction,
+# which predicts whether an edge exists between two particular nodes. This
+# tutorial shows an example of predicting whether a citation relationship,
+# either citing or being cited, between two papers exists in a citation
+# network.
+#
+# This tutorial formulates the link prediction problem as a binary classification
+# problem as follows:
+#
+# -  Treat the edges in the graph as *positive examples*.
+# -  Sample a number of non-existent edges (i.e. node pairs with no edges
+#    between them) as *negative* examples.
+# -  Divide the positive examples and negative examples into a training
+#    set and a test set.
+# -  Evaluate the model with any binary classification metric such as Area
+#    Under Curve (AUC).
+#
+# .. note::
+#
+#    The practice comes from
+#    `SEAL <https://papers.nips.cc/paper/2018/file/53f0d7c537d99b3824f0f99d62ea2428-Paper.pdf>`__,
+#    although the model here does not use their idea of node labeling.
+#
+# In some domains such as large-scale recommender systems or information
+# retrieval, you may favor metrics that emphasize good performance of
+# top-K predictions. In these cases you may want to consider other metrics
+# such as mean average precision, and use other negative sampling methods,
+# which are beyond the scope of this tutorial.
+#
+# Loading graph and features
+# --------------------------
+#
+# Following the :doc:`introduction <1_introduction>`, this tutorial
+# first loads the Cora dataset.
+#
+
+
+dataset = dgl.data.CoraGraphDataset()
+g = dataset[0]
+
+
+######################################################################
+# Prepare training and testing sets
+# ---------------------------------
+#
+# This tutorial randomly picks 10% of the edges for positive examples in
+# the test set, and leave the rest for the training set. It then samples
+# the same number of edges for negative examples in both sets.
+#
+
+# Split edge set for training and testing
+u, v = g.edges()
+
+eids = np.arange(g.num_edges())
+eids = np.random.permutation(eids)
+test_size = int(len(eids) * 0.1)
+train_size = g.num_edges() - test_size
+test_pos_u, test_pos_v = u[eids[:test_size]], v[eids[:test_size]]
+train_pos_u, train_pos_v = u[eids[test_size:]], v[eids[test_size:]]
+
+# Find all negative edges and split them for training and testing
+adj = sp.coo_matrix((np.ones(len(u)), (u.numpy(), v.numpy())))
+adj_neg = 1 - adj.todense() - np.eye(g.num_nodes())
+neg_u, neg_v = np.where(adj_neg != 0)
+
+neg_eids = np.random.choice(len(neg_u), g.num_edges())
+test_neg_u, test_neg_v = (
+    neg_u[neg_eids[:test_size]],
+    neg_v[neg_eids[:test_size]],
+)
+train_neg_u, train_neg_v = (
+    neg_u[neg_eids[test_size:]],
+    neg_v[neg_eids[test_size:]],
+)
+
+
+######################################################################
+# When training, you will need to remove the edges in the test set from
+# the original graph. You can do this via ``dgl.remove_edges``.
+#
+# .. note::
+#
+#    ``dgl.remove_edges`` works by creating a subgraph from the
+#    original graph, resulting in a copy and therefore could be slow for
+#    large graphs. If so, you could save the training and test graph to
+#    disk, as you would do for preprocessing.
+#
+
+train_g = dgl.remove_edges(g, eids[:test_size])
+
+
+######################################################################
+# Define a GraphSAGE model
+# ------------------------
+#
+# This tutorial builds a model consisting of two
+# `GraphSAGE <https://arxiv.org/abs/1706.02216>`__ layers, each computes
+# new node representations by averaging neighbor information. DGL provides
+# ``dgl.nn.SAGEConv`` that conveniently creates a GraphSAGE layer.
+#
+
+from dgl.nn import SAGEConv
+
+
+# ----------- 2. create model -------------- #
+# build a two-layer GraphSAGE model
+class GraphSAGE(nn.Module):
+    def __init__(self, in_feats, h_feats):
+        super(GraphSAGE, self).__init__()
+        self.conv1 = SAGEConv(in_feats, h_feats, "mean")
+        self.conv2 = SAGEConv(h_feats, h_feats, "mean")
+
+    def forward(self, g, in_feat):
+        h = self.conv1(g, in_feat)
+        h = F.relu(h)
+        h = self.conv2(g, h)
+        return h
+
+
+######################################################################
+# The model then predicts the probability of existence of an edge by
+# computing a score between the representations of both incident nodes
+# with a function (e.g. an MLP or a dot product), which you will see in
+# the next section.
+#
+# .. math::
+#
+#
+#    \hat{y}_{u\sim v} = f(h_u, h_v)
+#
+
+
+######################################################################
+# Positive graph, negative graph, and ``apply_edges``
+# ---------------------------------------------------
+#
+# In previous tutorials you have learned how to compute node
+# representations with a GNN. However, link prediction requires you to
+# compute representation of *pairs of nodes*.
+#
+# DGL recommends you to treat the pairs of nodes as another graph, since
+# you can describe a pair of nodes with an edge. In link prediction, you
+# will have a *positive graph* consisting of all the positive examples as
+# edges, and a *negative graph* consisting of all the negative examples.
+# The *positive graph* and the *negative graph* will contain the same set
+# of nodes as the original graph.  This makes it easier to pass node
+# features among multiple graphs for computation.  As you will see later,
+# you can directly feed the node representations computed on the entire
+# graph to the positive and the negative graphs for computing pair-wise
+# scores.
+#
+# The following code constructs the positive graph and the negative graph
+# for the training set and the test set respectively.
+#
+
+train_pos_g = dgl.graph((train_pos_u, train_pos_v), num_nodes=g.num_nodes())
+train_neg_g = dgl.graph((train_neg_u, train_neg_v), num_nodes=g.num_nodes())
+
+test_pos_g = dgl.graph((test_pos_u, test_pos_v), num_nodes=g.num_nodes())
+test_neg_g = dgl.graph((test_neg_u, test_neg_v), num_nodes=g.num_nodes())
+
+
+######################################################################
+# The benefit of treating the pairs of nodes as a graph is that you can
+# use the ``DGLGraph.apply_edges`` method, which conveniently computes new
+# edge features based on the incident nodes’ features and the original
+# edge features (if applicable).
+#
+# DGL provides a set of optimized builtin functions to compute new
+# edge features based on the original node/edge features. For example,
+# ``dgl.function.u_dot_v`` computes a dot product of the incident nodes’
+# representations for each edge.
+#
+
+import dgl.function as fn
+
+
+class DotPredictor(nn.Module):
+    def forward(self, g, h):
+        with g.local_scope():
+            g.ndata["h"] = h
+            # Compute a new edge feature named 'score' by a dot-product between the
+            # source node feature 'h' and destination node feature 'h'.
+            g.apply_edges(fn.u_dot_v("h", "h", "score"))
+            # u_dot_v returns a 1-element vector for each edge so you need to squeeze it.
+            return g.edata["score"][:, 0]
+
+
+######################################################################
+# You can also write your own function if it is complex.
+# For instance, the following module produces a scalar score on each edge
+# by concatenating the incident nodes’ features and passing it to an MLP.
+#
+
+
+class MLPPredictor(nn.Module):
+    def __init__(self, h_feats):
+        super().__init__()
+        self.W1 = nn.Linear(h_feats * 2, h_feats)
+        self.W2 = nn.Linear(h_feats, 1)
+
+    def apply_edges(self, edges):
+        """
+        Computes a scalar score for each edge of the given graph.
+
+        Parameters
+        ----------
+        edges :
+            Has three members ``src``, ``dst`` and ``data``, each of
+            which is a dictionary representing the features of the
+            source nodes, the destination nodes, and the edges
+            themselves.
+
+        Returns
+        -------
+        dict
+            A dictionary of new edge features.
+        """
+        h = torch.cat([edges.src["h"], edges.dst["h"]], 1)
+        return {"score": self.W2(F.relu(self.W1(h))).squeeze(1)}
+
+    def forward(self, g, h):
+        with g.local_scope():
+            g.ndata["h"] = h
+            g.apply_edges(self.apply_edges)
+            return g.edata["score"]
+
+
+######################################################################
+# .. note::
+#
+#    The builtin functions are optimized for both speed and memory.
+#    We recommend using builtin functions whenever possible.
+#
+# .. note::
+#
+#    If you have read the :doc:`message passing
+#    tutorial <3_message_passing>`, you will notice that the
+#    argument ``apply_edges`` takes has exactly the same form as a message
+#    function in ``update_all``.
+#
+
+
+######################################################################
+# Training loop
+# -------------
+#
+# After you defined the node representation computation and the edge score
+# computation, you can go ahead and define the overall model, loss
+# function, and evaluation metric.
+#
+# The loss function is simply binary cross entropy loss.
+#
+# .. math::
+#
+#
+#    \mathcal{L} = -\sum_{u\sim v\in \mathcal{D}}\left( y_{u\sim v}\log(\hat{y}_{u\sim v}) + (1-y_{u\sim v})\log(1-\hat{y}_{u\sim v})) \right)
+#
+# The evaluation metric in this tutorial is AUC.
+#
+
+model = GraphSAGE(train_g.ndata["feat"].shape[1], 16)
+# You can replace DotPredictor with MLPPredictor.
+# pred = MLPPredictor(16)
+pred = DotPredictor()
+
+
+def compute_loss(pos_score, neg_score):
+    scores = torch.cat([pos_score, neg_score])
+    labels = torch.cat(
+        [torch.ones(pos_score.shape[0]), torch.zeros(neg_score.shape[0])]
+    )
+    return F.binary_cross_entropy_with_logits(scores, labels)
+
+
+def compute_auc(pos_score, neg_score):
+    scores = torch.cat([pos_score, neg_score]).numpy()
+    labels = torch.cat(
+        [torch.ones(pos_score.shape[0]), torch.zeros(neg_score.shape[0])]
+    ).numpy()
+    return roc_auc_score(labels, scores)
+
+
+######################################################################
+# The training loop goes as follows:
+#
+# .. note::
+#
+#    This tutorial does not include evaluation on a validation
+#    set. In practice you should save and evaluate the best model based on
+#    performance on the validation set.
+#
+
+# ----------- 3. set up loss and optimizer -------------- #
+# in this case, loss will in training loop
+optimizer = torch.optim.Adam(
+    itertools.chain(model.parameters(), pred.parameters()), lr=0.01
+)
+
+# ----------- 4. training -------------------------------- #
+all_logits = []
+for e in range(100):
+    # forward
+    h = model(train_g, train_g.ndata["feat"])
+    pos_score = pred(train_pos_g, h)
+    neg_score = pred(train_neg_g, h)
+    loss = compute_loss(pos_score, neg_score)
+
+    # backward
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+
+    if e % 5 == 0:
+        print("In epoch {}, loss: {}".format(e, loss))
+
+# ----------- 5. check results ------------------------ #
+from sklearn.metrics import roc_auc_score
+
+with torch.no_grad():
+    pos_score = pred(test_pos_g, h)
+    neg_score = pred(test_neg_g, h)
+    print("AUC", compute_auc(pos_score, neg_score))
+
+
+# Thumbnail credits: Link Prediction with Neo4j, Mark Needham
+# sphinx_gallery_thumbnail_path = '_static/blitz_4_link_predict.png'

tutorials/blitz/5_graph_classification.py

@@ -0,0 +1,223 @@
+"""
+Training a GNN for Graph Classification
+=======================================
+
+By the end of this tutorial, you will be able to
+
+-  Load a DGL-provided graph classification dataset.
+-  Understand what *readout* function does.
+-  Understand how to create and use a minibatch of graphs.
+-  Build a GNN-based graph classification model.
+-  Train and evaluate the model on a DGL-provided dataset.
+
+(Time estimate: 18 minutes)
+"""
+
+import os
+
+os.environ["DGLBACKEND"] = "pytorch"
+import dgl
+import dgl.data
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+######################################################################
+# Overview of Graph Classification with GNN
+# -----------------------------------------
+#
+# Graph classification or regression requires a model to predict certain
+# graph-level properties of a single graph given its node and edge
+# features.  Molecular property prediction is one particular application.
+#
+# This tutorial shows how to train a graph classification model for a
+# small dataset from the paper `How Powerful Are Graph Neural
+# Networks <https://arxiv.org/abs/1810.00826>`__.
+#
+# Loading Data
+# ------------
+#
+
+
+# Generate a synthetic dataset with 10000 graphs, ranging from 10 to 500 nodes.
+dataset = dgl.data.GINDataset("PROTEINS", self_loop=True)
+
+
+######################################################################
+# The dataset is a set of graphs, each with node features and a single
+# label. One can see the node feature dimensionality and the number of
+# possible graph categories of ``GINDataset`` objects in ``dim_nfeats``
+# and ``gclasses`` attributes.
+#
+
+print("Node feature dimensionality:", dataset.dim_nfeats)
+print("Number of graph categories:", dataset.gclasses)
+
+
+from dgl.dataloading import GraphDataLoader
+
+######################################################################
+# Defining Data Loader
+# --------------------
+#
+# A graph classification dataset usually contains two types of elements: a
+# set of graphs, and their graph-level labels. Similar to an image
+# classification task, when the dataset is large enough, we need to train
+# with mini-batches. When you train a model for image classification or
+# language modeling, you will use a ``DataLoader`` to iterate over the
+# dataset. In DGL, you can use the ``GraphDataLoader``.
+#
+# You can also use various dataset samplers provided in
+# `torch.utils.data.sampler <https://pytorch.org/docs/stable/data.html#data-loading-order-and-sampler>`__.
+# For example, this tutorial creates a training ``GraphDataLoader`` and
+# test ``GraphDataLoader``, using ``SubsetRandomSampler`` to tell PyTorch
+# to sample from only a subset of the dataset.
+#
+
+from torch.utils.data.sampler import SubsetRandomSampler
+
+num_examples = len(dataset)
+num_train = int(num_examples * 0.8)
+
+train_sampler = SubsetRandomSampler(torch.arange(num_train))
+test_sampler = SubsetRandomSampler(torch.arange(num_train, num_examples))
+
+train_dataloader = GraphDataLoader(
+    dataset, sampler=train_sampler, batch_size=5, drop_last=False
+)
+test_dataloader = GraphDataLoader(
+    dataset, sampler=test_sampler, batch_size=5, drop_last=False
+)
+
+
+######################################################################
+# You can try to iterate over the created ``GraphDataLoader`` and see what it
+# gives:
+#
+
+it = iter(train_dataloader)
+batch = next(it)
+print(batch)
+
+
+######################################################################
+# As each element in ``dataset`` has a graph and a label, the
+# ``GraphDataLoader`` will return two objects for each iteration. The
+# first element is the batched graph, and the second element is simply a
+# label vector representing the category of each graph in the mini-batch.
+# Next, we’ll talked about the batched graph.
+#
+# A Batched Graph in DGL
+# ----------------------
+#
+# In each mini-batch, the sampled graphs are combined into a single bigger
+# batched graph via ``dgl.batch``. The single bigger batched graph merges
+# all original graphs as separately connected components, with the node
+# and edge features concatenated. This bigger graph is also a ``DGLGraph``
+# instance (so you can
+# still treat it as a normal ``DGLGraph`` object as in
+# `here <2_dglgraph.ipynb>`__). It however contains the information
+# necessary for recovering the original graphs, such as the number of
+# nodes and edges of each graph element.
+#
+
+batched_graph, labels = batch
+print(
+    "Number of nodes for each graph element in the batch:",
+    batched_graph.batch_num_nodes(),
+)
+print(
+    "Number of edges for each graph element in the batch:",
+    batched_graph.batch_num_edges(),
+)
+
+# Recover the original graph elements from the minibatch
+graphs = dgl.unbatch(batched_graph)
+print("The original graphs in the minibatch:")
+print(graphs)
+
+
+######################################################################
+# Define Model
+# ------------
+#
+# This tutorial will build a two-layer `Graph Convolutional Network
+# (GCN) <http://tkipf.github.io/graph-convolutional-networks/>`__. Each of
+# its layer computes new node representations by aggregating neighbor
+# information. If you have gone through the
+# :doc:`introduction <1_introduction>`, you will notice two
+# differences:
+#
+# -  Since the task is to predict a single category for the *entire graph*
+#    instead of for every node, you will need to aggregate the
+#    representations of all the nodes and potentially the edges to form a
+#    graph-level representation. Such process is more commonly referred as
+#    a *readout*. A simple choice is to average the node features of a
+#    graph with ``dgl.mean_nodes()``.
+#
+# -  The input graph to the model will be a batched graph yielded by the
+#    ``GraphDataLoader``. The readout functions provided by DGL can handle
+#    batched graphs so that they will return one representation for each
+#    minibatch element.
+#
+
+from dgl.nn import GraphConv
+
+
+class GCN(nn.Module):
+    def __init__(self, in_feats, h_feats, num_classes):
+        super(GCN, self).__init__()
+        self.conv1 = GraphConv(in_feats, h_feats)
+        self.conv2 = GraphConv(h_feats, num_classes)
+
+    def forward(self, g, in_feat):
+        h = self.conv1(g, in_feat)
+        h = F.relu(h)
+        h = self.conv2(g, h)
+        g.ndata["h"] = h
+        return dgl.mean_nodes(g, "h")
+
+
+######################################################################
+# Training Loop
+# -------------
+#
+# The training loop iterates over the training set with the
+# ``GraphDataLoader`` object and computes the gradients, just like
+# image classification or language modeling.
+#
+
+# Create the model with given dimensions
+model = GCN(dataset.dim_nfeats, 16, dataset.gclasses)
+optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+
+for epoch in range(20):
+    for batched_graph, labels in train_dataloader:
+        pred = model(batched_graph, batched_graph.ndata["attr"].float())
+        loss = F.cross_entropy(pred, labels)
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+num_correct = 0
+num_tests = 0
+for batched_graph, labels in test_dataloader:
+    pred = model(batched_graph, batched_graph.ndata["attr"].float())
+    num_correct += (pred.argmax(1) == labels).sum().item()
+    num_tests += len(labels)
+
+print("Test accuracy:", num_correct / num_tests)
+
+
+######################################################################
+# What’s next
+# -----------
+#
+# -  See `GIN
+#    example <https://github.com/dmlc/dgl/tree/master/examples/pytorch/gin>`__
+#    for an end-to-end graph classification model.
+#
+
+
+# Thumbnail credits: DGL
+# sphinx_gallery_thumbnail_path = '_static/blitz_5_graph_classification.png'

tutorials/blitz/6_load_data.py

@@ -0,0 +1,255 @@
+"""
+Make Your Own Dataset
+=====================
+
+This tutorial assumes that you already know :doc:`the basics of training a
+GNN for node classification <1_introduction>` and :doc:`how to
+create, load, and store a DGL graph <2_dglgraph>`.
+
+By the end of this tutorial, you will be able to
+
+-  Create your own graph dataset for node classification, link
+   prediction, or graph classification.
+
+(Time estimate: 15 minutes)
+"""
+
+
+######################################################################
+# ``DGLDataset`` Object Overview
+# ------------------------------
+#
+# Your custom graph dataset should inherit the ``dgl.data.DGLDataset``
+# class and implement the following methods:
+#
+# -  ``__getitem__(self, i)``: retrieve the ``i``-th example of the
+#    dataset. An example often contains a single DGL graph, and
+#    occasionally its label.
+# -  ``__len__(self)``: the number of examples in the dataset.
+# -  ``process(self)``: load and process raw data from disk.
+#
+
+
+######################################################################
+# Creating a Dataset for Node Classification or Link Prediction from CSV
+# ----------------------------------------------------------------------
+#
+# A node classification dataset often consists of a single graph, as well
+# as its node and edge features.
+#
+# This tutorial takes a small dataset based on `Zachary’s Karate Club
+# network <https://en.wikipedia.org/wiki/Zachary%27s_karate_club>`__. It
+# contains
+#
+# * A ``members.csv`` file containing the attributes of all
+#   members, as well as their attributes.
+#
+# * An ``interactions.csv`` file
+#   containing the pair-wise interactions between two club members.
+#
+
+import urllib.request
+
+import pandas as pd
+
+urllib.request.urlretrieve(
+    "https://data.dgl.ai/tutorial/dataset/members.csv", "./members.csv"
+)
+urllib.request.urlretrieve(
+    "https://data.dgl.ai/tutorial/dataset/interactions.csv",
+    "./interactions.csv",
+)
+
+members = pd.read_csv("./members.csv")
+members.head()
+
+interactions = pd.read_csv("./interactions.csv")
+interactions.head()
+
+
+######################################################################
+# This tutorial treats the members as nodes and interactions as edges. It
+# takes age as a numeric feature of the nodes, affiliated club as the label
+# of the nodes, and edge weight as a numeric feature of the edges.
+#
+# .. note::
+#
+#    The original Zachary’s Karate Club network does not have
+#    member ages. The ages in this tutorial are generated synthetically
+#    for demonstrating how to add node features into the graph for dataset
+#    creation.
+#
+# .. note::
+#
+#    In practice, taking age directly as a numeric feature may
+#    not work well in machine learning; strategies like binning or
+#    normalizing the feature would work better. This tutorial directly
+#    takes the values as-is for simplicity.
+#
+
+import os
+
+os.environ["DGLBACKEND"] = "pytorch"
+import dgl
+import torch
+from dgl.data import DGLDataset
+
+
+class KarateClubDataset(DGLDataset):
+    def __init__(self):
+        super().__init__(name="karate_club")
+
+    def process(self):
+        nodes_data = pd.read_csv("./members.csv")
+        edges_data = pd.read_csv("./interactions.csv")
+        node_features = torch.from_numpy(nodes_data["Age"].to_numpy())
+        node_labels = torch.from_numpy(
+            nodes_data["Club"].astype("category").cat.codes.to_numpy()
+        )
+        edge_features = torch.from_numpy(edges_data["Weight"].to_numpy())
+        edges_src = torch.from_numpy(edges_data["Src"].to_numpy())
+        edges_dst = torch.from_numpy(edges_data["Dst"].to_numpy())
+
+        self.graph = dgl.graph(
+            (edges_src, edges_dst), num_nodes=nodes_data.shape[0]
+        )
+        self.graph.ndata["feat"] = node_features
+        self.graph.ndata["label"] = node_labels
+        self.graph.edata["weight"] = edge_features
+
+        # If your dataset is a node classification dataset, you will need to assign
+        # masks indicating whether a node belongs to training, validation, and test set.
+        n_nodes = nodes_data.shape[0]
+        n_train = int(n_nodes * 0.6)
+        n_val = int(n_nodes * 0.2)
+        train_mask = torch.zeros(n_nodes, dtype=torch.bool)
+        val_mask = torch.zeros(n_nodes, dtype=torch.bool)
+        test_mask = torch.zeros(n_nodes, dtype=torch.bool)
+        train_mask[:n_train] = True
+        val_mask[n_train : n_train + n_val] = True
+        test_mask[n_train + n_val :] = True
+        self.graph.ndata["train_mask"] = train_mask
+        self.graph.ndata["val_mask"] = val_mask
+        self.graph.ndata["test_mask"] = test_mask
+
+    def __getitem__(self, i):
+        return self.graph
+
+    def __len__(self):
+        return 1
+
+
+dataset = KarateClubDataset()
+graph = dataset[0]
+
+print(graph)
+
+
+######################################################################
+# Since a link prediction dataset only involves a single graph, preparing
+# a link prediction dataset will have the same experience as preparing a
+# node classification dataset.
+#
+
+
+######################################################################
+# Creating a Dataset for Graph Classification from CSV
+# ----------------------------------------------------
+#
+# Creating a graph classification dataset involves implementing
+# ``__getitem__`` to return both the graph and its graph-level label.
+#
+# This tutorial demonstrates how to create a graph classification dataset
+# with the following synthetic CSV data:
+#
+# -  ``graph_edges.csv``: containing three columns:
+#
+#    -  ``graph_id``: the ID of the graph.
+#    -  ``src``: the source node of an edge of the given graph.
+#    -  ``dst``: the destination node of an edge of the given graph.
+#
+# -  ``graph_properties.csv``: containing three columns:
+#
+#    -  ``graph_id``: the ID of the graph.
+#    -  ``label``: the label of the graph.
+#    -  ``num_nodes``: the number of nodes in the graph.
+#
+
+urllib.request.urlretrieve(
+    "https://data.dgl.ai/tutorial/dataset/graph_edges.csv", "./graph_edges.csv"
+)
+urllib.request.urlretrieve(
+    "https://data.dgl.ai/tutorial/dataset/graph_properties.csv",
+    "./graph_properties.csv",
+)
+edges = pd.read_csv("./graph_edges.csv")
+properties = pd.read_csv("./graph_properties.csv")
+
+edges.head()
+
+properties.head()
+
+
+class SyntheticDataset(DGLDataset):
+    def __init__(self):
+        super().__init__(name="synthetic")
+
+    def process(self):
+        edges = pd.read_csv("./graph_edges.csv")
+        properties = pd.read_csv("./graph_properties.csv")
+        self.graphs = []
+        self.labels = []
+
+        # Create a graph for each graph ID from the edges table.
+        # First process the properties table into two dictionaries with graph IDs as keys.
+        # The label and number of nodes are values.
+        label_dict = {}
+        num_nodes_dict = {}
+        for _, row in properties.iterrows():
+            label_dict[row["graph_id"]] = row["label"]
+            num_nodes_dict[row["graph_id"]] = row["num_nodes"]
+
+        # For the edges, first group the table by graph IDs.
+        edges_group = edges.groupby("graph_id")
+
+        # For each graph ID...
+        for graph_id in edges_group.groups:
+            # Find the edges as well as the number of nodes and its label.
+            edges_of_id = edges_group.get_group(graph_id)
+            src = edges_of_id["src"].to_numpy()
+            dst = edges_of_id["dst"].to_numpy()
+            num_nodes = num_nodes_dict[graph_id]
+            label = label_dict[graph_id]
+
+            # Create a graph and add it to the list of graphs and labels.
+            g = dgl.graph((src, dst), num_nodes=num_nodes)
+            self.graphs.append(g)
+            self.labels.append(label)
+
+        # Convert the label list to tensor for saving.
+        self.labels = torch.LongTensor(self.labels)
+
+    def __getitem__(self, i):
+        return self.graphs[i], self.labels[i]
+
+    def __len__(self):
+        return len(self.graphs)
+
+
+dataset = SyntheticDataset()
+graph, label = dataset[0]
+print(graph, label)
+
+######################################################################
+# Creating Dataset from CSV via :class:`~dgl.data.CSVDataset`
+# ------------------------------------------------------------
+#
+# The previous examples describe how to create a dataset from CSV files
+# step-by-step. DGL also provides a utility class :class:`~dgl.data.CSVDataset`
+# for reading and parsing data from CSV files. See :ref:`guide-data-pipeline-loadcsv`
+# for more details.
+#
+
+
+# Thumbnail credits: (Un)common Use Cases for Graph Databases, Michal Bachman
+# sphinx_gallery_thumbnail_path = '_static/blitz_6_load_data.png'

tutorials/blitz/README.txt

@@ -0,0 +1,2 @@
+A Blitz Introduction to DGL
+===========================

tutorials/cpu/README.txt

@@ -0,0 +1,2 @@
+Training on CPUs
+=========================

tutorials/cpu/argo_tutorial.py

@@ -0,0 +1,151 @@
+"""
+Improve Scalability on Multi-Core CPUs
+=====================================================
+
+Graph Neural Network (GNN) training suffers from low scalability on multi-core CPUs. 
+Specificially, the performance often caps at 16 cores, and no improvement is observed when applying more than 16 cores [#f1]_.
+ARGO is a runtime system that offers scalable performance.  
+With ARGO enabled, we are able to scale over 64 cores, allowing ARGO to speedup GNN training (in terms of epoch time) by up to 4.30x and 3.32x on a Xeon 8380H and a Xeon 6430L, respectively [#f2]_.
+This chapter focus on how to setup ARGO to unleash the potential of multi-core CPUs to speedup GNN training.
+
+Installation
+`````````````````````````````
+
+ARGO utilizes the scikit-optimize library for auto-tuning. Please install scikit-optimize to run ARGO:
+.. code-block:: shell
+    conda install -c conda-forge "scikit-optimize>=0.9.0" 
+or
+.. code-block:: shell
+    pip install scikit-optimize>=0.9
+
+Enabling ARGO on your own GNN program
+```````````````````````````````````````````
+In this section, we provide a step-by-step tutorial on how to enable ARGO on a DGL program. 
+We use the *ogb_example.py* [#f3]_ as an example.
+.. note::
+    We also provide the complete example file *ogb_example_ARGO.py* [#f4]_ 
+    which followed the steps below to enable ARGO on *ogb_example.py*.
+
+Step 1
+---------------------------
+First, include all necessary packages on top of the file. Please place your file and *argo.py* [#f5]_ in the same directory.
+
+.. code-block:: python
+    import os
+    import torch.distributed as dist
+    from torch.nn.parallel import DistributedDataParallel
+    import torch.multiprocessing as mp
+    from argo import ARGO
+
+Step 2
+---------------------------
+Setup PyTorch Distributed Data Parallel (DDP)
+
+2.1. Add the initialization function on top of the training program, and wrap the ```model``` with the DDP wrapper
+.. code-block:: python
+    def train(...):
+        dist.init_process_group('gloo', rank=rank, world_size=world_size) # newly added
+        model = SAGE(...) # original code
+        model = DistributedDataParallel(model) # newly added
+        ...
+     
+2.2. In the main program, add the following before launching the training function
+.. code-block:: python
+    ...
+    os.environ['MASTER_ADDR'] = '127.0.0.1'
+    os.environ['MASTER_PORT'] = '29501'
+    mp.set_start_method('fork', force=True)
+    train(args, device, data) # original code for launching the training function
+
+Step 3
+---------------------------
+Enable ARGO by initializing the runtime system, and wrapping the training function
+.. code-block:: python
+    runtime = ARGO(n_search = 15, epoch = args.num_epochs, batch_size = args.batch_size) # initialization
+    runtime.run(train, args=(args, device, data)) # wrap the training function
+
+.. note::
+    ARGO takes three input parameters: number of searches *n_search*, number of epochs, and the mini-batch size. 
+    Increasing *n_search* potentially leads to a better configuration with less epoch time; 
+    however, searching itself also causes extra overhead. We recommend setting *n_search* from 15 to 45 for an optimal overall performance. 
+
+Step 4
+---------------------------
+Modify the input of the training function, by directly adding ARGO parameters after the original inputs.
+   
+This is the original function:
+.. code-block:: python
+   def train(args, device, data):
+   
+Add the following variables: *rank, world_size, comp_core, load_core, counter, b_size, ep*
+.. code-block:: python
+    def train(args, device, data, rank, world_size, comp_core, load_core, counter, b_size, ep):
+
+Step 5
+---------------------------
+Modify the *dataloader* function in the training function
+.. code-block:: python
+    dataloader = dgl.dataloading.DataLoader(
+            g,
+            train_nid,
+            sampler,
+            batch_size=b_size, # modified
+            shuffle=True,
+            drop_last=False,
+            num_workers=len(load_core), # modified
+            use_ddp = True) # newly added
+
+Step 6
+---------------------------
+Enable core-binding by adding *enable_cpu_affinity()* before the training for-loop, and also change the number of epochs into the variable *ep*: 
+.. code-block:: python
+    with dataloader.enable_cpu_affinity(loader_cores=load_core, compute_cores=comp_core): 
+        for epoch in range(ep): # change num_epochs to ep
+   
+Step 7
+---------------------------
+Last step! Load the model before training and save it afterward.  
+
+Original Program:
+.. code-block:: python
+    with dataloader.enable_cpu_affinity(loader_cores=load_core, compute_cores=comp_core): 
+        for epoch in range(ep): 
+        ... # training operations
+   
+Modified:
+.. code-block:: python
+    PATH = "model.pt"
+    if counter[0] != 0:
+        checkpoint = th.load(PATH)
+        model.load_state_dict(checkpoint['model_state_dict'])
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        epoch = checkpoint['epoch']
+        loss = checkpoint['loss']
+   
+    with dataloader.enable_cpu_affinity(loader_cores=load_core, compute_cores=comp_core): 
+        for epoch in range(ep): 
+        ... # training operations
+   
+    dist.barrier()
+    if rank == 0:
+        th.save({'epoch': counter[0],
+                    'model_state_dict': model.state_dict(),
+                    'optimizer_state_dict': optimizer.state_dict(),
+                    'loss': loss,
+                    }, PATH)
+   
+Step 8
+---------------------------
+Done! You can now run your GNN program with ARGO enabled.
+.. code-block:: shell
+    python <your_code>.py
+
+    
+.. rubric:: Footnotes
+
+.. [#f1] https://github.com/dmlc/dgl/blob/master/examples/pytorch/argo/argo_scale.png
+.. [#f2] https://arxiv.org/abs/2402.03671
+.. [#f3] https://github.com/dmlc/dgl/blob/master/examples/pytorch/argo/ogb_example.py
+.. [#f4] https://github.com/dmlc/dgl/blob/master/examples/pytorch/argo/ogb_example_ARGO.py
+.. [#f5] https://github.com/dmlc/dgl/blob/master/examples/pytorch/argo/argo.py
+"""

tutorials/cpu/cpu_best_practises.py

@@ -0,0 +1,88 @@
+"""
+CPU Best Practices
+=====================================================
+
+This chapter focus on providing best practises for environment setup
+to get the best performance during training and inference on the CPU.
+
+Intel
+`````````````````````````````
+
+Hyper-threading
+---------------------------
+
+For specific workloads as GNN’s domain, suggested default setting for having best performance
+is to turn off hyperthreading.
+Turning off the hyper threading feature can be done at BIOS [#f1]_ or operating system level [#f2]_ [#f3]_ .
+
+Alternative memory allocators
+---------------------------
+
+Alternative memory allocators, such as *tcmalloc*, might provide significant performance improvements by more efficient memory usage, reducing overhead on unnecessary memory allocations or deallocations. *tcmalloc* uses thread-local caches to reduce overhead on thread synchronization, locks contention by using spinlocks and per-thread arenas respectively and categorizes memory allocations by sizes to reduce overhead on memory fragmentation.
+
+To take advantage of optimizations *tcmalloc* provides, install it on your system (on Ubuntu *tcmalloc* is included in libgoogle-perftools4 package) and add shared library to the LD_PRELOAD environment variable:
+
+.. code-block:: shell
+
+  export LD_PRELOAD=/lib/x86_64-linux-gnu/libtcmalloc.so.4:$LD_PRELOAD
+
+OpenMP settings
+---------------------------
+
+As `OpenMP` is the default parallel backend, we could control performance
+including sampling and training via `dgl.utils.set_num_threads()`.
+
+If number of OpenMP threads is not set and `num_workers` in dataloader is set
+to 0, the OpenMP runtime typically use the number of available CPU cores by
+default. This works well for most cases, and is also the default behavior in DGL.
+
+If `num_workers` in dataloader is set to greater than 0, the number of
+OpenMP threads will be set to **1** for each worker process. This is the
+default behavior in PyTorch. In this case, we can set the number of OpenMP
+threads to the number of CPU cores in the main process.
+
+Performance tuning is highly dependent on the workload and hardware
+configuration. We recommend users to try different settings and choose the
+best one for their own cases.
+
+**Dataloader CPU affinity**
+
+.. note::
+
+    This feature is available for `dgl.dataloading.DataLoader` only. Not
+    available for dataloaders in `dgl.graphbolt` yet.
+
+
+If number of dataloader workers is more than 0, please consider using **use_cpu_affinity()** method
+of DGL Dataloader class, it will generally result in significant performance improvement for training.
+
+*use_cpu_affinity* will set the proper OpenMP thread count (equal to the number of CPU cores allocated for main process),
+affinitize dataloader workers for separate CPU cores and restrict the main process to remaining cores
+
+In multiple NUMA nodes setups *use_cpu_affinity* will only use cores of NUMA node 0 by default
+with an assumption, that the workload is scaling poorly across multiple NUMA nodes. If you believe
+your workload will have better performance utilizing more than one NUMA node, you can pass
+the list of cores to use for dataloading (loader_cores) and for compute (compute_cores).
+
+loader_cores and compute_cores arguments (list of CPU cores) can be passed to *enable_cpu_affinity* for more
+control over which cores should be used, e.g. in case a workload scales well across multiple NUMA nodes.
+
+Usage:
+    .. code:: python
+
+        dataloader = dgl.dataloading.DataLoader(...)
+        ...
+        with dataloader.enable_cpu_affinity():
+            <training loop or inferencing>
+
+**Manual control**
+
+For advanced and more fine-grained control over OpenMP settings please refer to Maximize Performance of Intel® Optimization for PyTorch* on CPU [#f4]_ article
+
+.. rubric:: Footnotes
+
+.. [#f1] https://www.intel.com/content/www/us/en/support/articles/000007645/boards-and-kits/desktop-boards.html
+.. [#f2] https://aws.amazon.com/blogs/compute/disabling-intel-hyper-threading-technology-on-amazon-linux/
+.. [#f3] https://aws.amazon.com/blogs/compute/disabling-intel-hyper-threading-technology-on-amazon-ec2-windows-instances/
+.. [#f4] https://software.intel.com/content/www/us/en/develop/articles/how-to-get-better-performance-on-pytorchcaffe2-with-intel-acceleration.html
+"""

tutorials/models/1_gnn/1_gcn.py

@@ -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>`_.

tutorials/models/1_gnn/4_rgcn.py

@@ -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>`_.

tutorials/models/1_gnn/6_line_graph.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>`_.

tutorials/models/1_gnn/9_gat.py

@@ -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

tutorials/models/1_gnn/README.txt

@@ -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.

tutorials/models/2_small_graph/3_tree-lstm.py

@@ -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.

tutorials/models/2_small_graph/README.txt

@@ -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.

tutorials/models/3_generative_model/5_dgmg.py

@@ -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>`__.
+#

tutorials/models/3_generative_model/README.txt

@@ -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.

tutorials/models/4_old_wines/2_capsule.py

@@ -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

tutorials/models/4_old_wines/7_transformer.py

@@ -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.

tutorials/models/4_old_wines/README.txt

@@ -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.

tutorials/models/README.txt

@@ -0,0 +1,2 @@
+Paper Study with DGL
+=========================================

tutorials/multi/1_graph_classification.py

@@ -0,0 +1,270 @@
+"""
+Single Machine Multi-GPU Minibatch Graph Classification
+=======================================================
+
+In this tutorial, you will learn how to use multiple GPUs in training a
+graph neural network (GNN) for graph classification. This tutorial assumes
+knowledge in GNNs for graph classification and we recommend you to check
+:doc:`Training a GNN for Graph Classification <../blitz/5_graph_classification>` otherwise.
+
+(Time estimate: 8 minutes)
+
+To use a single GPU in training a GNN, we need to put the model, graph(s), and other
+tensors (e.g. labels) on the same GPU:
+
+.. code:: python
+
+    import torch
+
+    # Use the first GPU
+    device = torch.device("cuda:0")
+    model = model.to(device)
+    graph = graph.to(device)
+    labels = labels.to(device)
+
+The node and edge features in the graphs, if any, will also be on the GPU.
+After that, the forward computation, backward computation and parameter
+update will take place on the GPU. For graph classification, this repeats
+for each minibatch gradient descent.
+
+Using multiple GPUs allows performing more computation per unit of time. It
+is like having a team work together, where each GPU is a team member. We need
+to distribute the computation workload across GPUs and let them synchronize
+the efforts regularly. PyTorch provides convenient APIs for this task with
+multiple processes, one per GPU, and we can use them in conjunction with DGL.
+
+Intuitively, we can distribute the workload along the dimension of data. This
+allows multiple GPUs to perform the forward and backward computation of
+multiple gradient descents in parallel. To distribute a dataset across
+multiple GPUs, we need to partition it into multiple mutually exclusive
+subsets of a similar size, one per GPU. We need to repeat the random
+partition every epoch to guarantee randomness. We can use
+:func:`~dgl.dataloading.pytorch.GraphDataLoader`, which wraps some PyTorch 
+APIs and does the job for graph classification in data loading.
+
+Once all GPUs have finished the backward computation for its minibatch,
+we need to synchronize the model parameter update across them. Specifically,
+this involves collecting gradients from all GPUs, averaging them and updating
+the model parameters on each GPU. We can wrap a PyTorch model with
+:func:`~torch.nn.parallel.DistributedDataParallel` so that the model
+parameter update will invoke gradient synchronization first under the hood.
+
+.. image:: https://data.dgl.ai/tutorial/mgpu_gc.png
+  :width: 450px
+  :align: center
+
+That’s the core behind this tutorial. We will explore it more in detail with
+a complete example below.
+
+.. note::
+
+   See `this tutorial <https://pytorch.org/tutorials/intermediate/ddp_tutorial.html>`__
+   from PyTorch for general multi-GPU training with ``DistributedDataParallel``.
+
+Distributed Process Group Initialization
+----------------------------------------
+
+For communication between multiple processes in multi-gpu training, we need
+to start the distributed backend at the beginning of each process. We use
+`world_size` to refer to the number of processes and `rank` to refer to the
+process ID, which should be an integer from `0` to `world_size - 1`.
+"""
+
+import os
+
+os.environ["DGLBACKEND"] = "pytorch"
+import torch.distributed as dist
+
+
+def init_process_group(world_size, rank):
+    dist.init_process_group(
+        backend="gloo",  # change to 'nccl' for multiple GPUs
+        init_method="tcp://127.0.0.1:12345",
+        world_size=world_size,
+        rank=rank,
+    )
+
+
+###############################################################################
+# Data Loader Preparation
+# -----------------------
+#
+# We split the dataset into training, validation and test subsets. In dataset
+# splitting, we need to use a same random seed across processes to ensure a
+# same split. We follow the common practice to train with multiple GPUs and
+# evaluate with a single GPU, thus only set `use_ddp` to True in the
+# :func:`~dgl.dataloading.pytorch.GraphDataLoader` for the training set, where
+# `ddp` stands for :func:`~torch.nn.parallel.DistributedDataParallel`.
+#
+
+from dgl.data import split_dataset
+from dgl.dataloading import GraphDataLoader
+
+
+def get_dataloaders(dataset, seed, batch_size=32):
+    # Use a 80:10:10 train-val-test split
+    train_set, val_set, test_set = split_dataset(
+        dataset, frac_list=[0.8, 0.1, 0.1], shuffle=True, random_state=seed
+    )
+    train_loader = GraphDataLoader(
+        train_set, use_ddp=True, batch_size=batch_size, shuffle=True
+    )
+    val_loader = GraphDataLoader(val_set, batch_size=batch_size)
+    test_loader = GraphDataLoader(test_set, batch_size=batch_size)
+
+    return train_loader, val_loader, test_loader
+
+
+###############################################################################
+# Model Initialization
+# --------------------
+#
+# For this tutorial, we use a simplified Graph Isomorphism Network (GIN).
+#
+
+import torch.nn as nn
+import torch.nn.functional as F
+
+from dgl.nn.pytorch import GINConv, SumPooling
+
+
+class GIN(nn.Module):
+    def __init__(self, input_size=1, num_classes=2):
+        super(GIN, self).__init__()
+
+        self.conv1 = GINConv(
+            nn.Linear(input_size, num_classes), aggregator_type="sum"
+        )
+        self.conv2 = GINConv(
+            nn.Linear(num_classes, num_classes), aggregator_type="sum"
+        )
+        self.pool = SumPooling()
+
+    def forward(self, g, feats):
+        feats = self.conv1(g, feats)
+        feats = F.relu(feats)
+        feats = self.conv2(g, feats)
+
+        return self.pool(g, feats)
+
+
+###############################################################################
+# To ensure same initial model parameters across processes, we need to set the
+# same random seed before model initialization. Once we construct a model
+# instance, we wrap it with :func:`~torch.nn.parallel.DistributedDataParallel`.
+#
+
+import torch
+from torch.nn.parallel import DistributedDataParallel
+
+
+def init_model(seed, device):
+    torch.manual_seed(seed)
+    model = GIN().to(device)
+    if device.type == "cpu":
+        model = DistributedDataParallel(model)
+    else:
+        model = DistributedDataParallel(
+            model, device_ids=[device], output_device=device
+        )
+
+    return model
+
+
+###############################################################################
+# Main Function for Each Process
+# -----------------------------
+#
+# Define the model evaluation function as in the single-GPU setting.
+#
+
+
+def evaluate(model, dataloader, device):
+    model.eval()
+
+    total = 0
+    total_correct = 0
+
+    for bg, labels in dataloader:
+        bg = bg.to(device)
+        labels = labels.to(device)
+        # Get input node features
+        feats = bg.ndata.pop("attr")
+        with torch.no_grad():
+            pred = model(bg, feats)
+        _, pred = torch.max(pred, 1)
+        total += len(labels)
+        total_correct += (pred == labels).sum().cpu().item()
+
+    return 1.0 * total_correct / total
+
+
+###############################################################################
+# Define the run function for each process.
+#
+
+from torch.optim import Adam
+
+
+def run(rank, world_size, dataset, seed=0):
+    init_process_group(world_size, rank)
+    if torch.cuda.is_available():
+        device = torch.device("cuda:{:d}".format(rank))
+        torch.cuda.set_device(device)
+    else:
+        device = torch.device("cpu")
+
+    model = init_model(seed, device)
+    criterion = nn.CrossEntropyLoss()
+    optimizer = Adam(model.parameters(), lr=0.01)
+
+    train_loader, val_loader, test_loader = get_dataloaders(dataset, seed)
+    for epoch in range(5):
+        model.train()
+        # The line below ensures all processes use a different
+        # random ordering in data loading for each epoch.
+        train_loader.set_epoch(epoch)
+
+        total_loss = 0
+        for bg, labels in train_loader:
+            bg = bg.to(device)
+            labels = labels.to(device)
+            feats = bg.ndata.pop("attr")
+            pred = model(bg, feats)
+
+            loss = criterion(pred, labels)
+            total_loss += loss.cpu().item()
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+        loss = total_loss
+        print("Loss: {:.4f}".format(loss))
+
+        val_acc = evaluate(model, val_loader, device)
+        print("Val acc: {:.4f}".format(val_acc))
+
+    test_acc = evaluate(model, test_loader, device)
+    print("Test acc: {:.4f}".format(test_acc))
+    dist.destroy_process_group()
+
+
+###############################################################################
+# Finally we load the dataset and launch the processes.
+#
+
+import torch.multiprocessing as mp
+from dgl.data import GINDataset
+
+
+def main():
+    if not torch.cuda.is_available():
+        print("No GPU found!")
+        return
+
+    num_gpus = torch.cuda.device_count()
+    dataset = GINDataset(name="IMDBBINARY", self_loop=False)
+    mp.spawn(run, args=(num_gpus, dataset), nprocs=num_gpus)
+
+
+if __name__ == "__main__":
+    main()

tutorials/multi/2_node_classification.py

@@ -0,0 +1,371 @@
+"""
+Single Machine Multi-GPU Minibatch Node Classification
+======================================================
+
+In this tutorial, you will learn how to use multiple GPUs in training a
+graph neural network (GNN) for node classification.
+
+This tutorial assumes that you have read the `Stochastic GNN Training for Node
+Classification in DGL <../../notebooks/stochastic_training/node_classification.ipynb>`__.
+It also assumes that you know the basics of training general
+models with multi-GPU with ``DistributedDataParallel``.
+
+.. note::
+
+   See `this tutorial <https://pytorch.org/tutorials/intermediate/ddp_tutorial.html>`__
+   from PyTorch for general multi-GPU training with ``DistributedDataParallel``.  Also,
+   see the first section of :doc:`the multi-GPU graph classification
+   tutorial <1_graph_classification>`
+   for an overview of using ``DistributedDataParallel`` with DGL.
+
+"""
+
+######################################################################
+# Importing Packages
+# ---------------
+#
+# We use ``torch.distributed`` to initialize a distributed training context
+# and ``torch.multiprocessing`` to spawn multiple processes for each GPU.
+#
+
+import os
+
+os.environ["DGLBACKEND"] = "pytorch"
+import time
+
+import dgl.graphbolt as gb
+import dgl.nn as dglnn
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+import torch.nn as nn
+import torch.nn.functional as F
+import torchmetrics.functional as MF
+from torch.distributed.algorithms.join import Join
+from torch.nn.parallel import DistributedDataParallel as DDP
+from tqdm.auto import tqdm
+
+
+######################################################################
+# Defining Model
+# --------------
+#
+# The model will be again identical to `Stochastic GNN Training for Node
+# Classification in DGL <../../notebooks/stochastic_training/node_classification.ipynb>`__.
+#
+
+
+class SAGE(nn.Module):
+    def __init__(self, in_size, hidden_size, out_size):
+        super().__init__()
+        self.layers = nn.ModuleList()
+        # Three-layer GraphSAGE-mean.
+        self.layers.append(dglnn.SAGEConv(in_size, hidden_size, "mean"))
+        self.layers.append(dglnn.SAGEConv(hidden_size, hidden_size, "mean"))
+        self.layers.append(dglnn.SAGEConv(hidden_size, out_size, "mean"))
+        self.dropout = nn.Dropout(0.5)
+        self.hidden_size = hidden_size
+        self.out_size = out_size
+        # Set the dtype for the layers manually.
+        self.float()
+
+    def forward(self, blocks, x):
+        hidden_x = x
+        for layer_idx, (layer, block) in enumerate(zip(self.layers, blocks)):
+            hidden_x = layer(block, hidden_x)
+            is_last_layer = layer_idx == len(self.layers) - 1
+            if not is_last_layer:
+                hidden_x = F.relu(hidden_x)
+                hidden_x = self.dropout(hidden_x)
+        return hidden_x
+
+
+######################################################################
+# Mini-batch Data Loading
+# -----------------------
+#
+# The major difference from the previous tutorial is that we will use
+# ``DistributedItemSampler`` instead of ``ItemSampler`` to sample mini-batches
+# of nodes.  ``DistributedItemSampler`` is a distributed version of
+# ``ItemSampler`` that works with ``DistributedDataParallel``.  It is
+# implemented as a wrapper around ``ItemSampler`` and will sample the same
+# minibatch on all replicas.  It also supports dropping the last non-full
+# minibatch to avoid the need for padding.
+#
+
+
+def create_dataloader(
+    graph,
+    features,
+    itemset,
+    device,
+    is_train,
+):
+    datapipe = gb.DistributedItemSampler(
+        item_set=itemset,
+        batch_size=1024,
+        drop_last=is_train,
+        shuffle=is_train,
+        drop_uneven_inputs=is_train,
+    )
+    datapipe = datapipe.copy_to(device)
+    # Now that we have moved to device, sample_neighbor and fetch_feature steps
+    # will be executed on GPUs.
+    datapipe = datapipe.sample_neighbor(graph, [10, 10, 10])
+    datapipe = datapipe.fetch_feature(features, node_feature_keys=["feat"])
+    return gb.DataLoader(datapipe)
+
+
+def weighted_reduce(tensor, weight, dst=0):
+    ########################################################################
+    # (HIGHLIGHT) Collect accuracy and loss values from sub-processes and
+    # obtain overall average values.
+    #
+    # `torch.distributed.reduce` is used to reduce tensors from all the
+    # sub-processes to a specified process, ReduceOp.SUM is used by default.
+    #
+    # Because the GPUs may have differing numbers of processed items, we
+    # perform a weighted mean to calculate the exact loss and accuracy.
+    ########################################################################
+    dist.reduce(tensor=tensor, dst=dst)
+    weight = torch.tensor(weight, device=tensor.device)
+    dist.reduce(tensor=weight, dst=dst)
+    return tensor / weight
+
+
+######################################################################
+# Evaluation Loop
+# ---------------
+#
+# The evaluation loop is almost identical to the previous tutorial.
+#
+
+
+@torch.no_grad()
+def evaluate(rank, model, graph, features, itemset, num_classes, device):
+    model.eval()
+    y = []
+    y_hats = []
+    dataloader = create_dataloader(
+        graph,
+        features,
+        itemset,
+        device,
+        is_train=False,
+    )
+
+    for data in tqdm(dataloader) if rank == 0 else dataloader:
+        blocks = data.blocks
+        x = data.node_features["feat"]
+        y.append(data.labels)
+        y_hats.append(model.module(blocks, x))
+
+    res = MF.accuracy(
+        torch.cat(y_hats),
+        torch.cat(y),
+        task="multiclass",
+        num_classes=num_classes,
+    )
+
+    return res.to(device), sum(y_i.size(0) for y_i in y)
+
+
+######################################################################
+# Training Loop
+# -------------
+#
+# The training loop is also almost identical to the previous tutorial except
+# that we use Join Context Manager to solve the uneven input problem. The
+# mechanics of Distributed Data Parallel (DDP) training in PyTorch requires
+# the number of inputs are the same for all ranks, otherwise the program may
+# error or hang. To solve it, PyTorch provides Join Context Manager. Please
+# refer to `this tutorial <https://pytorch.org/tutorials/advanced/generic_join.html>`__
+# for detailed information.
+#
+
+
+def train(
+    rank,
+    graph,
+    features,
+    train_set,
+    valid_set,
+    num_classes,
+    model,
+    device,
+):
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+    # Create training data loader.
+    dataloader = create_dataloader(
+        graph,
+        features,
+        train_set,
+        device,
+        is_train=True,
+    )
+
+    for epoch in range(5):
+        epoch_start = time.time()
+
+        model.train()
+        total_loss = torch.tensor(0, dtype=torch.float, device=device)
+        num_train_items = 0
+        with Join([model]):
+            for data in tqdm(dataloader) if rank == 0 else dataloader:
+                # The input features are from the source nodes in the first
+                # layer's computation graph.
+                x = data.node_features["feat"]
+
+                # The ground truth labels are from the destination nodes
+                # in the last layer's computation graph.
+                y = data.labels
+
+                blocks = data.blocks
+
+                y_hat = model(blocks, x)
+
+                # Compute loss.
+                loss = F.cross_entropy(y_hat, y)
+
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+                total_loss += loss.detach() * y.size(0)
+                num_train_items += y.size(0)
+
+        # Evaluate the model.
+        if rank == 0:
+            print("Validating...")
+        acc, num_val_items = evaluate(
+            rank,
+            model,
+            graph,
+            features,
+            valid_set,
+            num_classes,
+            device,
+        )
+        total_loss = weighted_reduce(total_loss, num_train_items)
+        acc = weighted_reduce(acc * num_val_items, num_val_items)
+
+        # We synchronize before measuring the epoch time.
+        torch.cuda.synchronize()
+        epoch_end = time.time()
+        if rank == 0:
+            print(
+                f"Epoch {epoch:05d} | "
+                f"Average Loss {total_loss.item():.4f} | "
+                f"Accuracy {acc.item():.4f} | "
+                f"Time {epoch_end - epoch_start:.4f}"
+            )
+
+
+######################################################################
+# Defining Traning and Evaluation Procedures
+# ------------------------------------------
+#
+# The following code defines the main function for each process. It is
+# similar to the previous tutorial except that we need to initialize a
+# distributed training context with ``torch.distributed`` and wrap the model
+# with ``torch.nn.parallel.DistributedDataParallel``.
+#
+
+
+def run(rank, world_size, devices, dataset):
+    # Set up multiprocessing environment.
+    device = devices[rank]
+    torch.cuda.set_device(device)
+    dist.init_process_group(
+        backend="nccl",  # Use NCCL backend for distributed GPU training
+        init_method="tcp://127.0.0.1:12345",
+        world_size=world_size,
+        rank=rank,
+    )
+
+    # Pin the graph and features in-place to enable GPU access.
+    graph = dataset.graph.pin_memory_()
+    features = dataset.feature.pin_memory_()
+    train_set = dataset.tasks[0].train_set
+    valid_set = dataset.tasks[0].validation_set
+    num_classes = dataset.tasks[0].metadata["num_classes"]
+
+    in_size = features.size("node", None, "feat")[0]
+    hidden_size = 256
+    out_size = num_classes
+
+    # Create GraphSAGE model. It should be copied onto a GPU as a replica.
+    model = SAGE(in_size, hidden_size, out_size).to(device)
+    model = DDP(model)
+
+    # Model training.
+    if rank == 0:
+        print("Training...")
+    train(
+        rank,
+        graph,
+        features,
+        train_set,
+        valid_set,
+        num_classes,
+        model,
+        device,
+    )
+
+    # Test the model.
+    if rank == 0:
+        print("Testing...")
+    test_set = dataset.tasks[0].test_set
+    test_acc, num_test_items = evaluate(
+        rank,
+        model,
+        graph,
+        features,
+        itemset=test_set,
+        num_classes=num_classes,
+        device=device,
+    )
+    test_acc = weighted_reduce(test_acc * num_test_items, num_test_items)
+
+    if rank == 0:
+        print(f"Test Accuracy {test_acc.item():.4f}")
+
+
+######################################################################
+# Spawning Trainer Processes
+# --------------------------
+#
+# The following code spawns a process for each GPU and calls the ``run``
+# function defined above.
+#
+
+
+def main():
+    if not torch.cuda.is_available():
+        print("No GPU found!")
+        return
+
+    devices = [
+        torch.device(f"cuda:{i}") for i in range(torch.cuda.device_count())
+    ]
+    world_size = len(devices)
+
+    print(f"Training with {world_size} gpus.")
+
+    # Load and preprocess dataset.
+    dataset = gb.BuiltinDataset("ogbn-arxiv").load()
+
+    # Thread limiting to avoid resource competition.
+    os.environ["OMP_NUM_THREADS"] = str(mp.cpu_count() // 2 // world_size)
+
+    mp.set_sharing_strategy("file_system")
+    mp.spawn(
+        run,
+        args=(world_size, devices, dataset),
+        nprocs=world_size,
+        join=True,
+    )
+
+
+if __name__ == "__main__":
+    main()

tutorials/multi/README.txt

@@ -0,0 +1,2 @@
+Training on Multiple GPUs
+=========================

tutorials/requirements.txt

@@ -0,0 +1,9 @@
+networkx>=2.1
+torch
+numpy
+seaborn
+matplotlib
+pygraphviz
+graphviz
+pandas
+rdflib