原来GNN这么好上手，OMG！用它！

2021-12-30 06:02:33

原创：阿里云计算平台事业部沈雯婷、艾宝乐

Overview

Graph-Learn(GL) 是阿里巴巴开源的高性能工业级大规模图学习系统，本文将对GL的用户接口做一个概览，并介绍GL丰富的图采样算法，以及GL灵活统一的GNNs模型框架，帮助用户快速上手GL。
项目地址：https://github.com/alibaba/graph-learn 。

图数据广泛存在于我们生活中，社交网络、互联网、交通网等等都是天然的图数据。在阿里内部场景，用户物品的交互数据，安全风控领域的用户、评论、交易等都构成图，图可以说是数据最为广泛的存在形式。

深度学习在文本，语音，图像等数据上取得了很大成功，GNNs(图神经网络)是深度学习在图结构数据上的应用。对于图这一非欧空间的数据，如何将它和深度学习结合起来是GNNs首先要解决的问题。GNNs早期的研究主要集中在Spectral-based 方法上，这种方法理论比较充分，但是难以用于大规模场景。近些年以GraphSage为代表的Spatial-based方法在性能，泛化性，灵活性等方面均体现出了优势。GraphSage通过采样子图进行按批次训练的方式在Pinterest推荐场景进行了实际应用，并取得了很好的效果。此外，像DeepWalk等Graph embedding算法和TransE等知识图谱模型的训练都是基于采样的方法进行设计和训练。

阿里很多场景的数据都是十亿级别节点，百亿边的规模，针对阿里众多场景的实际需求和特点，我们设计开发了简洁易用，高性能的工业级大规模图学习系统Graph-Learn(GL) 。GL主要面向spatial-based的GNNs算法，同时支持graph embedding, 知识图谱等常见图学习算法。

GL的整体模块如下图所示，底层由图引擎和NN引擎组成，通过数据模型管理采样后的图数据从而利用NN引擎的计算能力接入上层算法。

Graph Engine提供了分布式大规模图的查询、采样等操作。基于简洁的接口，你可以遍历图、得到邻居样本、得到所需属性，从而自己组织数据，构造模型。

Data Model是GL的基础数据模型，我们以种子节点（seed nodes/edges）和邻域（receptive fields /multi-hops neighbors）组成的子图为基本编程对象，并提供了对接上层神经网络模型的数据转换功能。基于数据模型，可以轻松管理样本，并专注于模型开发。

基于Data Model，Graph learning models通过抽象多种encoders模块可以将图数据转换为最终的embedding。GL提供了多个built-in的GNNs models，并封装了若干常用的encoders和对应的图卷积层，图聚合层。Graph learning models提供了e2e的模型训练评估过程，可以在单机和分布式环境一键运行，同时基于内置的各个模块也可以方便快速构建自己的训练过程和算法模型。

图引擎

简洁的接口：从裸数据到GNN样本，只差一行python

Graph Engine包含图对象模块、采样模块和查询模块，这些 Graph 上的模块接口通过一套Gremlin-like API表达。如何从裸数据构造单机或分布式的图、如何在图上游走、如何遍历、如何采样邻域、查询哪些field，以及如何用这些数据构造图神经网络模型所需的样本，整个过程只需要一句python表示，类似data-flow的查询语句：

gl.Graph().node().edge().init().V().outV().sample().by().values()

我们将这一行代码拆分到图对象接口、采样接口和查询接口中进行解释。

(1) GL图对象
图对象模块用于将结构化的图数据转换为逻辑图对象。GL的入口非常简单，载入graphlearn库，构建一个 Graph 逻辑对象，后续所有的操作都在这个 Graph 对象上进行。

import graphlearn as gl
g = gl.Graph()

图数据格式：灵活的schema，多变的数据类型

