（DataWhale）图神经网络Task02：消息传递图神经网络

消息传递范式MessagePassing

• 领域聚合
• 节点更新

$${\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$表示可微分的、具有排列不变性（函数输出结果与输入参数的排列无关）的函数。具有排列不变性的函数有，sum()函数、mean()函数和max()函数。

PyGMessagePassing基类初识与子类实例

MP基类

from torch_geometric.nn import MessagePassing，MessagePassing类中只给了message、aggregate、message_and_aggregate、update等函数名，实现方法要用户自定义。

• MessagePassing(aggr="add", flow="source_to_target", node_dim=-2)（对象初始化方法）：
  • aggr：定义要使用的聚合方案（“add”、"mean "或 “max”）；
  • flow：定义消息传递的流向（"source_to_target "或 “target_to_source”）；
  • node_dim：定义沿着哪个维度传播，默认值为-2，也就是节点表征张量x（[num_nodes, num_features]）的节点维度-第0维度，即第-2维度。
  • 注：MessagePassing(……)等同于MessagePassing.__init__(……)
• MessagePassing.propagate(edge_index, size=None, **kwargs)：
  • 开始消息传递，在此方法中message、update等方法被调用；
  • 它以edge_index（边的端点的索引）和flow（消息的流向）以及一些额外的数据为参数；
  • propagate()不局限于基于形状为[N, N]的对称邻接矩阵进行“消息传递过程”。基于非对称的邻接矩阵进行消息传递（当图为二部图时），需要传递参数size=(N, M)。
• MessagePassing.message(self, x_j: Tensor)：
  • 根据edge_index从邻接节点$j$向中心节点$i$“消息传递”，即该方法实现上述公式中的${\varphi }^{(k)}\left({\mathbf{x}}_i^{(k-1)},{\mathbf{x}}_j^{(k-1)},{\mathbf{e}}_{j,i}\right)$；
  • MessagePassing.message(...)方法可以接收传递给MessagePassing.propagate(edge_index, size=None, **kwargs)方法的所有参数；
  • 首先，确定要给节点$i$传递消息的边的集合：
    • 如果flow="source_to_target"，则是$(j,i)\in \mathcal{E}$的边的集合；
    • 如果flow="target_to_source"，则是$(i,j)\in \mathcal{E}$的边的集合。
  • 接着，为各条边创建要传递给节点$i$的消息，即实现$\varphi$函数。
• MessagePassing.aggregate(self, inputs: Tensor, index: Tensor, ptr: Optional[Tensor] = None, dim_size: Optional[int] = None)：
  • 将从源节点传递过来的消息聚合在目标节点上，即实现上述公式中的$\square_{j \in \mathcal{N}(i)} $；
  • 一般可选的聚合方式有sum, mean和max。
• MessagePassing.message_and_aggregate(self, adj_t: SparseTensor)：
  • 在一些场景里，邻接节点信息变换和邻接节点信息聚合这两项操作可以融合在一起。
• MessagePassing.update(self, inputs: Tensor):
  • 为每个节点$i\in \mathcal{V}$更新节点表征，即实现$\gamma$函数；
  • 此以aggregate方法的输出为第一个参数，并接收所有传递给propagate()方法的参数。

MP子类实例

通过继承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)}$邻接节点$j$的表征，$\mathbf{\Theta }$为权重矩阵，$\frac{1}{\sqrt{\mathrm{deg}(i)}\cdot \sqrt{\mathrm{deg}(j)}}$表示根据相邻节点度进行归一化处理。这个公式可以分为以下几个步骤：

1. 向邻接矩阵添加自环边；
2. 根据$\mathbf{\Theta }$对节点表征做线性转换；
3. 计算归一化系数$\frac{1}{\sqrt{\mathrm{deg}(i)}\cdot \sqrt{\mathrm{deg}(j)}}$；
4. 归一化邻接节点的节点表征；
5. 将相邻节点表征相加–sum聚合。

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')
        self.lin = torch.nn.Linear(in_channels, out_channels)  # 线性转换\Theta

    def forward(self, x, edge_index):
        # x : [N, in_channels]; edge_index ：[2, E]; N/E分别为节点和关系数量

        # Step 1: 向邻接矩阵添加自环边
        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))

        # Step 2: 对节点表征做线性转换
        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]  # [num_edges,]

        # Step 4-5: Message Passing
        return self.propagate(edge_index, x=x, norm=norm)

    def message(self, x_j, norm):
        # x_j : [E, out_channels]
        # Step 4: 归一化邻接节点的节点表征
        return norm.view(-1, 1) * x_j

