Press "Enter" to skip to content

王敏捷:亚马逊云科技图深度学习系统DGL介绍和展望

本站内容均来自兴趣收集,如不慎侵害的您的相关权益,请留言告知,我们将尽快删除.谢谢.

 

 

分享嘉宾 : 王敏捷 博士  亚马逊云科技

 

编辑整理 :郭艳 中山大学

 

出品平台:DataFunTalk

 

导读: 图作为一个非常普适的数据结构存在于很多场景,随之而来的图神经网络也被用来挖掘图上有用的信息以增进下游任务的表现。 但是根据现有工具,图神经网络的 细粒度公式表达和粗粒度的张量实现存在着巨大的鸿沟。 由此,我们构建了一座图算法与现有张量计算框架之间的桥梁:D GL,帮助用户更容易 在现有D L框架上高效的执行相应的图算法。

 

本文将围绕下面几点进行展开:

 

图 神经网络背景介绍

 

D GL 的提出背景

 

DGL的特点

 

01

 

图神经网络背景介绍

 

 

图作为一个非常powerful的数据结构存在于很多场景,例如微观层面上的药物和分子结构,宏观层面的社交网络、用户产品交互,以及知识图谱等等。有了这幺多的图数据结构,相应地,我们在机器学习领域有非常多的topic,其目的在于如何在图上挖掘有用的信息,来增进下游的任务。

 

下面举两个例子:

 

 

社交平台上的水军检测

 

在社交网络平台有一个典型的任务交水军检测,具体来说,我们根据社交网络上的用户和帖子的交互关系来抽象出一个异构图,其中用户是圆圈,帖子是方块,如果用户和帖子之间发生了一些交互关系,例如用户点赞、转发帖子,则他们之间有一条边。而水军检测的任务就是每当用户和帖子之间发生一个新的交互时,判断这个交互行为是恶意还是善意。可以看出,水军检测最终规约到机器学习中的边分类问题,即给定一个图以及一条边,判断边的标签是什幺。

 

 

药物重定位知识图谱

 

另外一个很经典的任务,即药物重定位知识图谱。由于药物的研发流程非常漫长,我们需要探究现在已有的药物是否能针对于新的病毒(例如cov’19)产生作用,来加快药物的研发。为此,我们构建了一个知识图谱——药物重定位知识图谱(DRKG)。把很多不同的数据集结和起来,变成一个比较复杂的知识图谱结构,上图中左边是知识图谱的schema。通过知识图谱,我们把药物重定位问题归约成了知识图谱的补全问题,即若发现一个新的病毒新的蛋白和一个药物的蛋白之间,可能有个missing-link,而missing-link的关系,恰好有治疗的潜在作用。那我们就有理由相信,已有的药物可能是新病毒的潜在的治疗药物。由此我们通过图上机器学习的方法,实现了药物重定位问题。

 

 

可以看出,图神经网络的应用场景非常广,同时它也是目前非常火的一个topic,每年关于图神经网络的论文数目的增长幅度非常巨大,同时在顶尖会议上,graph neural nerwork也是非常火的关键词之一。

 

 

什幺是图神经网络?图神经网络是以图为数据,学习图的点、边或者整张图的向量表示,以进行下游任务的一类深度神经网络。

 

例如判断用户的社区喜好问题中,图神经网络处理上的核心观点是收集这个用户邻居的信息,以此来得到该用户的社区。这也是一个很朴素的思想——消息传递,即在计算图神经网络的过程中,不仅要看节点本身的信息,还要看节点邻居/拓扑的信息,也叫context。

 

 

消息传递用数学公式来理解主要包括两类公式:

 

边上的函数——消息函数

 

消息函数正如上图中第一个公式所示,   是节点 v  的邻居三元组   的集合,其中参数 分别代表边 e  上源节点 u  和目标节点 v  的向量表征, 是边 e  上的权值表征。通过消息函数,我们得到了消息对象 , 也是一个向量。

 

