GNN-2-消息传递图神经网络

参考开源学习地址：datawhale

1. 消息传递范式介绍

• 消息传递范式是一种聚合邻接节点信息来更新中心节点信息的范式，它将卷积算子推广到了不规则数据领域，实现了图与神经网络的连接。
• 该范式包含这样三个步骤：(1)邻接节点信息变换、(2)邻接节点信息聚合到中心节点、(3)聚合信息变换。
• 因为简单且强大的特性，消息传递范式现被人们广泛地使用。基于此范式，我们可以定义聚合邻接节点信息来生成中心节点表征的图神经网络。
• 在PyG中，MessagePassing基类是所有基于消息传递范式的图神经网络的基类，它大大地方便了我们对图神经网络的构建。

下方图片展示了基于消息传递范式的聚合邻接节点信息来更新中心节点信息的过程：

• 图中黄色方框部分展示的是一次邻居节点信息传递到中心节点的过程：B节点的邻接节点（A,C）的信息经过变换后聚合到B节点，接着B节点信息与邻居节点聚合信息一起经过变换得到B节点的新的节点信息。同时，分别如红色和绿色方框部分所示，遵循同样的过程，C、D节点的信息也被更新。实际上，同样的过程在所有节点上都进行了一遍，所有节点的信息都更新了一遍。
• 这样的“邻居节点信息传递到中心节点的过程”会进行多次。如图中蓝色方框部分所示，A节点的邻接节点（B,C,D）的已经发生过一次更新的节点信息，经过变换、聚合、再变换产生了A节点第二次更新的节点信息。多次更新后的节点信息就作为节点表征。
  

2. MessagePassing基类初步分析

Pytorch Geometric(PyG)提供了MessagePassing基类，它封装了“消息传递”的运行流程。通过继承MessagePassing基类，可以方便地构造消息传递图神经网络。构造一个最简单的消息传递图神经网络类，我们只需定义message()方法（ 𝜙(…) ）、update()方法（ 𝛾(…) ），以及使用的消息聚合方案（aggr=”add”、aggr=”mean”或aggr=”max”。MessagePassing Base Class中这里最重要的3个函数是：

• MessagePassing.aggregate(…)：用于处理聚集到节点的信息的函数
• MessagePassing.message(…)：用于搭建传送到 node i 的节点消息，相对于𝜙(…)函
• MessagePassing.update(aggr_out, …): 用于更新节点的信息，相对于𝛾(…)

3.MessagePassing子类实例

我们以继承MessagePassing基类的GCNConv类为例，学习如何通过继承MessagePassing基类来实现一个简单的图神经网络。

GCNConv的数学定义为
x i ( k ) = ∑ j ∈ N ( i ) ∪ { i } 1 deg ⁡ ( i ) ⋅ deg ⁡ ( j ) ⋅ ( Θ ⋅ x j ( k − 1 ) ) , \mathbf{x}_i^{(k)} = \sum_{j \in \mathcal{N}(i) \cup \{ i \}} \frac{1}{\sqrt{\deg(i)} \cdot \sqrt{\deg(j)}} \cdot \left( \mathbf{\Theta} \cdot \mathbf{x}_j^{(k-1)} \right), xi(k)​=j∈N(i)∪{i}∑​deg(i) ​⋅deg(j) ​1​⋅(Θ⋅xj(k−1)​),
其中，邻接节点的表征${\mathbf{x}}_j^{(k-1)}$首先通过与权重矩阵$\mathbf{\Theta }$相乘进行变换，然后按端点的度$\mathrm{deg}(i),\mathrm{deg}(j)$进行归一化处理，最后进行求和。这个公式可以分为以下几个步骤：

1. 向邻接矩阵添加自环边。
2. 对节点表征做线性转换。
3. 计算归一化系数。
4. 归一化邻接节点的节点表征。
5. 将相邻节点表征相加（"求和 "聚合）。

步骤1-3通常是在消息传递发生之前计算的。步骤4-5可以使用MessagePassing基类轻松处理。该层的全部实现如下所示。

