5

GNN 教程：DGL框架-消息和GCN的实现

 2 years ago
source link: https://archwalker.github.io/blog/2019/07/07/GNN-Framework-DGL-GCN.html
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.
neoserver,ios ssh client

此为原创文章，转载务必保留出处

图神经网络的计算模式大致相似，节点的Embedding需要汇聚其邻接节点Embedding以更新，从线性代数的角度来看，这就是邻接矩阵和特征矩阵相乘。然而邻接矩阵通常都会很大，因此另一种计算方法是将邻居的Embedding传递到当前节点上，再进行更新。很多图并行框架都采用详细传递的机制进行运算(比如Google的Pregel)。而图神经网络框架DGL也采用了这样的思路。从本篇博文开始，我们DGL](https://docs.dgl.ai/index.html)做一个系统的介绍，我们主要关注他的设计，尤其是应对大规模图计算的设计。这篇文章将会介绍DGL的核心概念 — 消息传递机制，并且使用DGL框架实现GCN算法。

DGL 核心 — 消息传递

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

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

  • 汇聚函数（reduce function）：汇聚函数的目的是根据邻居传过来的消息更新跟新自身节点Embedding，对每个节点来说，它先从邮箱（v.mailbox[‘m’]）中汇聚消息函数所传递过来的消息（message），并清空邮箱（v.mailbox[‘m’]）内消息；然后该节点结合汇聚后的结果和该节点原Embedding，更新节点Embedding。

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

用消息传递的方式实现GCN

GCN 的线性代数表达

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

H(l+1)=σ(~D−12~A~D−12H(l)W(l))(1)(1)H(l+1)=σ(D~−12A~D~−12H(l)W(l))

从线性代数的角度，节点EmbeddingH(l+1)H(l+1)的的更新方式为首先左乘邻接矩阵~D−12~A~D−12D~−12A~D~−12以汇聚邻居Embedding，再为新Embedding做一次线性变换(右乘W(l)W(l))，简而言之：每个节点拿到邻居节点信息汇聚到自身 embedding 上在进行一次变换。具体 GCN 内容介绍可参考之前的博文。

从消息传递的角度分析

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

  1. 在 GCN 中每个节点都有属于自己的表示 hihi;
  2. 根据消息传递（message passing）的范式，每个节点将会收到来自邻居节点发送的Embedding；
  3. 每个节点将会对来自邻居节点的 Embedding进行汇聚以得到中间表示 ^hih^i ；
  4. 对中间节点表示 ^hih^i 进行线性变换，然后在利用非线性函数ff进行计算：hnewu=f(Wu^hu)hunew=f(Wuh^u);
  5. 利用新的节点表示 hnewuhunew 对该节点的表示 huhu进行更新。

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)，即定义一个全连接层（fully-connected layer）来对中间节点表示 ^hih^i 进行线性变换，然后在利用非线性函数ff进行计算：hnewu=f(Wu^hu)hunew=f(Wuh^u)。

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


report

Recommend

About Joyk


Aggregate valuable and interesting links.
Joyk means Joy of geeK