亚马逊马超：如何使用DGL进行大规模图神经网络训练？

演讲嘉宾 | 马超（亚马逊应用科学家）

整理 | 刘静

与传统基于张量（Tensor）的神经网络相比，图神经网络将图 (Graph) 作为输入，从图结构中学习潜在的知识，该方法在近些年已被证明在许多场景可以取得很好的效果。然而，使用传统的深度学习框架（比如 TensorFlow、Pytorch、MXNet）并不能方便地进行图神经网络的开发和训练，而 DGL 作为专门面向图神经网络的框架，可以很好地弥补这一缺陷。该框架在开源后于国内外引起了强烈的反响。

基于此，在CSDN主办的2019 AI开发者大会（AI ProCon 2019）上，亚马逊应用科学家马超，同时也是 DGL 项目的合作作者，发表了《使用 DGL 进行大规模图神经网络训练》的主题演讲。他从 API 使用、系统优化以及系统可扩展性等多个维度深入分享了 DGL 的设计思路。

以下为马超演讲全文，AI科技大本营整理（ID:rgznai100）：

现场大部分朋友都来自工业界，我想应该对传统的图算法了解更多，比如 PageRank 算法应该是我们学生时代学到的第一个传统图算法。但是突然之间，图神经网络算法火了起来。下面这张图是近些年 GNN 论文发表数量的情况，我们可以看到从 2016 年开始，GNN 论文发表数量呈指数上升趋势。

什么是图神经网络？

图神经网络在近年来的发展有目共睹。我们发现，光是去年的 ICML 上有近 50 篇论文都在讲图神经网络。而今年的 KDD 和 NeurIPS 上也有非常多有意思的图神经网络研究。所以我们要问的第一个问题是：为什么要在图上做深度学习？这个问题其实很好回答：因为现实生活中很大一部分数据都可能是图数据，而不是传统深度学习算法里的欧几里得（Euclid）结构数据。

相比之下，我们做语音、图像这类任务，输入数据都是规整的张量，例如图像中的像素一个个排列下来。但是现实生活中很多数据并不是这样子，比如社交网络。在社网络中，如果每个用户都被当作一个节点，用户之间的关系被当作节点之间的边，那么整个数据就是一个图，这是没办法用类似于像素这样的欧几里得结构来表示的；还有知识图谱 （Knowledge Graph）这个例子。

比如，我们使用百度百科时，关键词之间所组成关联图谱，也是图结构；此外，还有一些研究人员在正在尝试在医疗领域进行新药或者以及新化学结构的预测和合成，而化学结构本质上也是图。对整个图而言，没有一个节点的嵌入表示是孤立的，而传统 CNN 和 RNN 分别把空间和时间上的相临节点作为处理对象的方法并不能很好的刻画出整个图上的关系，局限性很大。

图神经网络能做什么？

在聊 GNN 之前，我们首先思考一下大家熟悉的卷积神经网络 CNN 在做什么。比如指定一张图片，CNN 帮助我们判断是猫还是狗，本质上是卷积神经网络学习了在某些图像上的嵌入表示。而给定一张图，如果用 GNN 来预测节点之间的关系，本质上是 GNN 学习了在图上的嵌入表示。比如我们有一个 NBA 粉丝社群的关系图，这时给你一个新的用户，那么即使我们本来不知道这个新用户应该属于哪个社群，但是仍然可以通过 GNN 来预测这个新用户是不是勇士队的球迷，因为 GNN 发现该用户很多的关系节点都是勇士队的球迷。

那么 GNN 具体的工作机制是怎样的？大家请看这张图。在这张图里，不同颜色的圆圈代表图上不同的节点。我们做图神经网络训练时，每一个节点会给它相邻的节点发消息 message，消息可以来自于边上的特征，也可以来自于节点上的特征。邻居节点拿到消息之后会进行聚合操作，然后通过一个神经网络的非线性变换，例如 ReLu，并把它的输出可以作为下一次发送消息的输入，然后继续发送给邻居节点，通过这样不断迭代学习最终学会图上的参数表示。

相信这个过程会让大家联想起 PageRank 算法的流程，其实 PageRank 算法和这里演示的 GNN 算法过程类似，只不过在 PageRank 里没有特征向量和神经网络的概念。

我们之前提到，从 2016 年开始出现了大量关于图神经网络的研究论文。但是即使你是一个图神经网络的专家，如果没有好的工具，想要去复现已有的论文工作或者创造新的算法，仍然十分困难。比如遇到内存泄露或者内存爆掉的问题时，大部分研究人员都束手无策。除此之外，还有一个问题就是现在的深度学习框架，无论是Tensorflow、Pytorch 还是 MXNet，默认都是基于稠密的 Tensor 表示。

