@[TOC](消息传递图神经网络)
消息传递范式
最右图为消息传递范式
Pytorch Geometric(PyG)中的MessagePassing基类
- 它封装了“消息传递”的运行流程。通过继承MessagePassing基类,可以方便地构造消息传递图神经网络。
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)),
其中,邻接节点的表征
x
j
(
k
−
1
)
\mathbf{x}_j^{(k-1)}
xj(k−1)首先通过与权重矩阵
Θ
\mathbf{\Theta}
Θ相乘进行变换,然后按端点的度
deg
(
i
)
,
deg
(
j
)
\deg(i), \deg(j)
deg(i),deg(j)进行归一化处理,最后进行求和。这个公式可以分为以下几个步骤:
- 向邻接矩阵添加自环边。
- 对节点表征做线性转换。
- 计算归一化系数。
- 归一化邻接节点的节点表征。
- 将相邻节点表征相加("求和 "聚合)。
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
- 所有逻辑在forward()方法中实现。
- 第一步通过
torch_geometric.utils.add_self_loops()
给边索引添加自循环边。 - 第二步通过
torch.nn.Linear
线性变换 propagate()
在forward()中调用- 在’MessagePassing’基类中若
message_and_aggregate()
被实现则同时出现的message
方法和aggregate
方法不会被执行。 - 对于
message
方法,如果是节点的属性的话,可以被拆分成属于中心节点的部分和属于邻接节点的部分,只需在变量名后面加上_i
或_j
。 update
方法接收聚合的输出作为第一个参数,此外还可以接收传递给propagate
方法的任何参数。
参考资料:
https://github.com/datawhalechina/team-learning-nlp/tree