GNN教程：开山之作GCN代码理解（DGL框架）

转载

目录

大纲

DGL 核心 — 消息传递

用消息传递的方式实现GCN

GCN 的线性代数表达

从消息传递的角度分析

具体实现

step 1，引入相关包

step 2，需要定义 GCN 的 message 函数和 reduce 函数

step 3，定义一个应用于节点的 node UDF(user defined function)

step 4，定义 GCN 的Embedding更新层

step 5，定义了一个包含两个 GCN 层的图神经网络分类器

step 6，加载 cora 数据集，并进行数据预处理

step 7，训练 GCN 神经网络

后话

Reference

---

大纲

本文为GNN教程的第七篇文章【使用DGL框架实现GCN算法】。图神经网络的计算模式大致相似，节点的Embedding需要汇聚其邻接节点Embedding以更新，从线性代数的角度来看，这就是邻接矩阵和特征矩阵相乘。然而邻接矩阵通常都会很大，因此另一种计算方法是将邻居的Embedding传递到当前节点上，再进行更新。很多图并行框架都采用详细传递的机制进行运算(比如Google的Pregel)。而图神经网络框架DGL也采用了这样的思路。

从本篇博文开始，我们使用DGL做一个系统的介绍，我们主要关注他的设计，尤其是应对大规模图计算的设计。这篇文章将会介绍DGL的核心概念 — 消息传递机制，并且使用DGL框架实现GCN算法。

DGL 核心 — 消息传递

DGL 的核心为消息传递机制（message passing），主要分为消息函数 （message function）和汇聚函数（reduce function）。如下图所示：

• 消息函数（message function）：传递消息的目的是将节点计算时需要的信息传递给它，因此对每条边来说，每个源节点将会将自身的Embedding（e.src.data）和边的Embedding(edge.data)传递到目的节点；对于每个目的节点来说，它可能会受到多个源节点传过来的消息，它会将这些消息存储在"邮箱"中。（gcn_msg：就是将边消息和邻居消息h传到中心节点，存到中心节点的邮箱mailbox中,记做m）

• 汇聚函数（reduce function）：汇聚函数的目的是根据邻居传过来的消息更新跟新自身节点Embedding，对每个节点来说，它先从邮箱（v.mailbox['m']）中汇聚消息函数所传递过来的消息（message），并清空邮箱（v.mailbox['m']）内消息（gcn_reduce：每个节点从邮箱中汇聚所有传过来的信息m，并清空，记做h）；然后该节点结合汇聚后的结果和该节点原Embedding，更新节点Embedding（需要做个线性变换NodeApplyModule，并且在线性变换的基础上加上自身信息来更新节点表示GCN：update_all+apply_nodes）。

下面我们以GCN的算法为例，详细说明消息传递的机制是如何work的。

用消息传递的方式实现GCN

GCN 的线性代数表达

GCN 的逐层传播公式如下所示：

简而言之：每个节点拿到邻居节点信息汇聚到自身 embedding 上在进行一次变换。具体 GCN 内容介绍可参考之前的文章

从消息传递的角度分析

上面的数学描述可以利用消息传递的机制实现为：

具体实现

step 1，引入相关包

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

step 2，需要定义 GCN 的 message 函数和 reduce 函数

 message 函数用于发送节点的Embedding，reduce 函数用来对收到的 Embedding 进行聚合。在这里，每个节点发送Embedding的时候不需要任何处理，所以可以通过内置的copy_scr实现，out='m'表示发送到目的节点后目的节点的mailbox用m来标识这个消息是源节点的Embedding。

目的节点的reduce函数很简单，因为按照GCN的数学定义，邻接矩阵和特征矩阵相乘，以为这更新后的特征矩阵的每一行是原特征矩阵某几行相加的形式，"某几行"是由邻接矩阵选定的，即对应节点的邻居所在的行。因此目的节点reduce只需要通过sum将接受到的信息相加就可以了。

gcn_msg = fn.copy_src(src='h', out='m')
gcn_reduce = fn.sum(msg='m', out='h')

step 3，定义一个应用于节点的 node UDF(user defined function)

class NodeApplyModule(nn.Module):
    def __init__(self, in_feats, out_feats, activation):
        super(NodeApplyModule, self).__init__()
        self.linear = nn.Linear(in_feats, out_feats)
        self.activation = activation

    def forward(self, node):
        h = self.linear(node.data['h'])
        h = self.activation(h)
        return {'h' : h}

step 4，定义 GCN 的Embedding更新层

以实现在所有节点上进行消息传递，并利用 NodeApplyModule 对节点信息进行计算更新。

class GCN(nn.Module):
    def __init__(self, in_feats, out_feats, activation):
        super(GCN, self).__init__()
        self.apply_mod = NodeApplyModule(in_feats, out_feats, activation)

    def forward(self, g, feature):
        g.ndata['h'] = feature
        g.update_all(gcn_msg, gcn_reduce)
        g.apply_nodes(func=self.apply_mod)
        return g.ndata.pop('h')

step 5，定义了一个包含两个 GCN 层的图神经网络分类器

我们通过向该分类器输入特征大小为 1433 的训练样本，以获得该样本所属的类别编号，类别总共包含 7 类。

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.gcn1 = GCN(1433, 16, F.relu)
        self.gcn2 = GCN(16, 7, F.relu)

    def forward(self, g, features):
        x = self.gcn1(g, features)
        x = self.gcn2(g, x)
        return x
net = Net()
print(net)

step 6，加载 cora 数据集，并进行数据预处理

from dgl.data import citation_graph as citegrh
def load_cora_data():
    data = citegrh.load_cora()
    features = th.FloatTensor(data.features)
    labels = th.LongTensor(data.labels)
    mask = th.ByteTensor(data.train_mask)
    g = data.graph
    # add self loop
    g.remove_edges_from(g.selfloop_edges())
    g = DGLGraph(g)
    g.add_edges(g.nodes(), g.nodes())
    return g, features, labels, mask

step 7，训练 GCN 神经网络

import time
import numpy as np
g, features, labels, mask = load_cora_data()
optimizer = th.optim.Adam(net.parameters(), lr=1e-3)
dur = []
for epoch in range(30):
    if epoch >=3:
        t0 = time.time()

    logits = net(g, 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)))

后话

本篇博文介绍了如何利用图神经网络框架DGL编写GCN模型，接下来我们会介绍如何利用DGL实现GraphSAGE中的采样机制，以减少运算规模。

Reference

1. DGL Basics

2. Graph Convolutional Network

3. PageRank with DGL Message Passing

4. DGL 作者答疑！关于 DGL 你想知道的都在这里