1.6版本详解手册 | 一文看懂昇思MindSpore Graph Learning

为何需要开发图学习框架？

图神经网络（GNN）在工业界和学术界都引发了极大兴趣。应用场景覆盖了药物识别与发现，推荐系统，交通流量预测，芯片设计等等领域。

GNN井喷式的应用，需要不断的快速迭代GNN模型的研究、开发和应用，然而我们观察到基于现有的GNN框架采用消息传播机制进行GNN算法开发，需要用户对所有节点的特征张量进行编程，与算法设计从中心节点出发聚合邻居节点的描述存在gap，给GNN算法开发造成了一定的难度门槛。而且现有框架没有充分考虑GNN算法流程的整体优化如GNN特有融合模式、反向重计算等，执行效率有待提升。

MindSpore Graph Learning

MindSpore Graph Learning是香港中文大学的James Cheng教授团队吴一迪博士等人和华为昇思MindSpore团队基于昇思MindSpore提出的图学习框架，尝试从易用性、高性能的角度寻求突破，创新性地提出了以点为中心的编程范式和针对图学习的编译优化策略。

MindSpore Graph Learning的整体架构图如下：

以节点为中心的编程模型

GNN算法设计，通常从中心节点的视角出发， 描述如何转换其邻居节点的特征向量并进行聚合，具有局部化、低维度的特点。以GCN算法为例，

     表示第层卷积计算后节点的特征向量，由其邻近节点的特征经过线性变换后求合，然后进行非线性转换得到。一个最为自然的Python写法是：

然而我们观察到，深度学习框架以张量为中心无法支持这样直观的表达方式。用户必须依赖GNN框架提供的接口，如消息传播机制的message和reduce来实现这样的聚合，但是仍然需要定义消息发送函数和消息聚合函数将聚合逻辑转成对图中所有节点特征组成的高维张量进行计算的代码逻辑，相比于直观的表达，代码编写的复杂度提升了一个维度。

因此，我们提出了以节点为中心的编程模型。我们展示GCN核心步骤的实现如下：

其中v为中心节点，调用v.innbs获取它的邻居节点列表，然后通过g.sum进行求和聚合。可以清楚的看到，这种表达方式最大程度的贴合了用户的思考习惯和使用习惯，使得编写GNN模型像编写普通Python程序一样简单。我们再以GAT模型为例，进一步展示我们的编程模型：

GAT的公式如左图所示，u和v代表了边（u，v）的起点和终点。我们可以看到在我们的GNN框架下可以最大程度上做到了对左侧公式的一一对应。

以节点为中心的编程模型极大降低了用户开发新模型的门槛，方便用户快速对模型进行实现和迭代。同时和基于message-passing的方式具有同等表达能力，可以支持当前已有的任意GNN模型。

我们利用此编程模型，迅速复现了如今最流行的GNN框架DGL的卷积库和模型库，目前已支持了10+经典模型，覆盖了从同构模型、异构模型、推荐模型、知识图谱到生命科学等多个领域，充分证明了新的编程模型的易用性和强大的表达能力。

如何实现以节点中心的编程呢？MindSpore Graph Learning需要做到以下几点：

• 识别新的编程模型代码并基于MindSpore执行。

• 新的编程模型与MindSpore以张量为中心的代码进行融合，做到两种不同编程模型的代码的无缝衔接。

• 复用MindSpore的特性如自动微分，自动并行，动态图静态图的自由切换，计算图编译，甚至将来可能有的新特性。

为了达成以上目标，MindSpore Graph Learning构建了一套新的parser利用源到源转换的方法将以节点为中心的代码转换为MindSpore原生支持的代码。 

用这种方法，不同编程模型都以张量进行数据交换，可以无缝衔接；并且源到源转换的结果是正常的MindSpore代码，保证了对MindSpore所有特性的兼容性，又不需要对现有框架做任何修改。

这种方法可行的核心关键是现有的图神经网络的邻居获取和聚合运算可以被归纳到Gather 和Scatter算子, 只要支持了这些图操作，就可以支持大部分GNN模型。同时不能转换的代码将直接推给MindSpore运行，从而支持更丰富的表达组合。

为了方便用户对生成的代码进行检查，我们还提供精确到行的代码转换对照。利用此工具，我们将GAT模型的核心代码转换展示在下图：

编译优化

GNN任务具有内存密集的特点，核心开销在于大量的张量在显存-缓存-核心层级间的搬运。减少内存搬运的传统做法是算子融合。我们观察到现有GNN框架中的算子融合优化方案是通过人工挑选几个常用的GNN算子组合，用手写算子组合前向和反向的方法来进行逐点优化。

这种方法有两大缺点，第一，支持新的GNN算子组合方式需要专家进行手动编写和优化，第二，只利用了部分算子融合的机会，没有考虑跨算子优化如算子自动融合，内存规划，同表达式移除等等。

对于GNN 任务，可以归纳成节点特征转换，源节点、目标节点以及边特征交互，转换后特征向中心节点汇聚更新的过程，因而存在频繁数据读写和大量内存占用。

为了减少上游算子产生的中间结果张量通过内存传给下游算子的内存搬运，基于MindSpore 的图算融合进行了算子融合。

区别于其他深度学习模型，GNN模型，无论是同构图还是异构图，通过scatter/gather转化为MindSpore张量计算后，都存在Gather-Injective（包括+/-/*/÷，激活函数）-Scatter (GIS)执行模式。

因此MindSpore 增加识别了新的算子融合模式GIS，将一个或多个Gather算子，Injective算子，Scatter算子融合成一个算子。算子融合之后，在AKG（自动算子生成）算子编译层面，通过polyhedral技术进行自动的算子优化和生成，包含进行特征自适应线程组合并行，节点并行加速计算，以及采用动态Scheduling进行负载均衡，该方案流程如下：

训练过程，融合后的计算图仍会有一些算子具有超大输出内存占用，这些超大输出张量的主要来源为反向计算依赖前向输出，比如一些边上计算的中间张量已超硬件显存可达127G，导致显存不足阻断训练。由于前向算子到反向算子跨度较大，无法通过普通的融合优化来消除这种超大张量。

而且这类超大张量通常也是经过Gather等操作后出现的内存膨胀。为此，我们提出了一种基于算子融合的重计算方案，通过自动识别超大张量对应的计算模式，并进行重计算复制后重新融合，可以达到消除超大张量的目的。

通过减少张量在显存内的实例化，我们获得了比现有最流行框架DGL平均3到4倍的性能提升。此外，这个融合模式除了能够覆盖现有框架中已有的算子融合优化，还能够覆盖大量新的算子组合。

近期展望

在后续的版本更新中会陆续增加大规模图学习性能优化、高效分布式图采样和训练，以及支持多种异构硬件后端等特性。这些新特性将使MindSpore Graph Learning更好支持GNN模型落地商品或新闻推荐、金融风控等基于大规模交互图的工业场景。

长远展望

扩展和丰富MindSpore Graph Learning的应用场景和模型库，构成GNN在生命科学、金融、流体模拟等场景模型库，更好地支撑基于GNN的科学研究和工业应用。