引言
在开篇中我们介绍了,为节点生成节点表征(Node Representation)是图计算任务成功的关键,我们要利用神经网络来学习节点表征。
消息传递范式是一种聚合邻接节点信息来更新中心节点信息的范式,它将卷积算子推广到了不规则数据领域,实现了图与神经网络的连接。消息传递范式因为简单、强大的特性,于是被人们广泛地使用。遵循消息传递范式的图神经网络被称为消息传递图神经网络。本节中,
- 首先我们将学习图神经网络生成节点表征的范式–消息传递(Message Passing)范式。
- 接着我们将初步分析PyG中的
MessagePassing
基类,通过继承此基类我们可以方便地构造一个图神经网络。 - 然后我们以继承
MessagePassing
基类的GCNConv
类为例,学习如何通过继承MessagePassing
基类来构造图神经网络。 - 再接着我们将对
MessagePassing
基类进行剖析。 - 最后我们将学习在继承
MessagePassing
基类的子类中覆写message(),aggreate(),message_and_aggreate()和update(),这些方法的规范。
消息传递范式介绍
下方图片展示了基于消息传递范式的聚合邻接节点信息来更新中心节点信息的过程:
- 图中黄色方框部分展示的是一次邻接节点信息传递到中心节点的过程:B节点的邻接节点(A,C)的信息经过变换后聚合到B节点,接着B节点信息与邻接节点聚合信息一起经过变换得到B节点的新的节点信息。同时,分别如红色和绿色方框部分所示,遵循同样的过程,C、D节点的信息也被更新。实际上,同样的过程在所有节点上都进行了一遍,所有节点的信息都更新了一遍。
- 这样的“邻接节点信息传递到中心节点的过程”会进行多次。如图中蓝色方框部分所示,A节点的邻接节点(B,C,D)的已经发生过一次更新的节点信息,经过变换、聚合、再变换产生了A节点第二次更新的节点信息。多次更新后的节点信息就作为节点表征。
消息传递图神经网络遵循上述的“聚合邻接节点信息来更新中心节点信息的过程”,来生成节点表征。用
x
(
k
−
1
)
i
∈
R
F
\mathbf{x}^{(k-1)}i\in\mathbb{R}^F
x(k−1)i∈RF表示
(
k
−
1
)
(k-1)
(k−1)层中节点
i
i
i的节点表征,
e
j
,
i
∈
R
D
\mathbf{e}{j,i} \in \mathbb{R}^D
ej,i∈RD 表示从节点
j
j
j到节点
i
i
i的边的属性,消息传递图神经网络可以描述为
x
i
(
k
)
=
γ
(
k
)
(
x
i
(
k
−
1
)
,
□
j
∈
N
(
i
)
ϕ
(
k
)
(
x
i
(
k
−
1
)
,
x
j
(
k
−
1
)
,
e
j
,
i
)
)
,
\mathbf{x}_i^{(k)} = \gamma^{(k)} \left( \mathbf{x}_i^{(k-1)}, \square_{j \in \mathcal{N}(i)} \phi^{(k)}\left(\mathbf{x}_i^{(k-1)}, \mathbf{x}_j^{(k-1)},\mathbf{e}_{j,i}\right) \right),
xi(k)=γ(k)(xi(k−1),□j∈N(i)ϕ(k)(xi(k−1),xj(k−1),ej,i)), 其中
□
\square
□表示可微分的、具有**排列不变性(函数输出结果与输入参数的排列无关)**的函数。具有排列不变性的函数有,sum()函数、mean()函数和max()函数。
γ
\gamma
γ和
ϕ
\phi
ϕ表示可微分的函数,如MLPs(多层感知器)。
注(1):神经网络的生成节点表征的操作称为节点嵌入(Node Embedding),节点表征也可以称为节点嵌入。为了统一此次组队学习中的表述,我们规定节点嵌入只代指神经网络生成节点表征的操作。
注(2):未经过训练的图神经网络生成的节点表征还不是好的节点表征,好的节点表征可用于衡量节点之间的相似性。通过监督学习对图神经网络做很好的训练,图神经网络才可以生成好的节点表征。我们将在第5节介绍此部分内容。
注(3),节点表征与节点属性的区分:遵循被广泛使用的约定,此次组队学习我们也约定,节点属性data.x是节点的第0层节点表征,第 h h h层的节点表征经过一次的节点间信息传递产生第 h + 1 h+1 h+1层的节点表征。不过,节点属性不单指data.x,广义上它就指节点的属性,如节点的度等。
MessagePassing基类初步分析
Pytorch Geometric(PyG)提供了MessagePassing
基类,它封装了“消息传递”的运行流程。通过继承MessagePassing
基类,可以方便地构造消息传递图神经网络。构造一个最简单的消息传递图神经网络类,我们只需定义message()方法(
ϕ
\phi
ϕ)、update()方法(
γ
\gamma
γ),以及使用的消息聚合方案(aggr=“add”、aggr="mean"或aggr=“max”)。这一切是在以下方法的帮助下完成的:
-
MessagePassing(aggr="add", flow="source_to_target", node_dim=-2)
(对象初始化方法):
aggr:定义要使用的聚合方案(“add”、"mean "或 “max”);
flow:定义消息传递的流向("source_to_target "或 “target_to_source”);
node_dim:定义沿着哪个维度传播,默认值为**-2**,也就是节点表征张量(Tensor)的哪一个维度是节点维度。节点表征张量x形状为[num_nodes, num_features],其第0维度(也是第-2维度)是节点维度,其第1维度(也是第-1维度)是节点表征维度,所以我们可以设置node_dim=-2。
注:MessagePassing(……)
等同于MessagePassing.__init__(……)
-
MessagePassing.propagate(edge_index, size=None, **kwargs)
:
开始传递消息的起始调用,在此方法中message、update等方法被调用
。
它以edge_index(边的端点的索引)和flow(消息的流向)以及一些额外的数据为参数。
请注意,propagate()不局限于基于形状为[N, N]的对称邻接矩阵进行“消息传递过程”。基于非对称的邻接矩阵进行消息传递(当图为二部图时),需要传递参数size=(N, M)。
如果设置size=None,则认为邻接矩阵是对称的。 -
MessagePassing.message(...)
:
首先确定要给节点 i i i传递消息的边的集合:
如果flow=“source_to_target”,则是 ( j , i ) ∈ E (j,i) \in \mathcal{E} (j,i)∈E的边的集合;
如果flow=“target_to_source”,则是 ( i , j ) ∈ E (i,j) \in \mathcal{E} (i,j)∈E的边的集合。
接着为各条边创建要传递给节点 i i i的消息,即实现 ϕ \phi ϕ函数。
MessagePassing.message(…)方法可以接收传递给MessagePassing.propagate(edge_index, size=None, **kwargs)方法的所有参数,我们在message()方法的参数列表里定义要接收的参数,例如我们要接收x,y,z参数,则我们应定义message(x,y,z)方法。
传递给propagate()方法的参数,如果是节点的属性的话,可以被拆分成属于中心节点的部分和属于邻接节点的部分,只需在变量名后面加上_i或_j。例如,我们自己定义的message方法包含参数x_i,那么首先propagate()方法将节点表征拆分成中心节点表征和邻接节点表征,接着propagate()方法调用message方法并传递中心节点表征给参数x_i
。而如果我们自己定义的meassage方法包含参数x_j,那么propagate()方法会传递邻接节点表征给参数x_j
。
我们用 i i i表示“消息传递”中的中心节点,用 j j j表示“消息传递”中的邻接节点。 -
MessagePassing.aggregate(...)
:
将从源节点传递过来的消息聚合在目标节点上,一般可选的聚合方式有sum, mean和max。 -
MessagePassing.message_and_aggregate(...)
:
在一些场景里,邻接节点信息变换和邻接节点信息聚合这两项操作可以融合在一起,那么我们可以在此方法里定义这两项操作,从而让程序运行更加高效。 -
MessagePassing.update(aggr_out, ...)
:
为每个节点 i ∈ V i \in \mathcal{V} i∈V更新节点表征,即实现 γ \gamma γ函数。此方法以aggregate方法的输出为第一个参数,并接收所有传递给propagate()方法的参数。
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)), 其中,邻接节点的表征
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)进行归一化处理,最后进行求和。这个公式可以分为以下几个步骤:
- 向邻接矩阵添加自环边。
- 对节点表征做线性转换。
- 计算归一化系数。
- 归一化邻接节点的节点表征。
- 将相邻节点表征相加("求和 "聚合)。
步骤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)。
edge_index 是一个形状为 [2, E] 的张量,其中 E 是边的数量。
row = edge_index[0] 表示边的源节点索引。
col = edge_index[1] 表示边的目标节点索引。
简写就是 row, col = edge_index
通过 row 和 col,可以获取每条边连接的两个节点,从而访问节点的度数信息。
deg = degree(col, x.size(0), dtype=x.dtype)
在 PyTorch Geometric 中,torch_geometric.utils.degree
是一个用于计算节点度数的工具函数。
它的输入是图中每条边的目标节点索引,输出是每个节点的度数。
通过 degree(col, x.size(0)),计算每个节点的度数。col 表示每条边的目标节点,因此按 col 统计得到每个节点的入度。x.size(0)
:表示图中节点的总数 N。表示张量 x 在第 0 维度(即行数)上的大小。在图神经网络(GCN)代码中,x.size(0) 通常用于获取节点的数量,因为 x 是节点特征矩阵。这个值用于指定度数数组的长度,确保每个节点(包括孤立节点)的度数都能被正确计算。dtype=x.dtype
:确保返回的度数张量 deg 的数据类型与输入特征张量 x 的类型一致。
deg 是一个长度为节点数 N 的张量,表示每个节点的度数。
norm = deg_inv_sqrt[row] * deg_inv_sqrt[col]
通过 row 和 col 索引得到相应节点的度数倒数平方根,并将其相乘。最终得到的 norm 是一个长度为 E(边的个数) 的张量。
该 norm 表示每条边的归一化因子,用于后续消息传递过程中的特征归一化。
在message()
方法中,我们需要通过norm对邻接节点表征x_j进行归一化处理。
x_j 是从邻接节点传来的特征,形状为 [E, out_channels],其中 E 是边的数量。
norm 是每条边的归一化因子,形状为 [E]。
-
归一化因子调整形状:
norm.view(-1, 1)
norm 是形状 [E] 的张量。使用 .view(-1, 1) 将其调整为 [E, 1] 的形状,使其可以广播(broadcast)到 x_j 的形状 [E, out_channels]。 -
归一化特征:
norm.view(-1, 1) * x_j
逐元素乘法,将归一化因子 norm 应用于每条边的特征 x_j。这一步相当于将邻接矩阵的归一化权重应用于节点特征,确保每条边的信息根据归一化因子缩放。
通过以上内容的学习,我们便掌握了创建一个仅包含一次“消息传递过程”的图神经网络的方法。如下方代码所示,我们可以很方便地初始化和调用它:
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)
通过串联多个这样的简单图神经网络,我们就可以构造复杂的图神经网络模型。我们将在第5节介绍复杂图神经网络模型的构建。
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