对于稠密张量的计算和自动求导，目前的框架已经可以做的很好了，而实现一个图神经网络程序却不是遵循这样的范式，大部分时间开发者都是在对稀疏数据进行操作。开发者从稠密数据的编程模式迁移到稀疏数据的编程模式上需要跨越巨大的 Gap。如果用户现在想编写图神经网络程序，可能需要花费很多时间进行程序上的优化，并且效果还不一定好。

正是因为存在这些问题，去年我们亚马逊 AWS 联合 NYU 以及 NYU上海，一起推出了 DGL 这个平台，用来帮助用户解决上述的这些问题。DGL 希望既能方便研究人员进行在图神经网络领域进行研究，又可以帮助工业界的用户进行业务部署。

DGL：专门面向图神经网络的框架

如果大家有关注我们的官网的话（https://www.dgl.ai/），可以在上面找到 dgl 的源代码、文档教程，博客资源，论坛讨论等等。我们的代码目前已获得 3000+ stars，450+ forks，60k+下载量，目前已经持续迭代到 0.4 版本。

DGL 发布之后，用户的主要构成来自中国和美国，这和我们团队构成是一样的。到目前为止，市面上也有了其他的图神经网络框架竞争者，那么DGL的目标是什么？其实主要是两个方面：1. 我们要对已有深度学习框架的向前和向后兼容做到最好；2.在第一个目标基础上尽可能提高框架的性能和可扩展性。

对于第一个目标，或者说是挑战。在面对各种各样的 GNN 算法需求是，DGL 框架能否提供一套灵活易用的 API 让研究人员在框架上面可以快速不断尝试新的想法。经过大量研究现有的图神经网络论文以及代码，DGL 最终选择使用 message passing 作为底层的API 抽象。

图神经网络框架最核心部分就是在处理 message passing，因此 DGL 把消息通信放在了最底层的API，用来适应广泛的算法需求，对于大部分 GNN 算法而言，无论算法如何变化，通常都很难脱离 message passing 的表示。下面这张图展示 DGL 为 message Passing API 提供了不同层级底层的封装。

当然，光有 message passing 是不够的，DGL 可以让用户实现自己需要的 UDF（用户定义函数）。下图是一个简单的 GCN 的 demo，展示了如何用 DGL 编写 GNN 程序。

DGL 另一个重要的设计是，我们从一开始就把 DGL 设计成 “框架上的框架”。为了避免重复造轮子，DGL 在设计上采用了类似 Keras 的做法，基于已有的主流深度学习框架之上进行开发，目前 DGL 已经可以很好的支持 MXNet 和 Pytorch 两个主流深度学习框架。

考虑到很大一部分用户可能就是想利用现成的代码把已有的算法拿来玩一玩，所以最近我们一口气写了很多现成的 NN moudule，方便大家可以直接调用。此外 DGL 还提供多种 GNN 模型的完整实现和运行脚本，欢迎大家都来玩一玩。包括我们在 0.4 版本即将推出的 Knowledge Graph model zoo.

DGL的性能优势

DGL 在性能上做了很多系统优化。例如消息融合，DGL 通过稀疏矩阵乘法的转化来做 message fusion。此外，为了让系统在稀疏数据上的性能达到最好，DGL还专门启动了一个名为 mini-Gun 的项目来做 kernel 优化。除此之外， DGL 内部还会通过 auto-batching 技术来加速批量图的并行处理过程。

谈论了单机上的性能，那么遇到大图怎么办？比如微博有好几亿用户，用户的关系可能是百亿级别，这样的图无论是对 GPU 还是单机来说都装不下。因此 DGL 做了分布式的架构。

这个架构有三层，最底层是采样器，中间这层的 trainer 是训练器，最上面是DGL-PS参数服务器。如果图非常大，我们就先把大图切分成很多小图，再把它分到不同的机器或者不同的GPU上，每次训练时只使用一小部分的 miini-batch 数据，并且通过参数服务器来存储训练过程中需要的模型参数，trainer 和参数服务器之间通过我们自己实现的网络协议进行数据交互。

最后我想说的是，现在深度学习的主要数据结构是 Tensor，但是 Tensor 是真正的数据结构吗？大家可以思考一下这个问题，如果把程序=算法+数据结构这句话扩展到深度学习领域中，会不会Graph就是深度学习的重要Date Structure？这也是为什么说一定要抓住图神经网络这个机会。黑盒式的神经网络一定不是长久走下去的道路，神经网络的发展在某一天可能会遇到可解释问题的瓶颈，图神经网络是否能为我提供一个神经网络可解释性的有效方法，值得我们思考。