DGL学习笔记——第二章消息传递范式

最新推荐文章于 2024-10-01 14:44:54 发布

斯外戈的小白

最新推荐文章于 2024-10-01 14:44:54 发布

阅读量744

点赞数

文章标签：学习人工智能 python

本文链接：https://blog.youkuaiyun.com/weixin_51589123/article/details/127989000

版权

提示：DGL用户指南学习中

文章目录

一、内置函数和消息传递API
二、编写高效的消息传递代码
总结

消息传递是实现GNN的一种通用框架和编程范式。它从聚合与更新的角度归纳总结了多种GNN模型的实现。

假设节点 𝑣 上的的特征为 𝑥𝑣∈ℝ𝑑1，边 (𝑢,𝑣) 上的特征为 𝑤𝑒∈ℝ𝑑2。消息传递范式定义了以下逐节点和边上的计算：

请添加图片描述

在上面的等式中， 𝜙 是定义在每条边上的消息函数，它通过将边上特征与其两端节点的特征相结合来生成消息。聚合函数 𝜌 会聚合节点接受到的消息。更新函数 𝜓 会结合聚合后的消息和节点本身的特征来更新节点的特征。

一、内置函数和消息传递API

在DGL中，消息函数接受一个参数 edges，这是一个 EdgeBatch 的实例，在消息传递时，它被DGL在内部生成以表示一批边。 edges 有 src、 dst 和 data 共3个成员属性，分别用于访问源节点、目标节点和边的特征。

聚合函数接受一个参数 nodes，这是一个 NodeBatch 的实例，在消息传递时，它被DGL在内部生成以表示一批节点。 nodes 的成员属性 mailbox 可以用来访问节点收到的消息。一些最常见的聚合操作包括 sum、max、min 等。

更新函数接受一个如上所述的参数 nodes。此函数对聚合函数的聚合结果进行操作，通常在消息传递的最后一步将其与节点的特征相结合，并将输出作为节点的新特征。

DGL在命名空间 dgl.function 中实现了常用的消息函数和聚合函数作为内置函数。一般来说，DGL建议尽可能使用内置函数，因为它们经过了大量优化，并且可以自动处理维度广播。

如果用户的消息传递函数无法用内置函数实现，则可以实现自己的消息或聚合函数(也称为用户定义函数 )。

内置消息函数可以是一元函数或二元函数。对于一元函数，DGL支持 copy 函数。对于二元函数， DGL现在支持 add、 sub、 mul、 div、 dot 函数。消息的内置函数的命名约定是 u 表示源节点， v 表示目标节点，e 表示边。这些函数的参数是字符串，指示相应节点和边的输入和输出特征字段名。关于内置函数的列表，请参见 DGL Built-in Function。例如，要对源节点的 hu 特征和目标节点的 hv 特征求和，然后将结果保存在边的 he 特征上，用户可以使用内置函数 dgl.function.u_add_v(‘hu’, ‘hv’, ‘he’)。而以下用户定义消息函数与此内置函数等价。

请添加图片描述

DGL支持内置的聚合函数 sum、 max、 min 和 mean 操作。聚合函数通常有两个参数，它们的类型都是字符串。一个用于指定 mailbox 中的字段名，一个用于指示目标节点特征的字段名，例如，dgl.function.sum(‘m’, ‘h’) 等价于如下所示的对接收到消息求和的用户定义函数：

请添加图片描述

在DGL中，也可以在不涉及消息传递的情况下，通过 apply_edges() 单独调用逐边计算。 apply_edges() 的参数是一个消息函数。并且在默认情况下，这个接口将更新所有的边。例如：
请添加图片描述
对于消息传递， update_all() 是一个高级API。它在单个API调用里合并了消息生成、消息聚合和节点特征更新，这为从整体上进行系统优化提供了空间。