import torch
from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops, degree

class GCNConv(MessagePassing):
    def __init__(self, in_channels, out_channels):
        super(GCNConv, self).__init__(aggr='add', flow='source_to_target')
        # "Add" aggregation (Step 5).
        # flow='source_to_target' 表示消息从源节点传播到目标节点
        self.lin = torch.nn.Linear(in_channels, out_channels)

    def forward(self, x, edge_index):
        # x has shape [N, in_channels]
        # edge_index has shape [2, E]

        # Step 1: Add self-loops to the adjacency matrix.
        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))

        # Step 2: Linearly transform node feature matrix.
        x = self.lin(x)

        # Step 3: Compute normalization.
        row, col = edge_index
        deg = degree(col, x.size(0), dtype=x.dtype)
        deg_inv_sqrt = deg.pow(-0.5)
        norm = deg_inv_sqrt[row] * deg_inv_sqrt[col]

        # Step 4-5: Start propagating messages.
        return self.propagate(edge_index, x=x, norm=norm)

    def message(self, x_j, norm):
        # x_j has shape [E, out_channels]
        # Step 4: Normalize node features.
        return norm.view(-1, 1) * x_j
        

GCNConv继承了MessagePassing并以"求和"作为领域节点信息聚合方式。该层的所有逻辑都发生在其forward()方法中。在这里，我们首先使用torch_geometric.utils.add_self_loops()函数向我们的边索引添加自循环边（步骤1），以及通过调用torch.nn.Linear实例对节点表征进行线性变换（步骤2）。propagate()方法也在forward方法中被调用，propagate()方法被调用后节点间的信息传递开始执行。

归一化系数是由每个节点的节点度得出的，它被转换为每条边的节点度。结果被保存在形状为[num_edges,]的变量norm中（步骤3）。

在message()方法中，我们需要通过norm对邻接节点表征x_j进行归一化处理。

通过以上内容的学习，我们便掌握了创建一个仅包含一次“消息传递过程”的图神经网络的方法。如下方代码所示，我们可以很方便地初始化和调用它：

from torch_geometric.datasets import Planetoid

dataset = Planetoid(root='dataset/Cora', name='Cora')
data = dataset[0]

# data.x: a matrix with each row representing the data in a node
# data.edge_index: matrix with shape [2, number of edges], each column representing edge from node to another node, value=index of node
net = GCNConv(data.num_features, 64)
h_nodes = net(data.x, data.edge_index)
print(h_nodes.shape)

通过串联多个这样的简单图神经网络，我们就可以构造复杂的图神经网络模型**。我们将在第5节介绍复杂图神经网络模型的构建。

以上主要内容来源于Implementing the GCN Layer。

from torch_geometric.datasets import Planetoid

dataset = Planetoid(root='./dataset/Cora', name='Cora')
data = dataset[0]
# GCNConv: 
#in_channels: dimension of input vector of linear layer
# out_channels: dimension of output vector of linear layer
#Note: the linear transform is performed before message passing to reduce the dimension of node representation

net = GCNConv(data.num_features, 64)


h_nodes = net(data.x, data.edge_index)
print(h_nodes.shape)

4.message方法的覆写

1. MessagePassing基类剖析

在__init__()方法中，我们看到程序会检查子类是否实现了message_and_aggregate()方法，并将检查结果赋值给fuse属性。

class MessagePassing(torch.nn.Module):
	def __init__(self, aggr: Optional[str] = "add", flow: str = "source_to_target", node_dim: int = -2):
        super(MessagePassing, self).__init__()
		# 此处省略n行代码
        # Support for "fused" message passing.
        self.fuse = self.inspector.implements('message_and_aggregate')
		# 此处省略n行代码

“消息传递过程”是从propagate方法被调用开始执行的。

