算法如何高效表达图计算?亚马逊云科技 DGL 图学习平台介绍
分享嘉宾:王敏捷博士 亚马逊云科技
编辑整理:宋佳林 中国人民公安大学
出品平台:DataFunTalk
导读: 本次分享题目为高效、易用、开放的图深度学习平台DGL介绍及展望,主要内容包括以下几大方面:
- 图数据和图神经网络
- DGL的优势
- 未来规划
01图数据和图神经网络
1. 图数据无处不在
我们为什么要研究图?从微观层面上看,药物分子结构是图数据结构,而在宏观层面上,像社交网络、用户产品交互网络、知识图谱等都是大规模图数据在工业级应用场景中的重要体现。因为图数据具有普适性,所以在图数据上进行机器学习任务非常重要的。
① 社交平台上的水军检测
以一个社交平台上做水军检测的案例为例,我们可以看到它是属于二部图的图类型,一边是表示用户的节点,另一边是表示帖子的节点,每当一个用户和帖子进行交互时,例如点赞转发等,就会在这个图上构建一条边。
水军检测是对每当有一条新的边即新的交互产生的时候,需要去判断这个交互究竟是一个善意交互,还是带有恶意交互,其实质是图上的边分类的问题。
② 药物重定位知识图谱(DRKG)
我们构建了一个全新的知识图谱,命名为DRKG。图中左边是知识图谱的一些schema,通过这个复杂的知识图谱,我们可以看到上面的不同节点类型和不同的边的类型,包括disease包括drug等。而药物重定位这个任务,是我们知道现在已有审批的这些药物下,当一个新的病毒产生的时候,如果我们想要再去研发新的药物,则需要花很久的时间。那么药物重定位的工作就是希望可以在现在已有审批的药物当中,发现哪些可以应对新的疾病。
为了解决这个问题,我们在知识图谱上进行挖掘,如果说我们发现一个药物的节点和一个疾病的节点之间很有可能是有边的,我们就有理由相信这个药物可能是潜在的应对病毒的治疗药物。
我们通过图数据挖掘发现,预测结果比较准确。在模型所推荐的41种药物当中,有11种已经被用于Conv-19的临床治疗,所以在图数据上进行挖掘,它有很多实际的应用和落地的场景。
2. 图神经网络
最近一个重要研究热点,就是将深度学习的这些算法和模型技巧应用在图数据当中,因此产生了一类新的神经网络,叫做图神经网络(Graph Neural Network)。ArXiv论文发表网站上,关于GNN论文的数量呈指数级增长,在ICLR2020机器学习的顶会中可以看到GNN是所有关键词当中最热门的话题,由此可见,不论是学界还是业界,对于图神经网络都非常关注,因此我们来了解一下什么是图神经网络。
① 什么是图神经网络
相信大家应该非常熟悉诸如卷积神经网络CNN、如transformer这些语言模型,其实核心的概念就是当给定一些数据,通过神经网络可以把数据映射到一个高维空间的向量表征,那么可以说这些向量表征包含了这些数据的这些内容,可以使用这些向量表征去做分类任务。
图神经网络本质是一样的,只是它的输入数据不同,输入数据是图数据,比如说有一个社交网络,它有四个用户之间的朋友关系连接,图神经网络可以对每个节点生成一个向量表征,然后使用向量表征可以输入到下游的分类去判断用户的喜好。
图神经网络的一个核心的理念是,当计算向量表征时,不仅要考虑节点本身的信息,也要考虑节点的拓扑信息,它所采用的方式是消息传递,也就是每当计算本节点消息向量表征的时候,同时也会去看该节点邻居的节点的信息,把它收集起来做聚合等变化,这样图神经网络计算出来的向量表征就可以包含不仅仅是本节点的信息,也有拓扑结构的信息。
② 图神经网络基于消息传递
我们可以用两个公式来表示消息传递的过程。
- 边上公式 :消息函数接受了三个参数,分别是边上的原节点,目标节点和边本身的向量表征,计算的结果就是在边上所要传递的消息对象。
- 点上公式 :第一个是累和函数的作用是去聚合收集到的原节点周围的发送过来的消息对象,接着与原节点表征进行结合,通过更新函数去计算这个节点的下一层的向量表征。
消息函数,累和函数和更新函数都可以根据场景进行定制化。
③ 图卷积网络
对于图卷积网络GCN公式的表达,我们看一下其中比较重要的概念:
- 矩阵形式中
邻接矩阵:A表示了图结构的信息,其中公式中+I是一个自环的处理
度数矩阵:D矩阵对角线上为每个节点的度
节点表示:H
可学习的权重:W
将这些矩阵相乘做一个Sigmoid非线性变换,就可以获得下一层所要计算的新节点的表示。
- 传递模式的公式中
消息函数:使用邻居节点J对于某一个节点I的邻居节点的表征hj乘上它的可行性权重W位
累和函数:相加累和
更新函数:更新向量表征
④ 图神经网络的编写
可见,消息传递模式是一个普适性的模式去描述很多种不同的图神经网络。但是很可惜的是,如果我们使用现在已有的深度学习框架,像Pytorch,Tensorflow去编写消息传递模型的话,其实是很困难且低效的,因为该算法和我们现在常见的深度学习算法有很多的不同。上面提到的消息传递计算,其实是定义了一种非常细粒度的计算,它讲的是在边上如何计算和发送消息,在点上如何聚合和使用消息,所以粒度非常细,但是如果我们去使用现在已有的这些计算框架去做,它往往提供的是张量的编程接口,就会粒度十分粗。
那么,为什么在现在的深度学习框架会做接口设计,是因为对于底层的硬件,框架希望表达更多的计算,使得能够更高效地并行。消息传递计算是一个细粒度的定义,但是接口是个粗粒度的计算,当中就有一个鸿沟,因此我们开发了DGL框架。
02 DGL的优势
如果用一句话来概括DGL的话,它就是一座桥梁,横跨在这个图和张量的领域里。我们希望通过让用户能够很直白的使用图的一些算法去表达图计算包括图神经网络的计算,然后DGL可以自动把其转换成比较高效张量的计算算子去进行加速和并行。
DGL有哪些特点?
第一, 灵活性 。我们拥有灵活应用的编程接口,其核心的理念是以图为本,让编程变得非常贴近图计算原生的语义。
第二, 高效的底层系统设计 。图计算在底层的张量硬件上可以并行,使用算子融合等技术对消息传递进行加速。
第三, 拥有巨图训练的性能 。目前DGL支持十亿级的巨图,同时能够高效地利用多机多GPU集群进行并行的计算。
第四,作为一款开源的软件,DGL具有 完善的开源生态 ,它和很多的开源软件有良好的互通性。
1. 灵活易用的编程接口
DGL核心的理念,以图为本。用一句话概括就是图是程序中的“一等公民”,也就是说DGL的所有的函数和NN模块都可以接收和返回图的对象,像非常核心的消息传递的API,它有边上的计算和点上的计算,对应的我们可以把它包装成图对象的内部的成员函数和对应的编程计算。
我们以编写图卷积网络为例,对于新手小白可能不太清楚卷积模块怎么去实现,这时使用DGL可以直接Import这些模块去轻松构建网络。DGL以图为本的概念,也体现在卷积模块的调用过程当中,比如说forward函数,可以看到它有两个参数,第一个参数是Graph,第二参数是节点的向量表征,其核心的概念就是图和张量同为核心的数据结构,所以它都成了GCN的forward参数。
通过这两个输入参数,GCN模块就会计算出节点的更新的表征,可以把它输入到下游的任务中。模型定义是整个训练当中的一小部分,对于模型的前后处理,DGL也提供了100多个围绕图的不同层次的API,这些API比如说包括像构图的算法,比如说做CV的同学会常常用到点云数据类型,有了点之后就可以去构图了。此外,像KNN这些经典的构图算法在DGL当中是有高效的实现的。此外DGL也提供了像图的增删修改,比如删掉一些边增加一些点,以及一些动态图模型的实现。此外还有子图的抽取采样,这对后面我们所提到的巨图训练是非常有帮助的。
图遍历算法和图变换算法,以及图和各类数据格式的转换,比如Networkx,数据科学家所使用的常见的图工具包,包括DGL的Graph和Networkx Graph的转换,可以实现一键转换。此外也包括Scipy Sparse Matrix,这些DGL都提供良好的API去实现转换。
DGL的API也不仅仅优于数量更优于质量,API不仅仅支持同构图,也支持异构图。同时这些API都支持GPU的加速,所以它有非常高效的实现,最后,DGL可以支持多平台,这也是它非常大的设计特色。
我们来关注模型实现,DGL提供了良好的API去支持消息传递。用户可以通过ndata和edata去访问图里的节点和边的特征数据,同时可以使用定义消息函数和累和函数,最后update_all来触发这些消息传递。因此,用简单几行就可以实现GraphSAGE图神经网络的模型定义。这些接口对于更高阶的用户或者研究人员其实是非常有帮助的。
2. 高效的底层系统设计
可以说,消息传递是GNN计算的重要的一个核心,在很多的模型当中,消息传递的计算可以占到GNN训练的50%以上,所以说它的性能的优化是非常重要的。传统的基于gather/scatter的原语的消息传递模式,其实它有很大的问题,比如说pattern geometric使用gather/scatter的原语会把节点的信息汇聚到到边上,在边上做过消息计算之后再把累和到node的Buffer里面。这就会导致在计算过程当中产生大量的消息计算,而这个消息计算消息对象,其实它的数量是和这个图的边的数目是有关的,而我们知道一个稀疏图的边的数量数目是远远大于点的数目的,所以在这个过程当中产生了大量的内存带宽和计算冗余。
DGL首先提出了算子融合的技术,对消息传递进行优化,使得在整个计算的过程当中是不会产生消息对象。
首先我们看一下图卷积网络简化的定义,可以看到边上的计算其实就是把它原节点的相应表征拿过来作为消息对象,点上计算就是把消息对象进行一个累和,其实质是一个稀疏-稠密矩阵乘法乘法算子SpMM。稀疏矩阵A就是这个图的邻接矩阵,而它的稠密矩阵X表示了整个图中节点表征,当把A和X相乘的时候,其结果就相当于做了一个消息传递的计算。因为避免了中间的消息对象的生成,此时的计算是非常高效的。
我们再看图注意力网络,当要去计算边上的注意力权重时,常见的做法就是说对于一个边的源节点和目标节点,将向量表征拿过来做一个内积即为注意力权重。因为边是稀疏的,所以计算也是稀疏的。这里,我们利用了采样-稠密矩阵乘法算子SDDMM,把X和Y做外积相当于对它们的每一行和每一列分别做向量内积。在上面的图注意力网络计算也避免了将点上的数据拷贝到边上的带宽开销,因此可以变得非常高效。
在DGL当中,我们拓展了算子,让它能支持各种各样不同的消息传递所涉及到的计算模式。
拓展主要是有两个方面:
第一个方面:传统的Spmm和SDDMM是只支持乘加计算的,我们的g-SpMM和g-SDDMM可以更加灵活的。
第二个方面:传统的算子是支持标量计算单元,而我们是支持张量计算单元的。
3. 优秀的巨图训练性能
在性能评测方面,我们和另外一款非常有名的图神经网络框架PyG做了一些简单的对比,可以看到我们刚刚所提到的算子融合技术的效果是非常显著的,它在CPU和GPU上都有非常明显提高,同时它也节省了很多内存,特别对GAT这个模型。这让用户可以训练更大的图。
我们知道其实真实数据的图的规模是非常巨大的,甚至因此学术圈也越来越关注大规模图数据问题。在OGB评测标准中的一些数据集可以看到,像这些数据集,其实已经有十亿量级的规模,如果比较工业界的图,我们发现工业界图的规模是更大的,有百亿甚至千亿量级的规模,所以说大规模图训练这个问题非常具有挑战性。
GNN的巨图训练和传统的深度神经网络有所不同,首先它要对目标节点随机选取部分邻居节点并进行拓展,这样可以采样抽取出一个子图,进而抽取这个子图对应特征,比如说它的节点和边的特征,最后在子图上训练网络并更新参数,然后重复步骤1直至收敛。
DGL设计了不同的模块去解决不同的问题,比如对于第二个子图采样的过程,DGL有一个专门Sampling模块去解决这个问题。对于抽取子图特征,我们就设计了KVStore Component。最后在网络参数训练,使用了Training component。
此外,DGL对于通信做了一些优化技巧,比如当Sample 一个子图之后想要去拿这个子图的特征时,发现特征其实可能存在本地,这时仅需要shared memory就能把这个特征拿过来,可以大大节省传输的代价。
此外,DGL也设计了灵活应用的针对图节点和图特征的KVStore,让用户使用图的过程就像单机一样。
最后我们也针对GNN训练多环节的特点设计了流水线,最大程度将数据传输和计算进行并行,通过这些优化,DGL要支持多机多GPU的训练,同时支持多机图数据的存储和划分,在编程接口上保持了单机的高度一致性,目前DGL支持十亿量级的图数据。
4. 完善的开源生态
DGL作为一个开源项目,大家可以通过Github网址进行访问。
回顾DGL的发展历程,可以看到我们在很多的技术点上布局非常早,同时很多技术点,也都是在不断地更新迭代的过程当中。
DGL项目在上线初获得了广泛的关注和好评,在Github收获了8K+的Star以及1.7K的Fork,DGL论文引用数也与日俱增。可以说,在学界DGL是全球领先的图学习框架之一,而在业界,DGL在使用率上就全面领先。可以看到很多知名的学者都对DGL有非常高的评价。
当然,DGL作为一个开源生态,它更多的是围绕很多的开源社区在不断地构建出来,这其中包括像知识图谱嵌入、可视化、异质图、自然语言处理等等。对于DGL所开发的这些包,从DGL作为一个开源工具来说,我们更乐于看到这样的结果。
我将介绍两个最近DGL刚刚发布且非常有趣的开源项目。
(1) GNN可视化工具GNNLens
我们知道图数据的主要问题在于图数据本身非常难以理解,所以在研究过程当中更倾向需要把它给可视化出来,因此我们联合香港中文大学一起开发了GNNLens。GNNLens可以将图可视化显示,并对图数据进行一些操作。我们把可解释性集成到了图可视化工具当中,也希望通过这一点能够帮助研究人员更好地去理解图数据。
(2) 异质图神经网络工具包OpenHGNN
它是由北京邮电大学石川教授团队所开发的。OpenHGNN使用起来特别的方便,训练一个模型只需要使用一行的命令,针对不同的模型、不同的数据、不同的任务,可以实现一键训练。
该工具首批发布有16个SOTA的异质图神经网络模型。此外,工具内置了AutoML进行参数的自动调优,可以很方便去试验一些异质图神经网络的效果。
在开源社区当中DGL有广泛的开源社区的合作伙伴,这其中包括像NVDIA,英特尔这些硬件的大厂,他们给我们贡献了很多的代码来做底层优化,也有来自高校,像MIT,NYU等等也都为DGL贡献过代码,以及360的开源Operator等等。
在开源社区建设方面,DGL投入了很多,因为图神经网络还是一个非常新的领域,我们希望能够把来自学界和业界的不同的研究成果汇聚在一起,让大家及时地沟通。我们每月定期组织用户群的分享会,来邀请学界和业界的研究者分享图神经网络的学习成果。
此外,DGL团队也在学术顶会上做了手把手教程,这些所有的教程材料我们全部公开,大家可以去DGL AI上去找到。
03 未来规划
对于未来的规划,我们即将发布1.0的稳定版本,这其中包括更好的用户文档,更友好应用的内部接口,更多更快的图算子。总体来说,我们希望能够降低大家使用DGL的障碍,提高使用DGL的性能。如果大家有任何想法,请把你宝贵的意见留下来,我们团队也会紧密关注大家的需求。
最后,欢迎大家使用并贡献DGL。欢迎大家加入用户论坛、微信群、知乎专栏,或者加入我们,我们研究院的实习岗位常年开放。如果你对我们整个团队有兴趣的话,也可以把你的简历发送到下面这个邮箱,我们非常欢迎你的加入。
04 精彩问答
Q:在DGL每次训练的时候,当把整个图load到本地进行缓存,假设Matrix或者load的feature特别大,应该怎么办?
A:第一种方式是使用内存映射memorymap,或者使用DGL分布式架构,将图划分到不同的机器上,可以减小内存开销。
Q:如果在训练过程中,模型一开始没有办法定义图的边,而是要动态的加边,这种情况下DGL怎样应对?
A:前面所提到的DGL的一些API,比如增删修改边的API可以很好地满足需求。使用API在训练的过程当中不断加新的点,加新的边,然后在新的图上继续做message passing,算出新的表征,然后继续进一步的工作。
分享嘉宾: