.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "tutorials/large/L4_message_passing.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_tutorials_large_L4_message_passing.py: Writing GNN Modules for Stochastic GNN Training =============================================== All GNN modules DGL provides support stochastic GNN training. This tutorial teaches you how to write your own graph neural network module for stochastic GNN training. It assumes that 1. You know :doc:`how to write GNN modules for full graph training <../blitz/3_message_passing>`. 2. You know :doc:`how stochastic GNN training pipeline works `. .. GENERATED FROM PYTHON SOURCE LINES 15-49 .. code-block:: Python import os os.environ["DGLBACKEND"] = "pytorch" import dgl import numpy as np import torch from ogb.nodeproppred import DglNodePropPredDataset dataset = DglNodePropPredDataset("ogbn-arxiv") device = "cpu" # change to 'cuda' for GPU graph, node_labels = dataset[0] # Add reverse edges since ogbn-arxiv is unidirectional. graph = dgl.add_reverse_edges(graph) graph.ndata["label"] = node_labels[:, 0] idx_split = dataset.get_idx_split() train_nids = idx_split["train"] node_features = graph.ndata["feat"] sampler = dgl.dataloading.MultiLayerNeighborSampler([4, 4]) train_dataloader = dgl.dataloading.DataLoader( graph, train_nids, sampler, batch_size=1024, shuffle=True, drop_last=False, num_workers=0, ) input_nodes, output_nodes, mfgs = next(iter(train_dataloader)) .. rst-class:: sphx-glr-script-out .. code-block:: none /home/ubuntu/regression_test/dgl/python/dgl/dataloading/dataloader.py:1149: DGLWarning: Dataloader CPU affinity opt is not enabled, consider switching it on (see enable_cpu_affinity() or CPU best practices for DGL [https://docs.dgl.ai/tutorials/cpu/cpu_best_practises.html]) dgl_warning( .. GENERATED FROM PYTHON SOURCE LINES 50-62 DGL Bipartite Graph Introduction -------------------------------- In the previous tutorials, you have seen the concept *message flow graph* (MFG), where nodes are divided into two parts. It is a kind of (directional) bipartite graph. This section introduces how you can manipulate (directional) bipartite graphs. You can access the source node features and destination node features via ``srcdata`` and ``dstdata`` attributes: .. GENERATED FROM PYTHON SOURCE LINES 62-68 .. code-block:: Python mfg = mfgs[0] print(mfg.srcdata) print(mfg.dstdata) .. rst-class:: sphx-glr-script-out .. code-block:: none {'year': tensor([[2017], [2015], [2016], ..., [2020], [2017], [2019]]), 'feat': tensor([[-0.1797, -0.2894, -0.2900, ..., 0.2316, 0.0188, -0.1484], [-0.0015, 0.0283, -0.3957, ..., 0.0549, -0.1626, -0.0565], [-0.1152, 0.0897, -0.0365, ..., 0.1176, 0.0125, -0.0534], ..., [-0.1217, 0.1248, -0.1144, ..., -0.0179, -0.1192, -0.0170], [-0.1477, 0.0288, -0.2111, ..., 0.0983, -0.0368, -0.0379], [ 0.0081, 0.3123, -0.1327, ..., 0.0163, -0.1030, -0.1363]]), 'label': tensor([16, 5, 2, ..., 24, 16, 16]), '_ID': tensor([ 42905, 31998, 46238, ..., 126880, 163148, 134259])} {'year': tensor([[2017], [2015], [2016], ..., [2013], [2015], [2015]]), 'feat': tensor([[-0.1797, -0.2894, -0.2900, ..., 0.2316, 0.0188, -0.1484], [-0.0015, 0.0283, -0.3957, ..., 0.0549, -0.1626, -0.0565], [-0.1152, 0.0897, -0.0365, ..., 0.1176, 0.0125, -0.0534], ..., [-0.1347, -0.0979, -0.0540, ..., 0.1124, 0.0021, -0.2762], [-0.2496, 0.0399, -0.1484, ..., -0.2209, 0.0317, -0.1959], [-0.1497, -0.0160, -0.1847, ..., 0.0413, -0.0658, -0.1807]]), 'label': tensor([16, 5, 2, ..., 16, 16, 16]), '_ID': tensor([42905, 31998, 46238, ..., 3414, 27311, 99465])} .. GENERATED FROM PYTHON SOURCE LINES 69-72 It also has ``num_src_nodes`` and ``num_dst_nodes`` functions to query how many source nodes and destination nodes exist in the bipartite graph: .. GENERATED FROM PYTHON SOURCE LINES 72-76 .. code-block:: Python print(mfg.num_src_nodes(), mfg.num_dst_nodes()) .. rst-class:: sphx-glr-script-out .. code-block:: none 13017 4182 .. GENERATED FROM PYTHON SOURCE LINES 77-80 You can assign features to ``srcdata`` and ``dstdata`` just as what you will do with ``ndata`` on the graphs you have seen earlier: .. GENERATED FROM PYTHON SOURCE LINES 80-85 .. code-block:: Python mfg.srcdata["x"] = torch.zeros(mfg.num_src_nodes(), mfg.num_dst_nodes()) dst_feat = mfg.dstdata["feat"] .. GENERATED FROM PYTHON SOURCE LINES 86-91 Also, since the bipartite graphs are constructed by DGL, you can retrieve the source node IDs (i.e. those that are required to compute the output) and destination node IDs (i.e. those whose representations the current GNN layer should compute) as follows. .. GENERATED FROM PYTHON SOURCE LINES 91-95 .. code-block:: Python mfg.srcdata[dgl.NID], mfg.dstdata[dgl.NID] .. rst-class:: sphx-glr-script-out .. code-block:: none (tensor([ 42905, 31998, 46238, ..., 126880, 163148, 134259]), tensor([42905, 31998, 46238, ..., 3414, 27311, 99465])) .. GENERATED FROM PYTHON SOURCE LINES 96-99 Writing GNN Modules for Bipartite Graphs for Stochastic Training ---------------------------------------------------------------- .. GENERATED FROM PYTHON SOURCE LINES 102-110 Recall that the MFGs yielded by the ``DataLoader`` have the property that the first few source nodes are always identical to the destination nodes: |image1| .. |image1| image:: https://data.dgl.ai/tutorial/img/bipartite.gif .. GENERATED FROM PYTHON SOURCE LINES 110-118 .. code-block:: Python print( torch.equal( mfg.srcdata[dgl.NID][: mfg.num_dst_nodes()], mfg.dstdata[dgl.NID] ) ) .. rst-class:: sphx-glr-script-out .. code-block:: none True .. GENERATED FROM PYTHON SOURCE LINES 119-122 Suppose you have obtained the source node representations :math:`h_u^{(l-1)}`: .. GENERATED FROM PYTHON SOURCE LINES 122-126 .. code-block:: Python mfg.srcdata["h"] = torch.randn(mfg.num_src_nodes(), 10) .. GENERATED FROM PYTHON SOURCE LINES 127-139 Recall that DGL provides the `update_all` interface for expressing how to compute messages and how to aggregate them on the nodes that receive them. This concept naturally applies to bipartite graphs like MFGs -- message computation happens on the edges between source and destination nodes of the edges, and message aggregation happens on the destination nodes. For example, suppose the message function copies the source feature (i.e. :math:`M^{(l)}\left(h_v^{(l-1)}, h_u^{(l-1)}, e_{u\to v}^{(l-1)}\right) = h_v^{(l-1)}`), and the reduce function averages the received messages. Performing such message passing computation on a bipartite graph is no different than on a full graph: .. GENERATED FROM PYTHON SOURCE LINES 139-147 .. code-block:: Python import dgl.function as fn mfg.update_all(message_func=fn.copy_u("h", "m"), reduce_func=fn.mean("m", "h")) m_v = mfg.dstdata["h"] m_v .. rst-class:: sphx-glr-script-out .. code-block:: none tensor([[ 0.4055, 0.0349, 0.1661, ..., 0.1347, -0.5621, -0.1685], [-0.2778, -0.4114, 0.2081, ..., -0.2586, -0.3285, -0.3846], [-0.1454, 0.3454, 0.5464, ..., 0.8526, -0.4063, 0.2639], ..., [-1.1680, 0.4332, -0.7955, ..., -0.0908, 0.1632, -0.1003], [ 0.1939, 0.1668, -0.0243, ..., 0.0857, 0.3781, -0.1000], [ 0.0922, -0.6378, 0.5134, ..., 0.3548, -0.0424, -0.1413]]) .. GENERATED FROM PYTHON SOURCE LINES 148-152 Putting them together, you can implement a GraphSAGE convolution for training with neighbor sampling as follows (the differences to the :doc:`full graph counterpart <../blitz/3_message_passing>` are highlighted with arrows ``<---``) .. GENERATED FROM PYTHON SOURCE LINES 152-230 .. code-block:: Python import torch.nn as nn import torch.nn.functional as F import tqdm 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 MFG. h : (Tensor, Tensor) The feature of source nodes and destination nodes as a pair of Tensors. """ with g.local_scope(): h_src, h_dst = h g.srcdata["h"] = h_src # <--- g.dstdata["h"] = h_dst # <--- # update_all is a message passing API. g.update_all(fn.copy_u("h", "m"), fn.mean("m", "h_N")) h_N = g.dstdata["h_N"] h_total = torch.cat([h_dst, h_N], dim=1) # <--- return self.linear(h_total) 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, mfgs, x): h_dst = x[: mfgs[0].num_dst_nodes()] h = self.conv1(mfgs[0], (x, h_dst)) h = F.relu(h) h_dst = h[: mfgs[1].num_dst_nodes()] h = self.conv2(mfgs[1], (h, h_dst)) return h sampler = dgl.dataloading.MultiLayerNeighborSampler([4, 4]) train_dataloader = dgl.dataloading.DataLoader( graph, train_nids, sampler, device=device, batch_size=1024, shuffle=True, drop_last=False, num_workers=0, ) model = Model(graph.ndata["feat"].shape[1], 128, dataset.num_classes).to(device) with tqdm.tqdm(train_dataloader) as tq: for step, (input_nodes, output_nodes, mfgs) in enumerate(tq): inputs = mfgs[0].srcdata["feat"] labels = mfgs[-1].dstdata["label"] predictions = model(mfgs, inputs) .. rst-class:: sphx-glr-script-out .. code-block:: none 0%| | 0/89 [00:00` *and* :doc:`stochastic training `. Say you start with a GNN module that works for full-graph training only: .. GENERATED FROM PYTHON SOURCE LINES 247-287 .. code-block:: Python 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().__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) .. GENERATED FROM PYTHON SOURCE LINES 288-366 **First step**: Check whether the input feature is a single tensor or a pair of tensors: .. code:: python if isinstance(h, tuple): h_src, h_dst = h else: h_src = h_dst = h **Second step**: Replace node features ``h`` with ``h_src`` or ``h_dst``, and assign the node features to ``srcdata`` or ``dstdata``, instead of ``ndata``. Whether to assign to ``srcdata`` or ``dstdata`` depends on whether the said feature acts as the features on source nodes or destination nodes of the edges in the message functions (in ``update_all`` or ``apply_edges``). *Example 1*: For the following ``update_all`` statement: .. code:: python g.ndata['h'] = h g.update_all(message_func=fn.copy_u('h', 'm'), reduce_func=fn.mean('m', 'h_N')) The node feature ``h`` acts as source node feature because ``'h'`` appeared as source node feature. So you will need to replace ``h`` with source feature ``h_src`` and assign to ``srcdata`` for the version that works with both cases: .. code:: python g.srcdata['h'] = h_src g.update_all(message_func=fn.copy_u('h', 'm'), reduce_func=fn.mean('m', 'h_N')) *Example 2*: For the following ``apply_edges`` statement: .. code:: python g.ndata['h'] = h g.apply_edges(fn.u_dot_v('h', 'h', 'score')) The node feature ``h`` acts as both source node feature and destination node feature. So you will assign ``h_src`` to ``srcdata`` and ``h_dst`` to ``dstdata``: .. code:: python g.srcdata['h'] = h_src g.dstdata['h'] = h_dst # The first 'h' corresponds to source feature (u) while the second 'h' corresponds to destination feature (v). g.apply_edges(fn.u_dot_v('h', 'h', 'score')) .. note:: For homogeneous graphs (i.e. graphs with only one node type and one edge type), ``srcdata`` and ``dstdata`` are aliases of ``ndata``. So you can safely replace ``ndata`` with ``srcdata`` and ``dstdata`` even for full-graph training. **Third step**: Replace the ``ndata`` for outputs with ``dstdata``. For example, the following code .. code:: python # Assume that update_all() function has been called with output node features in `h_N`. h_N = g.ndata['h_N'] h_total = torch.cat([h, h_N], dim=1) will change to .. code:: python h_N = g.dstdata['h_N'] h_total = torch.cat([h_dst, h_N], dim=1) .. GENERATED FROM PYTHON SOURCE LINES 369-372 Putting together, you will change the ``SAGEConvForBoth`` module above to something like the following: .. GENERATED FROM PYTHON SOURCE LINES 372-419 .. code-block:: Python class SAGEConvForBoth(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().__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 or tuple[Tensor, Tensor] The input node feature. """ with g.local_scope(): if isinstance(h, tuple): h_src, h_dst = h else: h_src = h_dst = h g.srcdata["h"] = h_src # 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_dst, h_N], dim=1) return self.linear(h_total) # Thumbnail credits: Representation Learning on Networks, Jure Leskovec, WWW 2018 # sphinx_gallery_thumbnail_path = '_static/blitz_3_message_passing.png' .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 0.646 seconds) .. _sphx_glr_download_tutorials_large_L4_message_passing.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: L4_message_passing.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: L4_message_passing.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: L4_message_passing.zip ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_