专题解读 | 图神经网络内核库tf_geometric

国人自研的tf_geometric是一个与Tensorflow1.x和2.x兼容的构建图神经网络的内核库。内核库由构建高效GNN的基础设施组成，包括图数据结构、图MapReduce框架、图Mini-Batch策略等。这些设计使tf_geometric能够支持单图计算、多图计算、图小批量、分布式训练等；因此，tf_geometric可以用于各种图深度学习任务，如归纳节点分类、归纳节点分类、链接预测和图分类。下面介绍tf_geometric的相关论文。

Efficient Graph Deep Learning in TensorFlow with tf_geometric

下面是tf_geometric库的框架图：

tf_geometric的内核由五个部分组成：

• Graph Data Structure
• Graph Map-Reduce Framework
• Graph Mini-batch Strategy
• Distribution Strategy
• OOP & Functional API

以下依次介绍每个内核模块。

1.Graph Data Structure

tf_geometric图结构分为Graph和BatchGraph两种，结构如下图所示。BatchGraph是在聚合多个Graph的信息形成表示。这些信息传入神经网络进行训练后，按需求和相应的处理可以完成图级别或者节点级别的任务。

2.Graph Map-Reduce Framework

这是tf_geometric非常有特色的一个设计。

Map-Reduce是一种用于处理大规模数据的编程模型和处理框架，最早由Google提出并应用于其分布式计算系统中。它的设计目标是为了能够高效地处理海量数据，并且能够充分利用大规模计算集群的并行处理能力。

而tf_geometric的作者将这种优秀的数据计算思想融入到图神经网络算法库的设计当中，使得模型的数据流更加符合线性计算的思维，方便用户更好地处理流式数据和复杂计算。

下图显示了一个Map-Reduce工作流程的例子，它计算了图注意网络（GAT）的归一化注意分数。在这个例子中，每条彩色边包含一个节点和它的一个邻居节点的特征向量。

每个箭头对应着一步基本操作，可以看出所有的操作都可以细分为Gather、Map和Reduce操作。其中Map操作完成了边数据到边数据的映射，Reduce则完成了边数据到节点数据的映射。一般的模型设计通常都可以将数据流的处理细分为这些操作。而通过这样的设计也使得算法设计拥有更好地扩展性和可读性。

3.Graph Mini-batch Strategy

与一般的深度学习模型一样，GNN可以从图的小批量训练和推理中受益。给定一批图，tf_geometric将它们组合成一个BatchGraph，并在其上应用图神经网络。由于标签信息在BatchGraph中也被合并，合并后的信息可以直接作为批次的节点/图的标签。tf_geometric中的mini-batch策略很灵活，不仅可以对输入图数据进行mini-batch，还可以对中间输出图进行mini-batch。在mini-batch过程中，由于将图形组合成批量图形的操作是可微调的，梯度将从批量图形传回给定的一批图形。此外，mini-batch构建操作足够快，可以在每个前向传播过程中执行。

4.Distribution Strategy

为了利用TensorFlow强大的分布能力，tf_geometric中的所有GNN模型都被实现为标准的TensorFlow模型，这些模型可以在数据处理上以最小的代码变化进行分布。tf_geometric GNN模型的分布可以通过TensorFlow分布策略轻松处理，通过选择不同的分布策略，模型可以以不同的方式分布。然而，图数据的分布不能通过简单地应用分布策略来解决。这是因为分布策略的内置数据分片机制是为规则张量设计的，它不能处理不规则图数据。尽管如此，用户仍然可以通过用tf.data.Dataset定制分布式图数据集的分布策略来轻松分配图数据。定制通常只需要对代码进行一些小的修改，以便进行本地数据处理。下图是tf_geometric中给出的实现图分布式策略的代码，借助于Tensorflow的分布能力，实现起来较为简单且清晰。

5.OOP & Functional API

tf_geometric为传统的OOP的算法实现提供了Functional即函数式的API调用，以此来方便用户定制化复杂算法实现（如图元学习）。

基于内核库，tf_geometric为不同的任务实现了各种流行的GNN模型，包括节点级模型，如图卷积网络（GCN）、图注意网络（GAT）、简单图卷积（SGC）、神经预测的近似个性化传播（APPNP）和深度图信息（DGI），以及图级模型。为了避免冗余，tf_geometric中的所有GNN模型首先被实现为函数式API，而OOP API只是相应函数式API的包装器。此外，tf_geometric还提供了重现文献中报道的GNN性能的演示。这些演示包含了数据加载、训练和评估的完整代码。它们以一种优雅的方式实现，也作为tf_geometric的风格指南。

Dataset Management Mechanism

1.Dataset Classes and Dataset Instances

Dataset有一个相应的数据集类，一个数据集类的不同实例（数据集实例）可以代表同一数据集的不同配置。每个数据集实例都可以从网络上自动下载原始数据集，然后将其预处理成方便的数据格式，这不仅有利于tf_geometric，也有利于其他图深度学习框架。此外，数据集类提供了一个缓存机制，它允许只处理每个原始数据集一次，并从缓存中即时加载。数据集类不仅仅是原始图数据集的简单包装器，它们也可能在预处理中涉及复杂的特征工程。

2.Built-in Datasets

内置数据集，由很多公共基准数据集组成，在图深度学习研究中经常使用。此外，内置的数据集涵盖了各种图深度学习的数据集。

3.Customizable Datasets

用户可以通过简单地对内置的抽象数据集类进行子分类来定制他们的数据集。内置的抽象数据集类管理数据集处理的工作流程，并且已经封装了一般流程的实现，如下载、文件管理和缓存。这些一般的过程可以通过子类中定义的配置参数来定制，比如数据集的URL和预处理结果是否应该被缓存。由于不同数据集的数据预处理过程通常是不同的，数据预处理被定义为超类中的抽象方法，用户可以通过覆盖抽象方法来实现他们在子类中的数据集。

Utilities

实现图深度学习模型需要一些重要的工具。这些工具包括常见的图数据处理、类型转换、图采样等。tf_geometric提供了这些utils模块，其中大部分不仅为tf_geometric设计，也为一般的图深度学习实现设计。

此外，论文给出了对比其他算法库的解读。

对比PyG和DGL，作者提出，tf_geometric不仅提供了OOP的API，并且提供了Functional API的设计用来更好地适应前沿神经网络算法的设计。

对比Spektral和StellarGraph（两者都是基于Tensorflow实现的比较著名的图神经网络算法库），tf_geometric兼容了Tensorflow1.x和2.x的版本，在两个版本都有比较好的实现。此外tf_geometric还实现了一套用户友好且便捷的缓存系统，能够更大程度减少重复计算带来的性能损失。

下图是论文中作者给出的tf_geometric在节点分类和图分类任务上的性能指标。