上述GCNConv是一个仅包含一次“消息传递过程”的图神经网络，在Cora数据集上测试如下：

from torch_geometric.datasets import Planetoid

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

net = GCNConv(data.num_features, 64)
h_nodes = net(data.x, data.edge_index)
print(h_nodes.shape)   # torch.Size([2708, 64])

PyGMessagePassing自定义函数message、aggregate、message_and_aggregate、update覆写示例

• 增加一个额外的节点属性，即节点的度deg，更新原GCNConv类的meassge方法。

  class GCNConv(MessagePassing):
      def __init__(self, in_channels, out_channels):
          # 省略···
          
      def forward(self, x, edge_index):
          # 省略···
  
          return self.propagate(edge_index, x=x, norm=norm, deg=deg.view((-1, 1)))
  
      def message(self, x_j, norm, deg_i):
          # x_j : [E, out_channels]; 
          # deg_i : [E, 1]
  
          return norm.view(-1, 1) * x_j * deg_i
  

• aggregate方法的覆写。

  class GCNConv(MessagePassing):
      def __init__(self, in_channels, out_channels):
          # 省略···
  
      def forward(self, x, edge_index):
          # 省略···
  
          # Step 4-5: Message Passing
          return self.propagate(edge_index, x=x, norm=norm, deg=deg.view((-1, 1)))
  
      def message(self, x_j, norm, deg_i):
          # x_j : [E, out_channels]; 
          # deg_i : [E, 1]
  
          return norm.view(-1, 1) * x_j * deg_i
  
      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方法中一块实现“消息传递”与“消息聚合”。

  from torch_sparse import SparseTensor
  
  class GCNConv(MessagePassing):
      def __init__(self, in_channels, out_channels):
          # 省略···
  
      def forward(self, x, edge_index):
          # 省略···
  
          # Step 4-5: Message Passing
          adjmat = SparseTensor(row=edge_index[0], col=edge_index[1], value=torch.ones(edge_index.shape[1]))
          # 此处传的不再是edge_idex，而是SparseTensor类型的Adjancency Matrix
          return self.propagate(adjmat, x=x, norm=norm, deg=deg.view((-1, 1)))
  
      # 虽然同时覆写了`message`和`aggregate`，但只有`message_and_aggregate`被执行
      def message(self, x_j, norm, deg_i):
          return norm.view(-1, 1) * x_j * deg_i
  
      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)
  
      def message_and_aggregate(self, adj_t, x, norm):
          print('`message_and_aggregate` is called')
          # TODO 此处未实现真实的消息传递与消息聚合联合操作
  

• update方法的覆写

  class GCNConv(MessagePassing):
      def __init__(self, in_channels, out_channels):
          # 省略···
  
      def forward(self, x, edge_index):
          # 省略···
  
          # Step 4-5: Messages Passing
          return self.propagate(edge_idex, x=x, norm=norm, deg=deg.view((-1, 1)))
  
      def message(self, x_j, norm, deg_i):
          # x_j : [E, out_channels]
          # deg_i : [E, 1]
          return norm.view(-1, 1) * x_j * deg_i
  
      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)
  
      def update(self, inputs, deg):
          print(deg)
          return inputs
  

作业2

MessagePassing基类的运行流程：

• 检查子类是否实现了message_and_aggregate()方法，并将检查结果赋值给fuse属性：True OR False；
• 调用propagate方法，开始执行“消息传递过程”；
• propagate()方法首先检查edge_index是否为SparseTensor类型以及子类实现了message_and_aggregate()方法：
  • 是，依次执行子类的message_and_aggregate(),update()方法；
  • 否，依次执行子类的message(),aggregate(),update()方法。

请复现一个一层的图神经网络的构造，总结通过继承MessagePassing基类来构造自己的图神经网络类的规范。

• TODO