GL图数据格式灵活，支持float，int，string类型的属性，支持带权重、标签。在现实的场景中，数据格式多变，通过config文件描述非常复杂，容易写错；有些GNN系统不支持多种类型的属性。
下面的示例描述了数据中存在string和float类型的两列属性和权重、标签列。

decoder = gl.Decoder(attr_types=["string", "float"], weighted=True, labeled=True)

数据源载入和拓扑描述：同构、异构、多数据源，通通支持

GL提供了 node 和 edge 两个简单的接口来支持顶点和边的数据源载入，同时在这两个接口中描述图的拓扑结构，比如“buy”边的源顶点类型是“user”，目的顶点类型是“item”。这个拓扑结构在异构图中十分重要，是后续采样路径meta-path的依据，也是查询的实体类型的基础。

g.node(data_source, node_type, decoder) \
 .node(data_source, node_type, decoder) \
 .edge(data_source, (src_node_type, dst_node_type, edge_type), decoder)

Graph Engine启动：快速拉起大规模分布式图引擎

GL提供单机的版本，通过init 接口快速启动Graph Engine，至此，图对象已经构造完毕，查询、采样操作就可以在 Graph 上进行了。

g.init(）

在大规模的场景下，图顶点可达亿级别，边可达千亿级别，不管是图结构还是图属性，都无法完全载入单机内存。GL提供分布式的Graph Engine，速度非常快，使用上也非常简单，只需要在 init 中加几个参数。

g.init({"server_count": N, "client_count": M}, task_name, job_index)

(2) 采样接口
采样在GL中通过游走路径和采样策略进行描述。
如下示例中，表达的是一个遍历和二跳邻居采样，即遍历图获得64个用户，对每个用户根据边的权重采样50个他们购买的商品的相似商品。

q = g.V("user").bath(64).outV("buy").sample(5).outV("similar-to").sample(10).by("edge_weight").values()

除了采样正邻居，GL也提供了负采样，只需要将示例中的 outV 改为 outNeg 即可。
GL提供了非常丰富的内置的采样和负采样策略，也支持自定义采样策略的实现，我们将在下一章详细解析大规模采样算法在GL中的实现。

(3) 查询接口
GL提供了 Nodes 和 Edges 两个基础的数据类型，作为遍历、采样的结果。为了获取 Nodes 的int类型的属性，可以调用如下接口进行查询，得到的是numpy array数据结构。

nodes = g.run(q)
nodes[2].int_atrs # nodes is a list of Nodes, include Nodes of user,
                  # Nodes of item for 1 hop, Nodes of item of 2 hop.

g.run(q) 可以多次执行，遍历图中的顶点和他们的边，直到遍历完毕。因此，上述采样和查询的结果可以作为generator接入tf或pytorch等NN引擎作为数据源，从而深度定制GNN模型。
GL也提供了数据模型的封装，来接管数据采样流程和样本组织。

丰富的采样：GNN落地工业级场景的基石

图采样算法是GNN落地工业级大规模场景的基石，GL提供了非常丰富灵活的采样算法。
本节将对GL的采样实现和各个采样算子做简单介绍，采样在GNNs中起着重要作用，并已经在很多场景得到了实际应用。

采样和查询是GL图引擎部分对外提供服务的基本接口，采样和查询构建于分布式图存储之上，以Python API的形式对外提供，如下图所示。查询包括图的节点和边的遍历，节点和边的属性、权重、标签等查询。采样分为邻居采样和负采样，邻居采样提供采样子图的功能，负采样主要面向非监督学习场景，系统直接提供采样负样本的功能，不需要用户额外产生负样本。

基于阿里内部大量业务场景的需求总结和实际应用，系统集成了若干种常用的采样和负采样方法(Built-in Samplers)。虽然这些采样方法基本覆盖了大部分需求，但是因为采样方法和算法效果关系密切，甚至直接关系到模型的成败，因此，开放采样编程接口，方便用户进行其他自定义采样方法(Custom Sampler)的探索是我们系统的一个重要功能以及GNN创新的一大方向。

GL的图引擎是CS架构，具体的采样实现流程如下图所示：给定一批需要采样的源节点，首先通过Partitioner去查询包含这些源节点的子图在哪个server上，图中src id为1，3的被partition到了server 1上，2，4被partition到了server 2上。然后，将这些src ids发送到对应的server上，查询子图，并调用本地采样方法去执行采样操作。目前系统集成了若干种常见的采样算子，也可以支持用户自定义自己的采样算子。本地采样结束后，需要将采样后的结果按照原始src ids的顺序进行拼接返回。Partition目前使用系统内置的hash划分方法，将来我们会支持不同的Partition策略来进行更加高效的图划分，提高采样和计算性能。

对于负采样来说，没有partition和stitch过程，一个batch的输入数据，会随机或者按照分布选择一个server，然后在这个选定server上进行负采样，这样可以做到全局负采样的效果。

我们底层的数据结构使用Tensor的形式实现，支持常见数值类型和string类型的Tensor。基于Tensor这种规整的格式，系统可以将采样的Partition，Stitch和分布式执行过程自动化，因此，新增一个采样算子时只需要写具体的本地采样逻辑即可。

GL采用了分布式图存储和采样，使得规模可以轻松扩展到上百台机器，支撑起十亿节点，百亿边的日常任务。此外，针对不同的采样策略，我们采用了热点缓存，慢机迁移等众多优化，使得在规模提升的同时，能够保持高性能地采样。

邻居采样
GL支持异构图，同构图是异构图的一种特例，因此图采样需要指定meta path，系统按照指定的meta path采样得到多跳邻居，不同的采样算子按照不同算法来返回邻居。目前built-in的采样算子包括以下几种：

random：随机采样邻居，每次在给定节点的所有邻居里随机选择一个邻居。
edge_weight：按照边的权重为概率分布进行采样。
topk：以edge_weight为大小对邻居进行排序，并返回topk个。邻居数不足要求个数会循环填充。
in_degree：按照邻居的入度大小为概率分布进行采样。
full neighbor: 返回给定点的所有邻居。

前四种采样都需要指定meta path，并且需要指定需要采的邻居个数，采样结果以dense的形式返回，这是由于采样后按batch训练的时候需要样本对齐。当然对于原始GCN等需要获取源节点的所有邻居的算法，我们支持采样返回节点的所有邻居，而不做上采样/下采样，我们称这种采样为full neighbor采样，采样后的结果会以sparse的形式返回。

以上图为例，假设需要采样src id为1的节点的邻居，指定的邻居个数是2。那么random会从[2, 3, 4, 5]里随机采样两个；edge_weight以边权重为概率分布返回两个；in_degree采样以入度为概率分布返回两个；topk会返回边权重最大的两个，即[5, 4]。如果指定采样个数为6，则topk会做循环填充，返回结果为[5, 4, 3, 2, 5, 4]；full neighbor采样的话则会返回实际个数的全部邻居[2, 3, 4, 5]。

以上采样算法里按概率分布的采样均采用了AliasMethod实现，因此复杂度都是常数时间。单机在512 batch size、采样10个1hop、15个2hop邻居id的情况下，耗时在毫秒级别。此外针对不同的策略系统有不同的优化，比如对于topk，我们在构图时就会预先按照edge weight为大小对底层存储表进行排序。

上面介绍的算法以node-wise(layer sampling)采样为主，node-wise采样随着采样的跳数增加，邻居数目会急剧上升，在多跳(>2)情况下尤其明显，因此近年来layer-wise的采样也被应用于不同GNN算法，我们也正在探索开发layer-wise(graph sampling)的采样，希望在多跳场景下带来进一步的性能优化和内存节省。

负采样
负采样就是给定源节点，返回和它不相连的目标节点。负采样同时支持本地和全局负采样，对于一个batch的样本要采样其负样本时，可以按照一定的概率分布选择一个server，然后在该server上采样。如果选择本地负采样，则只在本机进行负采样。目前built-in的负采样算子包括以下几种：

random：随机采样和给定源节点不相连的目的节点。
in_degree：按照目的节点的入度分布，返回与源节点不相连的目的节点。
按照入度分布的in_degree负采样实现上采用AliasMethod[5]，会在首次采样前提前构建好Alias Table，因此采样是常数时间复杂度。此外负采样我们默认是严格负采样，但是出于性能和极端情形的考虑，我们支持用户端设置一个阈值来控制严格程度。负采样可能对算法效果也会产生很大影响，因此负采样算子如何和图划分结合，如何高效负采样，以及按何种分布负采样都是值得探索和研究的地方。

自定义采样
上面介绍了GL开发过程中，结合实际业务场景需求系统已经集成的邻居采样和负采样算子。虽然这些built-in的采样算子满足了大部分GNNs和其他图学习算法的需求，但是，考虑到每个具体业务场景的采样、负采样需求不同，而且采样、负采样策略对算法的构建和效果有着重要影响，因此我们也允许用户自定义采样算子。自定义采样算子主要流程包括：
a）在C++文件中注册并实现新Sampler

class MySampler : public Operator {
public:
  Status Process(const OpRequest* req, OpResponse* res) override {
    //TODO
  }
};
REGISTER_OPERATOR("my_sampler", MySampler);

b）在Python端调用已注册的算子

g.sample(count).by("my_sampler")...

算法模型

底层的图引擎提供了分布式存储和采样查询功能，基于图引擎，GL实现了一套统一的图学习框架，可以和常见的深度学习框架比如TensorFlow, PyTorch进行结合。

Data Model：那些Dirty work系统帮你做了

图引擎提供的采样结果如何高效得和顶层神经网络模型结合是实现图学习算法的重要一环，为了简化和优化这一过程，GL封装了一套适配不同深度学习框架的数据模型。首先，将采样后的结果以EgoGraph的形式进行封装，为了兼容各个不同的深度学习引擎，EgoGraph 中的数据类型为numpy array。EgoGraph在后面NN部分使用时需要转换成对应的tensor格式EgoTensor, 这一转换过程通过EgoFlow来管理和优化。

Graph Learning Models：灵活、高效、统一

所有的图学习模型都需要编码器来对点、边或者子图进行编码从而得到点、边或者子图的embedding。GL首先实现了若干特征编码器，用来对采样后的EgoTensor进行编码，然后，使用图编码器来对原始embedding进行编码得到最终的embedding。对于大多数GNNs算法，图编码器实际是aggregate和combine过程，通过不同的卷积层可以轻松搭建一个图编码器。

基于这些编码器，GL封装了多个built-in的模型，包括：GraphSAGE，Bipartite GraphSAGE，GCN，GAT，DeepWalk，LINE，TransE，同时也封装了很多常见loss函数，opimizer以及训练预测过程，用户可以快速进行本地和分布式训练。在GL中运行以上模型，只需要一键执行python脚本。
比如我们要执行一个二部图的GraphSAGE算法，样例数据已经准备好了，在examples/data目录下，运行脚本就可以得到数据，模型训练可直接调用如下脚本。

cd examples/tf/bipartite_graphsage
python train_unsupervised.py

我们也将模型进行了抽象，包括Layers、Aggregators的复用，方便用户搭建自定义不同的图学习模型。GNNs是近年来图数据分析与应用的热点研究问题，学术届和工业界都在不断提出新的模型，我们也将继续探索，不断完善模型，并提升大规模下的性能。

总结和展望

GL上的GNNs模型正在快速发展迭代，我们也将把更多在阿里的大规模业务上经过验证的GNNs模型开源出来，欢迎大家加入到GL的共建中。

项目地址：https://github.com/alibaba/graph-learn
欢迎感兴趣的同学加入我们的团队 kun.zhao@alibaba-inc.com 。

码农公寓