点上的函数——更新函数和累和函数

 

其中累和函数 将节点 v  邻居的消息对象进行累加,得到累加后的消息对象,而更新函数 结合累加后的消息对象与节点 v  本身的信息,得到节点 v  新一步的向量表征 。

 

可以看出,图神经网络非常灵活,因为其消息函数、更新函数以及累和函数都可以定制化,甚至可以在这些函数中设置参数, 使 其本身变成一个神经网络模块。 因此,图神经网络的研究非常热,因为使用者可以根据不同的场景来设计相应的消息函数、更新函数和累 和 函数,以达到不一样的效果。

 

 

以图卷积网络GCN为例,在GCN论文中以矩阵的形式来表达(上图中第一个公式), 分别是图的度数矩阵和邻接矩阵, 是节点的向量表征,是一个稠密矩阵, 是一个可学习权重,再经过非线性函数 ,就是GCN的矩阵形式。

 

而矩阵形式和消息传递模式是一一对应的,通过将GCN的矩阵形式转写成消息传递公式即得到上图中第二个公式。消息函数通过将节点 i 的每一个邻居 j  投影到一个低维度空间后再乘上正则化系数    得到对应的消息对象,而累和函数把这些消息对象进行累和,再通过非线性函数 完成了对节点 i  表征的更新。以上就是整个图卷积网络消息传递的表示。

 

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

 

02

 

DGL的提出背景

 

 

理解了图神经网络基于消息传递的抽象,而图神经网络同时也是一类新的神经网络,那幺能否通过现有的深度学习框架来实现消息传递模式呢?

 

答案是非常难实现。因为消息传递其实是一种细粒度的计算,其定义的是在每条边和每个点上进行了哪些计算。例如消息函数定义的是在一条边上,如何计算和发送消息,在点上,聚合函数和更新函数表达的是如何聚合和使用所收到的信息。而常见的深度学习框架,例如tensorflow和pytorch等,它们都是张量的编程接口,需要一种比较粗粒度的思考方式。而两种粒度的思考方式是有鸿沟的,正是因为这个鸿沟,导致用现有的神经网络框架编写图神经网络比较困难。

 

 

看到了这个鸿沟,三年前在图神经网络刚刚兴起时,我们开发了一个新的框架——DGL(DEEP GRAPH LIBRARY)。一句话来说,DGL是一座图算法与现有的张量计算框架之间的桥梁,帮助用户更容易在现有DL框架的基础上高效地执行相应的图算法。

 

03

 

DGL的特点

 

 

而DGL的特点主要包括以下四个方面:

 

灵活易用的编程接口

 

 

本着以“图”为本的思想,在DGL中,图是程序里的“一等公民”,所有的DGL函数和NN模块都可以接受和返回图对象。例如前文提到的消息传递非常重要的两个公式,在DGL中也包装成了图对象的成员函数,边上和点上的函数分别对应成员函数g.apply_edges()和g.update_all()。两个成员函数中的参数是前文提到的消息函数 、累和函数 以及更新函数 。 虽然只提供了两个API,但是这两个API其实是非常灵活的,它可以支持任意消息传递的计算。

 

 

下面举一个较为简单的例子来讲如何使用DGL。DGL包装了很多high level的例如GCN的模块,用户可以使用nn.module模块来构建图神经网络。如上图,我们用DGL中的GraphConv模块定义了两层图卷积网络。在forward中,本着以“图”为本的思想,以一个graph对象和节点的原始向量in_feature为输入。而在forward函数的GraphCov模块中,以graph对象和节点的向量表征作为参数,通过图卷积的计算获得了节点新的向量表征。在经过一个例如PyTorch的激活函数relu和另外一个图卷积模块后,完成了整个计算。

 

