图数据切分与模型数据载入的问题解析

什么是 GNN/图神经网络？

图神经网络是一类 ”用于处理可以表示为图的数据“ 的人工神经网络。我们通常能看到的例子有: 社交网络的商品的推荐，蛋白质三维结构的预测和芯片设计优化等等。到目前为止，之所以图神经网络能够攻克了这些涉及领域比较核心的问题，包括其涉及的理论知识， 模型设计，benchmark 数据集效果验证等等一系列工作， 都能反映出了它结合图数据表达的能力和神经网络模型泛化的能力。

那我们又要基于什么样的前提，在 GNN 的范畴开始我们自己的探索和应用呢? 我希望通过对两个“假设”的描述 ，让大家更了解这个前沿技术背后的动机。

第一个假设：关于数据样本的分布、频率和空间位置表示

6 年前，谷歌 AI 趣味实验的一个展示平台上 有一个机器学习任务非常吸引我。他们当时的目标是通过⻦叫声分辨出⻦的种类。最左边所有的小图片，每一个代表北美一种⻦类叫声的声纹信息， 右边则是我们找到的对应的“鸳鸯”的照片。这个机器学习可视化项目，通过不使用标记训练的 t-SNE，将高维声纹信息转换为二维的空间位置来展示数据 并且验证了相似的种群有相似的叫声这个假设。

放在现在来看，我们有理由相信。当时研究人员的梦想 是可以用深度学习的模型去解决的， ⻦类有成千上万种，一种⻦的不同叫声也有很大的差别， 只要我们有大量的数据样本和标签可以去做训练，这个事情就是可行的。

我们知道深度学习，能解决这样直观但又复杂的事情， 恰恰是因为我们认定了这样的假设; 并且着眼于如何处理大量的数据，使用更深层、更精妙的神经网络模型设计，来自动挖掘出数据的规律，并最终给出判断。

假设二：一个种群有一个大致的活动范围

假如我们把问题换做：我想知道在哪一个城市可以听到这样的⻦叫声呢? 解答这个问题的方法是使用图算法。我们可以看到数据中记录的位置信息，属于地理位置的不同区块; 根据数据采集的密度决定区块的大小，把小的区块与种群关联， 再用社群发现的算法，得出不同物种的空间分布。

这个生物地理的问题，其实给定另一个假设，即一个种群有一个大致的活动范围。只不过这样的假设比之前样本特征分布的假设，更加容易以一种显式的方式体现。趋向于普遍事实的“假设”是可以被作为数据存储下来的， 比如我们常说的知识图谱中的三元组关系是对某一个事实关系的描述一样。

图神经网络

而图神经网络模型，恰恰是在这两个前提假设的融合之上建立的。如果我们也能够确信，我们手上的数据蕴含了这样的假设， 即样本特征分布的假设与逻辑关系的假设， 我们一定可以去尝试，并且有信心能构建出有效的 GNN 模型。

回到今天的主题，当我们有数据集需要做模型训练时，我们是如何对其进行训练集、验证集测试集的切分? 我们又是怎样对大到无法放入内存的数据进行分批次的训练?

如果放在 GNN 的场景下，我们需要考虑的是，如何切分图数据集，又如何为 GNN 模型分批加载子图。要知道在研究和应用上，GNN 是机器学习领域最火热的方向，各式各样的理论和方法层出不穷。有一点可以确信的是，无论从理论还是工程方面，这两个步骤并没有太高深，却也是 GNN 项目中比较重要的部分。

我也会在最后通过 TigerGraph 云实例，演示我们的图机器学习工作台，并在一个运用 GAT 模型进行反欺诈检测的 Demo 中，展示这两个步骤的作用。

图数据切分

首先，我们来讨论如何对图数据进行切分:

对于有监督机器学习模型，我们对样本的变量都会有一个独立同分布的假设。通过这个假设我们可以随机的按比例去切分训练集和验证集，用于调整模型的拟合，并在测试集上观察泛化的效果。或者可以按照时间发生的先后顺序去切分， 因为我们希望新的数据的特征同样能被模型识别到。然而通常来说，一定会遇到样本的分布发生变化; 样本存在分组; 有标记样本数量的平衡性等现象。

按照统计经验的方式切分，模型会在真实的数据上发生变化，影响了实际业务的判断。而这样的现象，当我们用图的思维去构建数据时，会更加显而易⻅的发现。可以假设一个生活中的例子来构建一个图机器学习任务:

我们希望通过人群访问地理位置的记录去预测场所的标签。这似乎与之前的种群分类的问题很像， 但显然我们不能直接通过之前的图算法给出结论: 因为标签本身是有含义的，并且每个地点可能有多个标签。

那如果我们为传统机器学习的模型提供图特征呢? 比如我们把【人群】的特征，归并到【地理位置】中去， 我们会在切分数据的时候发现一系列的问题。可能我们会发现训练集中【人群】的多少，访问时间段等信息与【场所】没有太强的相关性 ;或者在验证集上【地理位置】的特征的分布发生了偏移; 更重要的，我们如何保证测试集上新的数据，【地理位置】更新的生命周期，涉及到的【访问】关系 不会显著减少? 我们尝试各种方法:延⻓样本的时间跨度、在数据集上做交叉验证，缩短模型上线后的运行周期，甚至尝试剔除【人群】相关的变量。

这显然不是我们希望看到的，在【人群】这类特征上来看，按时间切分显然是不合适的。