class MessagePassing(torch.nn.Module):
    # 此处省略n行代码
    def propagate(self, edge_index: Adj, size: Size = None, **kwargs):
    	# 此处省略n行代码
        # Run "fused" message and aggregation (if applicable).
        if (isinstance(edge_index, SparseTensor) and self.fuse and not self.__explain__):
            coll_dict = self.__collect__(self.__fused_user_args__, edge_index, size, kwargs)

            msg_aggr_kwargs = self.inspector.distribute('message_and_aggregate', coll_dict)
            out = self.message_and_aggregate(edge_index, **msg_aggr_kwargs)

            update_kwargs = self.inspector.distribute('update', coll_dict)
            return self.update(out, **update_kwargs)
        # Otherwise, run both functions in separation.
        elif isinstance(edge_index, Tensor) or not self.fuse:
            coll_dict = self.__collect__(self.__user_args__, edge_index, size, kwargs)

            msg_kwargs = self.inspector.distribute('message', coll_dict)
            out = self.message(**msg_kwargs)
    		# 此处省略n行代码
            aggr_kwargs = self.inspector.distribute('aggregate', coll_dict)
            out = self.aggregate(out, **aggr_kwargs)

            update_kwargs = self.inspector.distribute('update', coll_dict)
            return self.update(out, **update_kwargs)

propagate()方法首先检查edge_index是否为SparseTensor类型以及是否子类实现了message_and_aggregate()方法，如是就执行子类的message_and_aggregate方法；否则依次执行子类的message(),aggregate(),update()三个方法。

2. message方法的覆写

5.aggregate方法的覆写

在前面的例子的基础上，我们增加如下的aggregate方法。通过观察运行结果我们可以看到，我们覆写的aggregate方法被调用，同时在super(GCNConv, self).__init__(aggr='add')中传递给aggr参数的值被存储到了self.aggr属性中。

6. message_and_aggregate方法的覆写

在一些案例中，“消息传递”与“消息聚合”可以融合在一起。对于这种情况，我们可以覆写message_and_aggregate方法，在message_and_aggregate方法中一块实现“消息传递”与“消息聚合”，这样能使程序的运行更加高效。

7. 覆写update方法

8. Message Passing 机制总结

消息传递图神经网络遵循上述的“聚合邻接节点信息来更新中心节点信息的过程”，来生成节点表征。用${\mathbf{x}}_i^{(k-1)}\in {\mathbb{R}}^F$表示$(k-1)$层中节点$i$的节点表征，${\mathbf{e}}_{j,i}\in {\mathbb{R}}^D$ 表示从节点$j$到节点$i$的边的属性，消息传递图神经网络可以描述为
$${\mathbf{x}}_i^{(k)}={\gamma }^{(k)}\left({\mathbf{x}}_i^{(k-1)},{\square }_{j\in \mathcal{N}(i)}{\varphi }^{(k)}\left({\mathbf{x}}_i^{(k-1)},{\mathbf{x}}_j^{(k-1)},{\mathbf{e}}_{j,i}\right)\right),$$
其中$\square$表示可微分的、具有排列不变性（函数输出结果与输入参数的排列无关）的函数。具有排列不变性的函数有，和函数、均值函数和最大值函数，e.g., sum, mean or max。$\gamma$和$\varphi$表示可微分的函数，如MLPs（多层感知器）。此处内容来源于CREATING MESSAGE PASSING NETWORKS。

用GCN的公式举个栗子，就是

x i ( k ) = ∑ j ∈ N ( i ) ∪ { i } 1 deg ⁡ ( i ) ⋅ deg ⁡ ( j ) ⋅ ( Θ ⋅ x j ( k − 1 ) ) , \mathbf{x}_i^{(k)} = \sum_{j \in \mathcal{N}(i) \cup \{ i \}} \frac{1}{\sqrt{\deg(i)} \cdot \sqrt{\deg(j)}} \cdot \left( \mathbf{\Theta} \cdot \mathbf{x}_j^{(k-1)} \right), xi(k)​=j∈N(i)∪{i}∑​deg(i) ​⋅deg(j) ​1​⋅(Θ⋅xj(k−1)​),

参考：GNN-2-MessagePassing