GNN学习笔记（二）：消息传递图神经网络（Message Passing GCN）

目录

一、MessagePassing基类初步分析

 二、MessagePassing子类实例

 三、message方法的覆写

 四、aggregate方法的覆写

五、message_and_aggregate方法的覆写

六、update方法的覆写

---

消息传递范式是一种聚合邻接节点信息来更新中心节点信息的范式，它将卷积算子推广到了不规则数据领域，实现了图与神经网络的连接。该范式包含这样三个步骤：

1. 邻接节点信息变换
2. 邻接节点信息聚合到中心节点
3. 聚合信息变换

因为简单且强大的特性，消息传递范式现被人们广泛地使用。基于此范式，我们可以定义聚合邻接节点信息来生成中心节点表征的图神经网络。在PyG中，MessagePassing基类是所有基于消息传递范式的图神经网络的基类，它大大地方便了我们对图神经网络的构建。

「来源于Datawhale」

一、MessagePassing基类初步分析

MessagePassing是torch_geometric中GNN模型的基类，要继承这个类，需要复写三个函数：propagate(edge_index, size=None)、message()、update()，其中propagate在执行的过程中会调用message和update。假设顶点V1和顶点v2,v3,v4,.....vn有边相连，propagate做的事情是将v2,v3,v4,.....vn的信息加（默认‘add’，也可以‘mean’，‘max’）到v1上。GCN的实现，三个函数都是在MessagePassing的基础上实现的。

1. __init__()。对象初始化方法。

MessagePassing(aggr="add", flow="source_to_target", node_dim=-2):

2. propagate()。开始传递消息的起始调用，在此方法中message、update等方法被调用 

MessagePassing.propagate(edge_index, size=None, **kwargs):

注意：

• size=none 意味着邻接矩阵为对称阵。

• kwargs: 图其他属性或额外的数据。

3. message()。首先确定要给节点传递消息的边的集合（运行flow="source_to_target"或者flow="target_to_source"），接着为各条边创建要传递给节点的消息，即实现函数。

MessagePassing.message(...):

4. aggregate()，即将从源节点传递过来的消息聚合在目标节点上，常用加和（sum）、平均（mean）以及取最大（max）的方式。

MessagePassing.aggregate(...):

5. message_and_aggregate()，即融合邻接节点信息变换和邻接节点信息聚合，如果在此方法里定义这项操作可以使得程序运行更加高效。

MessagePassing.message_and_aggregate(...):

6. update()。为每个节点更新节点表征，即实现​函数。此方法以aggregate方法的输出为第一个参数，并接收所有传递给propagate()方法的参数。

MessagePassing.update(aggr_out, ...): 

以上内容来源于The “MessagePassing” Base Class。

 二、MessagePassing子类实例

“我们以继承MessagePassing基类的GCNConv类为例，学习如何通过继承MessagePassing基类来实现一个简单的图神经网络”——这部分的内容可参见网站中的消息传递图神经网络.md：https://github.com/datawhalechina/team-learning-nlp/tree/master/GNN/Markdown版本

唠叨小主粘贴链接中的代码，便于温故而知新，如下所示。

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:向邻接矩阵添加自环边.
        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))

        # Step 2: 调用`torch.nn.Linear`实例对节点表征进行线性变换.
        x = self.lin(x)

        # Step 3: 计算归一化系数.
        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: 归一化邻接节点的节点表征.将相邻节点表征相加（"求和 "聚合）.
        return self.propagate(edge_index, x=x, norm=norm)

    def message(self, x_j, norm):
        # x_j has shape [E, out_channels]
        # Step 4: 通过`norm`对邻接节点表征`x_j`进行归一化处理.
        return norm.view(-1, 1) * x_j

from torch_geometric.datasets import Planetoid

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

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

 GCNConv继承了MessagePassing并以"求和"作为领域节点信息聚合方式。该层的所有逻辑都发生在其forward()方法中。

 三、message方法的覆写

我们希望meassge方法还能接收中心节点的度，所以对前面GCNConv的message方法进行改造得到新的GCNConv类，如下：

（1）在def forward(self, x, edge_index):中将

return self.propagate(edge_index, x=x, norm=norm)

更改为：

 return self.propagate(edge_index, x=x, norm=norm, deg=deg.view((-1, 1)))

（2）在 def message(self, x_j, norm):中将

 return norm.view(-1, 1) * x_j

更改为：

return norm.view(-1, 1) * x_j * deg_i

 四、aggregate方法的覆写

（1）在class GCNConv(MessagePassing):中添加

def aggregate(self, inputs, index, ptr, dim_size):
        print('self.aggr:', self.aggr)
        print("`aggregate` is called")
        return super().aggregate(inputs, index, ptr=ptr, dim_size=dim_size)
        

五、message_and_aggregate方法的覆写

（1）在原代码def forward(self, x, edge_index):中，增加下面这条代码，放置在 “Step 4-5”下面

 adjmat = SparseTensor(row=edge_index[0], col=edge_index[1], value=torch.ones(edge_index.shape[1]))

（2）并更改下面这条代码

return self.propagate(edge_index, x=x, norm=norm)

更改为：

return self.propagate(adjmat, x=x, norm=norm, deg=deg.view((-1, 1)))

从代码的区别中，我们可以看到， 此处传的不再是edge_idex，而是SparseTensor类型的Adjancency Matrix。

（3）在class GCNConv(MessagePassing):中添加

def message_and_aggregate(self, adj_t, x, norm):
        print('`message_and_aggregate` is called')
        # 没有实现真实的消息传递与消息聚合的操作

六、update方法的覆写

（1）在class GCNConv(MessagePassing):中添加

def update(self, inputs, deg):
        print(deg)
        return inputs

以上就是MessagePassing基类的运行流程。

七、MessagePassing基类的运行流程——总结

首先导入必要的包，其次创建一个类，继承MessagePassing基类，然后初始化，定义forward（）方法用于模型的表示，接着定义massage（）方法、aggregate（）方法、message_and_aggregate()方法、update()方法。

参考资料：

https://github.com/datawhalechina/team-learning-nlp/tree/master/GNN/Markdown版本

https://pytorch-geometric.readthedocs.io/en/latest/notes/create_gnn.html

https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html

https://pytorch-geometric.readthedocs.io/en/latest/notes/create_gnn.html

https://pytorch-geometric.readthedocs.io/en/latest/notes/create_gnn.html