消息传递是实现GNN
的一种通用框架和编程范式。它从聚合与更新的角度归纳总结了多种GNN
模型的实现,它的思路是:
u
相连的边上的信息函数聚合起来,并结合u
的已有节点特征,来更新v的节点特征。
消息传递的数学公式如下:
m e i t + 1 = Φ ( x v i t , x u t , ω e i t ) x u t + 1 = Ψ ( x u t , ρ ( { m e i t + 1 , i ∈ N ( u ) } ) \begin{aligned} m_{e_i}^{t+1} &= \Phi(x_{v_i}^t, x_u^t, \omega_{e_i}^t)\\ x_u^{t+1} &= \Psi(x_u^t, \rho(\{m_{e_i}^{t+1},i \in N(u)\}) \end{aligned} meit+1xut+1=Φ(xvit,xut,ωeit)=Ψ(xut,ρ({meit+1,i∈N(u)})
聚合
函数,把所有的边的消息聚合起来处理,通常是可微分、具有排列不变性的函数(即和 m e i t + 1 m_{e_i}^{t+1} meit+1的顺序无关);更新
函数,结合聚合后的消息和节点本身的特征来更新节点的特征;\quad
Pytorch Geometric(PyG)提供了MessagePassing
基类,它封装了“消息传递”的运行流程。通过继承MessagePassing
基类,可以方便地构造消息传递图神经网络。而关键地就是如上文所说,定义消息函数、聚合函数和更新函数,下面先介绍一下MP基类和该类的一些方法。
\quad
MessagePassing(arrr = "add", flow = "source_to_target", node_dim = -2)
MessagePassing.message()
MessagePassing.propagate(edge_index, size: Size = None, **kwargs)
message
、update
等方法被调用;MessagePassing.aggregate(...)
MessagePassing.update(aggr_out, ...)
aggregate
方法的输出为第一个参数,并接收所有传递给propagate()
方法的参数。函数调用流程如下(不执行mp.message_and_aggregate
的情况下):
\quad
我们以继承MessagePassing
基类的GCNConv
类为例,学习如何通过继承MessagePassing
基类来实现一个简单的图神经网络。
消息函数:
m e j t + 1 = Φ ( x v j t , x u t , ω e j t ) = 1 d e g ( j ) ⋅ 1 d e g ( u ) Θ ⋅ x j t m_{e_j}^{t+1} = \Phi(x_{v_j}^t, x_u^t, \omega_{e_j}^t) = \frac{1}{\sqrt{deg(j)}}·\frac{1}{\sqrt{deg(u)}}\Theta·x_j^t mejt+1=Φ(xvjt,xut,ωejt)=deg(j)1⋅deg(u)1Θ⋅xjt
聚合函数
ρ = ∑ j ∈ N ( u ) m e j t + 1 \quad \rho = \sum_{j\in N(u)}m_{e_j}^{t+1} ρ=j∈N(u)∑mejt+1
更新函数:
Ψ ( ρ ) = ρ \Psi(\rho)=\rho Ψ(ρ)=ρ
由上可知:
x u t + 1 = ∑ j ∈ N ( u ) 1 d e g ( j ) ⋅ 1 d e g ( u ) Θ ⋅ x j t x_u^{t+1} = \sum_{j\in N(u)}\frac{1}{\sqrt{deg(j)}}·\frac{1}{\sqrt{deg(u)}}\Theta·x_j^t xut+1=j∈N(u)∑deg(j)1⋅deg(u)1Θ⋅xjt
GCNConv
实现步骤主要通过torch_geometric.utils.add_self_loops
方法实现。
def add_self_loops(edge_index, edge_weight: Optional[torch.Tensor] = None,
fill_value: float = 1., num_nodes: Optional[int] = None):
pass
return edge_index, edge_weight
edge_index
中;edge_weight
不是None,加入的自环边权重补充为1,即fill_value
;edge_index
和edge_weight
;num_nodes
即节点数量。这一步就对应了 Θ ⋅ x j t \Theta·x_j^t Θ⋅xjt,主要通过一个线性层torch.nn.Linear
实现。
\quad
也就是算 d i j = 1 d e g ( i ) ⋅ 1 d e g ( j ) d_{ij} = \frac{1}{\sqrt{deg(i)}}·\frac{1}{\sqrt{deg(j)}} dij=deg(i)1⋅deg(j)1,然后得到一个张量norm
,里面元素的个数和扩充过得edge_index
中边的个数相等,且对应,即:若edge_index
的第k个边为 ( i , j ) (i,j) (i,j),norm
中第k个为 d i j d_{ij} dij。
而对于节点的度可以使用torch_geometric.utils.degree
获得
degree(index, num_nodes: Optional[int] = None, dtype: Optional[int] = None)
对于节点的次,上述函数是通过统计index里面的各节点的个数得到的。
在GCNconv
的实现过程中,这一部分的代码如下:
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]
row
和col
没有区别;row
呢,是初始节点,这时计算出的是出次,而对于col
算出来的即入次,两者往往不相同。实例,对于如下的图,我们分别使用row
和col
计算度数,看看有什么区别:
import torch
from torch_geometric.data import Data
from torch_geometric.utils import degree,add_self_loops
print(edge_index)
edge_index2,_ = add_self_loops(edge_index,num_nodes = 5)
print(edge_index2)
'''
tensor([[0, 4, 3, 1, 1, 0],
[4, 3, 2, 2, 0, 1]])
tensor([[0, 4, 3, 1, 1, 0, 0, 1, 2, 3, 4],
[4, 3, 2, 2, 0, 1, 0, 1, 2, 3, 4]])
'''
#下面分别使用row和col计算度
r2,c2 = edge_index2
print(r2)
print(c2)
'''
tensor([0, 4, 3, 1, 1, 0, 0, 1, 2, 3, 4])
tensor([4, 3, 2, 2, 0, 1, 0, 1, 2, 3, 4])
'''
d_out = degree(r2,5)
print("d_out:",d_out)
d_in = degree(c2,5)
print("d_in :",d_in)
'''
d_out: tensor([3., 3., 1., 2., 2.])
d_in : tensor([2., 2., 3., 2., 2.])
'''
\quad
也就是把上一步norm
中的 d i j d_{ij} dij(相当于j节点的权重了)和对应的 Θ ⋅ x j t \Theta·x_j^t Θ⋅xjt乘起来就搞定了。
\quad
这一步就是聚合函数了,把上一步中的全部加起来就好了。
\quad
1、【深度学习实战】Pytorch Geometric实践——利用Pytorch搭建GNN
2、消息传递图神经网络