而要实现一个整个的深度学习网络训练的pipline,DGL不仅仅提供刚才提到的GNN的NN模块,还提供了在数据处理、对loss的处理以及在图上做argumentation等环节上需要用到的一些模块以及一些基于图上的API。譬如构图算法、图上的增删改查,和对于巨图来说非常重要的子图抽取和采样,还有图遍历、图变换算法以及图和各类数据格式的转换等。

 

除了在API的数量上DGL提供了更多的coverage,让用户使用起来更加方便外,DGL对这些API的灵活性上也做了很多设计。例如这些API都支持异构图,即支持图有不同类型的节点和不同类型的边,同时大多数的API都支持GPU加速,以及DGL还是目前唯一支持多平台(PyTorch,TensorFlow,MXNet)的开源图机器学习框架。

 

 

知道了如何用D GL 的N N 模块搭建神经网络框架,那用户如何定制化自己的图卷积层呢 ? D GL 也提供了比较方便的消息传递函数。 以定义经典的Graph SAGE 举例,如上一个例子,forward函数同样以g raph 对象和节点的原向量表征为输入。 通过 ndata,edata 来访问节点和边的特征数据,然后将节点的数据与图做一个绑定,再计算消息传递。 消息函数直接把原节点copy过来作为消息对象,累和函数中采取mean函数,即对收到的消息做平均,最后调用update _all 函数触发消息传递,触发后,更新过的消息对象存在g .ndata 中,再把其取出来进行后续的操作。

 

以上就是关于DGL接口方面的设计,那如何高效地实现刚刚提到的消息传递呢?

 

高效的底层系统设计

 

 

消息传递计算是非常重要的。我们根据不同图和dataset,会发现在很多情况下,消息传递计算update_all计算可以达到整个GNN训练的50%以上。

 

而传统的实现方式是基于gather/scatter原语。什幺是gather/scatter?在把消息从一个节点a传递到节点b时, gather的意思先把a的feature,即节点a的消息,gather到a所有的出边上,再在这个边上进行消息的计算,后续通过一个scatter的原语,把计算完的消息累加到节点b的buffer里去。但是gather/scatter模式会在计算过程当中产生大量的冗余消息对象,占用了大量的内存带宽。因为在gather/scatter模式中,消息对象的数目和图的边的数目相关联,原始节点的feature的大小和节点数目相关联,而对于一个图来说,其边的数目往往是远大于节点的数目,所以这就产生了相应的带宽问题。

 

为了使得消息传递能更高效地实现,DGL首先提出了算子融合技术,对消息传递进行优化。

 

 

再次以图卷积网络举例,将其消息传递函数简化成两个公式,点上的和边上的,可以看出消息函数就是直接取原节点的feature,而累和函数就是将消息对象直接进行相应的累加。而这一模式对应一个算子,叫稀疏-稠密矩阵乘法算子(SpMM)。稀疏矩阵A(邻接矩阵)和稠密矩阵 X(节点向量表征矩阵)相乘,得到的结果正好是整个消息传递计算得到的结果。而算子相对于gather/scatter模式来说,不会在计算过程中产生中间结果,而是直接算出累加过的结果,因此它是非常高效的。

 

 

同样地,我们举另外一个例子——图注意力网络。图注意网络和刚刚提到图卷积网络区别在于边上要计算注意力权重,计算注意力权重可以看作消息函数,而注意力权重是边上源节点和目标节点的内积。而注意力的消息传递也可以用采样稠密矩阵乘法这一算子来进行高效的计算,采样稠密矩阵乘法是稠密乘法的一个拓展,两个稠密矩阵X和Y分别是图上源节点和目标节点的向量表征。采样稠密矩阵乘法是对X和Y^T的矩阵乘法通过A进行稀疏化。由于避免了将点数据拷贝到边上,采样-稠密矩阵乘法算子也是非常高效的。

 

 

