5.5 使用边权重
在带权图中,每条边都关联着一个具有语义意义的标量权重。例如,边权重可以是连接强度或置信度分数。自然地,人们可能希望在模型开发中利用边权重。
带边权重的消息传递
大多数图神经网络 (GNN) 通过且仅通过消息传递机制在前向计算中整合图拓扑信息。消息传递操作可以看作一个以邻接矩阵和附加输入特征作为输入参数的函数。对于无权图,邻接矩阵中的条目可以是零或一,其中取一的条目表示一条边。如果图是带权的,非零条目可以取任意标量值。这等效于将每条消息乘以其对应的边权重,如 GAT 中所述。
使用 DGL,可以通过以下方式实现:
将边权重保存为边特征
在消息函数中将原始消息乘以边特征
考虑下面使用 DGL 进行消息传递的示例。
import dgl.function as fn
# Suppose graph.ndata['ft'] stores the input node features
graph.update_all(fn.copy_u('ft', 'm'), fn.sum('m', 'ft'))
可以修改它以支持边权重,如下所示。
import dgl.function as fn
# Save edge weights as an edge feature, which is a tensor of shape (E, *)
# E is the number of edges
graph.edata['w'] = eweight
# Suppose graph.ndata['ft'] stores the input node features
graph.update_all(fn.u_mul_e('ft', 'w', 'm'), fn.sum('m', 'ft'))
使用带边权重的 NN 模块
可以通过修改 NN 模块中的所有消息传递操作来使其支持边权重。以下代码片段是一个支持边权重的 NN 模块示例。
如果 DGL 的内置 NN 模块在前向函数中接受一个可选的 edge_weight
参数,则它们支持边权重。
可能需要对原始边权重进行归一化。在这方面,DGL 提供了 EdgeWeightNorm()
。