我们是否应该重新考虑修改图模型本身呢? 答案是肯定的，我们可以尝试把人群的特征作为【地理位置】间【转移】的关系进行建模。这样，我们的选择就变多了:我们可以首先得到一个同构图，人群移动的信息就被包含在了地点之间【转移】的关系里面，就像种群迁移一样。我们还可以将转移所需的时间，作为距离和关联性的特征，也放在边关系里。试想骑手们在赶单子和上了一天班的你坐地铁回家，一定会产生不同的分布。甚至，我们还可以建立多条关系比如所有早高峰，午餐高峰等不同时间段的统计等等。这样构成的网络，就包含了【人群】带来的真正有用的信息。

回到刚刚的切分数据的问题本身，我们可以思考一下为什么需要【地理位置】的标签呢? 其中一个最直接的原因是，大部分的地理位置标签我们是没法直接拿到的，并且也正是我们希望模型给出预测的。

比如开辟一个新地块的【地理位置】标签，或者说 【人群】活动较少的地方这些【地理位置】的标签，就不会像大家常去的有 POI 的地方容易拿到标签。在这种情况下，按照地理位置的”活跃度“切分数据，恰恰是最合理的。

那么对于可以表示为图的数据，还有哪些切分方式呢?

我们可以看到 OGB这个专⻔为GNN模型准备的公开数据集，规范了一些切分方式。之前提到的人群迁移的这个例子，就是借鉴了这个数据集中【产品分类任务】的构建关系和切分的方法。当然从实际出发，我们也可以用不同的方式切分数据集，让他们同时落在一张图或者完全不相关的图上。这些切分方法，其实并没有在意样本的数据分布是否均匀和时间发生的顺序，而是根据实际情况出发，在建立关联关系的过程中，显示的把逻辑表述出来，数据本身的变化和无法及时获取的问题得以规避。最终，我们把发现分布的事交给各式各样的图神经网络模型，发挥他们不同的泛化能力。

总结一下，在图上做机器学习的任务，“切分数据” 这一步骤的目的，并非为了让模型更容易训练， 而是让模型训练的逻辑假设更加清晰，样本主体的特征保持稳定，模型结果更容易验证，并且最终投入生产应用。

模型数据载入

接下来，我们来看一下在训练阶段 GNN模型对我们的数据载入有哪些要求。

从基本的模块来看，图神经网络模型和其他神经网络模型一样可以堆叠层数，但通常会有一个置换等变层。这一层的目标，就是收集模型训练样本的邻居，把他们的信息传递过来。根据模型的不同，可能是单跳的邻居，也可能是多跳的邻居。聚合时，可以用固定权重去聚合这些信息，也可以是一个参数的形式来学习邻居之间的重要程度。

从这个过程来看，图神经网络，和其它一般的神经网络不一样的是: 如果我们的输入是图片、声音和文字这样的高维数据，我们可以想方设法做一些预处理来减少输入维度。而对于组成网络的数据而言，因为我们的初衷是希望尽可能地学习到图拓扑结构。我们很难在模型以外预处理他们(虽然也是有办法的)。

这就导致像 Pytorch / Tensorflow 这样的框架需要去做除了计算以外的事情，比如特征索引和邻居索引的工作。

从数据集的⻆度来看，我们必须正视拓扑结构对 GNN 模型架构的影响。大批量数据集本身，是没有办法放入一台机器的内存中的。对于深度学习，分批次的载入数据到更快的计算单元的内存(GPU)， 是让深度学习模型的训练扩展到大量数据至关重要的步骤。而大型的图数据集，尤其是密度很高的图数据集， 我们还需要额外的处理方式：比如通过邻居采样的方式去规避邻域爆炸的问题。

当然，不是将数据进行分批就一劳永逸了。我们仍需要考虑数据传输效率问题。举一个 PyTorch 的数据载入模块 DataLoader 为例，这个图表模拟了模型在进行分批训练时，当载入速度(比如2秒)和批次消耗速度(比如每个批次的正向和反向传播只需要 100 毫秒) 相差过大时，整体的训练速度会比平衡时慢两倍多。

如果说，我的载入速度能降到700 毫秒， 但是我开了 7 个 worker，这样的话就可以把消耗和载入速度平衡了。我们还要考虑每个批次的载入速度差异，差异过大，即使有很多worker 和管道，也会降低训练的速度。这样看来，在关联信息的图上做这样的控制，保证批次数据的载入速度和稳定性也是十分重要的。

从架构层面看，我们会发现特征索引和图的索引是 GNN 模型中不可避免的工作。

我们看到 PyG（基于PyTorch的图神经网络架构 ） 在对 GNN 的实例描述中，提及了这样几点发现：

● 获取具有足够 CPU DRAM 来同时存储图形和特征，是非常具有挑战性的；

● 使用数据并行性进行训练，需要在每个计算节点中复制图特征和图拓扑；

● 图和特征很容易就超过了单台机器的内存。

因此他们额外抽象出了 FeatureStore, GraphStore 这样远程的抽象层，由开发者决定如何连接提供有效的图特征采样的外部存储。作为同时具有扩展性和计算性能的图数据库，我们也明白这是TigerGraph需要出力的地方， 并同时达到一个实时的、安全的、稳定的、更方便获取特征的平台，加速整个 GNN 项目的演化和上线。

审核编辑：刘清