tf_geometric — 基于TensorFlow的友好高效的图神经网络（GNN）库

tf_geometric（tfg）是一个高效且友好的图神经网络库，同时支持TensorFlow 1.x 和 2.x。tfg同时提供面向对象接口（OOP API）和函数式接口（Functional API），可以用它们来构建有趣的GNN模型。

tf_geometric简介

tf_geometric（tfg）是一个高效且友好的图神经网络（GNN）库，同时支持TensorFlow 1.x 和 2.x。

Github主页: https://github.com/CrawlScript/tf_geometric
开发文档: https://tf-geometric.readthedocs.io
论文: 《Efficient Graph Deep Learning in TensorFlow with tf_geometric》（https://arxiv.org/abs/2101.11552，已被ACM Multimedia 2021 Open Source Software Competition接收）

tf_geometric总体架构

下图展示了tf_geometric的框架。从框架图可以看出，tfg提供了用于构建图神经网络的内核（Kernel），同时也提供了许多流行的GNN模型的具体实现。

tf_geometric Kernel: tfg内核包含用于构建高效GNN的各种基础设施，例如图数据结构、图map-reduce框架、图mini-batch策略等，这些基础设施使得tfg可以支持单图计算、多图计算、图mini-batch、分布式训练等。在内核的基础上，tfg实现了多种面向不同任务的GNN模型。

OOP and Functional API: tfg同时提供了面向对象（OOP）和函数式（Functional）接口。目前一些流行的GNN框架如PyG、DGL等主要支持OOP API，缺乏对Functional API的支持。Functional API的欠缺使得这些GNN框架难以应对一些进阶的GNN任务（如Graph Meta-learning）。tfg对Functional API有非常好的支持，因此可以灵活地使用tfg构建各种复杂的GNN模型。

Other API: 为了方便GNN的实现，tfg还提供了许多工具类，如数据集自动管理（现成数据集、数据集自动下载、解压、预处理等）、图采样等。

tf_geometric的图数据结构

tfg提供了非常友好且高效的图数据结构，可以支持高效的单图和多图GNN计算。值得注意的是，tfg提供的图数据结构同时支持TensorFlow和numpy，因此即使用户在使用PyG等PyTorch GNN框架，也可以使用tfg的图数据结构（使用numpy模式）来高效友好地进行单图、多图数据的组织。

下图简单展示了tfg的两种图数据结构：tfg.Graph（单图）和tfg.BatchGraph（多图），通过简单的语句tfg.BatchGraph.from_graphs([graph0, graph1, graph2])就可以轻松将三个单图虚拟成一个大图（tfg.BatchGraph），在大图上使用GCN、GAT等模型等价于在每个单图上分别这些GNN（即图mini-batch）。另外，可以在tfg.BatchGraph上直接执行各种图池化（Graph Pooling）操作，以并行对多个单图进行图池化。

tf_geometric的图map-reduce

许多GNN模型都可以看成是一个map-reduce计算流，其中map对应变换（Transformation）操作，reduce对应聚合（Aggregation）操作。实现高效GNN的关键之一是如何在不规整且稀疏的图数据上进行高效地聚合（基于邻节点的消息传播、图池化等）：

示例—消息传播：

示例—图池化：

明显，由于图数据的不规整性和稀疏性，上述操作难以用深度学习框架提供的常规的基于矩阵的聚合操作（如tf.reduce_mean）等实现。为此，tfg借助了TensorFlow中处理不规整稀疏数据的tf.math.segment_xxxx操作实现了各种高层的面向图数据的reduce操作。如下图所示，tf.math.segment_sum可以在不规整稀疏数据上根据索引号进行加法聚合。

tfg基于tf.math.segment_xxxx提供的若干基本接口，构建了一些高级reduce方法。例如tfg实现了segment_softmax方法，它可以在图数据中，为每个节点的邻节点进行softmax归一化操作（如下图所示）：

高效且友好的API

tf_geometric使用消息传递机制来实现图神经网络：相比于基于稠密矩阵的实现，它具有更高的效率；相比于基于稀疏矩阵的实现，它具有更友好的API。除此之外，tf_geometric还为复杂的图神经网络操作提供了简易优雅的API。下面的示例展现了使用tf_geometric构建一个图结构的数据，并使用多头图注意力网络（Multi-head GAT）对图数据进行处理的流程：

运行结果：

开箱即用的Demo

tfg提供了许多开箱即用的Demo，这些Demo不仅复现了原文中的效果，还具有非常高的执行效率（充分利用了tf.function的性能优化）：

一键安装

tfg对系统、环境等没有特殊的要求，一般能够运行Python 3.5+且能够运行TensorFlow的环境就可以使用pip一键安装tfg。

更多详细介绍

更多关于tf_geometric的详细介绍，可以访问项目的Github首页、开发文档和论文进行了解。

Github主页: https://github.com/CrawlScript/tf_geometric
开发文档: https://tf-geometric.readthedocs.io
论文: 《Efficient Graph Deep Learning in TensorFlow with tf_geometric》（https://arxiv.org/abs/2101.11552，已被ACM Multimedia 2021 Open Source Software Competition接收）

参考链接：

成为VIP会员查看完整内容