在DGL当中,对于消息传递模式的实现方式就是刚刚提到的两个算子。同时DGL也在两个方面通过扩展SpMM和SDDMM算子来支持更高效的消息传递计算。第一是传统的两个算子只有乘加操作,和在DGL中,两个算子更加灵活一些,例如可以支持max,min求最大值和最小值等累和操作。第二方面,DGL将传统SpMM和SDDMM算子标量的计算单元拓展成允许张量计算单元。

 

同时,我们也在论文中证明,只要一个图神经网络模型可以用消息传递框架表达,那幺它向前传播和向后传播SpMM和SDDMM两类算子。这对DGL的系统设计来说是一个非常好的抽象,因为对底层优化来说,只要关注两类算子如何优化即可,而不用管上层算法的演进变化。

 

 

上图是一个简单的性能评测,在不同的数据集、不同的模型,CPU和GPU当中,与PyG进行了一个对比,通过刚刚提到的算子融合技术,DGL在CPU和GPU上都有很好的加速,同时能节省更多内存。

 

优秀的巨图训练性能

 

 

刚刚我们提到了再单GPU内核上计算的优化,而另一个非常重要的topic是如何处理巨图。目前,学术界也越来越关注如何处理巨图,例如非常有名的OGB(Open Graph Benchmark),它现在不断地在推出更大规模的图数据集,当然这个数据集是10亿左右,与工业界百亿甚至是千亿的图仍有差距。

 

 

那对于巨图来说,GNN是如何去训练的呢。在此之前,我们需要知道由于图数据中节点与其邻居存在相应的依赖关系,因此GNN的分布式训练或者并行训练与常见的 的数据例如图像和文本是有区别的,会更加复杂一些。而GNN的小批次训练主要基于子图采样的思想,包括以下几个步骤:

 

1、对目标节点随机地选取其部分的邻居节点,再选取邻居节点的邻居节点,不断地迭代

 

2、抽取第一步得到的边,得到一个子图

 

3、根据第二步中得到的子图,获取子图中的节点和边上的特征。

 

4、在子图上训练网络并更新参数

 

5、重复上面4步,直至收敛

 

可以看出图神经网络小批次训练的流程比图像和文本的流程要复杂很多,因此,在整个DGL框架的设计中,我们针对不同的环节设计了不同的模块,例如针对子图采样设计了Sampling Component,针对子图特征的抽取设计了一个KVStore Component,而在最后获得子图结构和子图的特征后的训练环节,我们使用的是传统的数据并行的方式Training Componet,即每个GPU获得一个子图进行训练,而不同GPU的参数通过参数服务器进行同步更新。

 

针对巨图,DGL也做了相应的一些优化以及相应的成果,如下图所示:

 

 

完善的开源生态

 

DGL是完全开源的,其项目主页和Github仓库地址如下所示:

 

 

 

在DGL上线之初,也得益于图神经网络的热度,DGL就获得了广泛的关注和好评。

 

 

同时,我们也深刻意识到,DGL不是整个开源生态的全部,而是开源生态的一个核心部分。我们希望有更多基于DGL的开源生态,来优化或实现某一类特定的任务。

 

下面我们主要介绍两个工作,一个工作是GNNLens。GNNLens是一个针对GNN的可视化工具。因为我们知道图是一个非常复杂的数据结构,而GNNLens就提供了在训练过程中,将图可视化的需求。大家也可以访问下图中的网址的demo,来看GNNLens是如何工作的。

 

 

demo

 

另外一个工作是由北京邮电大学石川教授的团队开发的,一个面向异质图神经网络的工具包,叫做OpenHGNN。OpenHGNN是一个很好的针对异质图的模型库,同时它也内置了AutoML进行参数自动调优。对异质图感兴趣的同学,可了解一下。

 

 

除了上面两个工作,我们的开源社区建设还包括以下方面:

 

 

同时也非常欢迎大家使用并贡献GDL,实习岗位也常年开放,有实习意向的同学也可以进行投递。

Be First to Comment

发表评论

您的电子邮箱地址不会被公开。