[论文解读] Efficient Graph Deep Learning in TensorFlow with tf_geometric

Jun Hu, Shengsheng Qian+5 more|arXiv (Cornell University)|Jan 27, 2021

Advanced Graph Neural Networks参考文献 34被引用 9

一句话总结

本文介绍了 tf_geometric，一个与 TensorFlow 兼容的高效图深度学习库，支持 TensorFlow 1.x 和 2.x。该库提供了针对图操作（如消息传递、图池化和小批量训练）的优化内核库，同时集成了主流 GNN 的实现，在节点分类和图分类任务中性能与官方基准持平。

ABSTRACT

We introduce tf_geometric, an efficient and friendly library for graph deep learning, which is compatible with both TensorFlow 1.x and 2.x. tf_geometric provides kernel libraries for building Graph Neural Networks (GNNs) as well as implementations of popular GNNs. The kernel libraries consist of infrastructures for building efficient GNNs, including graph data structures, graph map-reduce framework, graph mini-batch strategy, etc. These infrastructures enable tf_geometric to support single-graph computation, multi-graph computation, graph mini-batch, distributed training, etc.; therefore, tf_geometric can be used for a variety of graph deep learning tasks, such as transductive node classification, inductive node classification, link prediction, and graph classification. Based on the kernel libraries, tf_geometric implements a variety of popular GNN models for different tasks. To facilitate the implementation of GNNs, tf_geometric also provides some other libraries for dataset management, graph sampling, etc. Different from existing popular GNN libraries, tf_geometric provides not only Object-Oriented Programming (OOP) APIs, but also Functional APIs, which enable tf_geometric to handle advanced graph deep learning tasks such as graph meta-learning. The APIs of tf_geometric are friendly, and they are suitable for both beginners and experts. In this paper, we first present an overview of tf_geometric's framework. Then, we conduct experiments on some benchmark datasets and report the performance of several popular GNN models implemented by tf_geometric.

研究动机与目标

• 解决在稀疏和不规则图数据上高效且可扩展地实现图神经网络（GNN）的挑战。
• 提供一个统一的高性能图深度学习框架，支持 TensorFlow 1.x 和 2.x。
• 支持多种 GNN 任务，包括归纳式和直推式节点分类、链接预测和图分类。
• 提供灵活的 API（面向对象和函数式），以满足初学者和高级用户在复杂 GNN 工作负载中的需求。
• 通过自定义内核库和高效数据结构，优化核心 GNN 操作，如消息聚合和图池化。

提出的方法

• 设计图数据结构和映射-归约框架，以高效处理稀疏和不规则图数据。
• 实现图小批量策略，通过高效的采样和批处理支持大规模图的训练。
• 引入内核库栈，支持单图、多图和分布式训练工作负载。
• 同时支持面向对象和函数式 API，以支持高级用例，如图元学习和模型组合。
• 提供核心 GNN 组件的优化实现，包括消息传递和分层图池化（如 DiffPool、SAGPool h）。
• 集成辅助工具用于数据集管理、图采样和交叉验证，以简化模型评估。

实验结果

研究问题

• RQ1 如何在 TensorFlow 上高效实现稀疏和不规则图数据的图神经网络？
• RQ2 在单图和多图设置下，支持可扩展 GNN 训练需要哪些系统级优化？
• RQ3 一个同时支持面向对象和函数式 API 的统一库，能否有效支持初学者友好和高级的 GNN 研究？
• RQ4 在标准基准上，tf_geometric 实现的 GNN 模型性能与官方实现相比如何？
• RQ5 在图分类任务中，架构修改（如均值-最大池化）能在多大程度上提升性能？

主要发现

• tf_geometric 在节点分类任务中实现了具有竞争力的性能，GCN、GAT、SGC、APPNP 和 DGI 模型在 Cora、CiteSeer 和 PubMed 数据集上的结果与官方结果匹配或非常接近。
• 在 Cora 数据集上，tf_geometric 的 GAT 模型达到 83.0±0.8% 的准确率，与 Velickovic 等人（2018）报告的 83.0% 一致。
• 在图分类任务中，均值-最大池化模型在 NCI1 数据集上达到 76.03±0.7% 的准确率，在某些情况下优于更复杂的模型（如 Set2Set 和 SortPool）。
• 该库的优化架构使性能优于基线实现，尤其是在使用多层均值-最大池化等增强池化策略时。
• 面向对象和函数式 API 的集成支持灵活建模，可实现元学习和模块化 GNN 设计等高级任务。
• 该库成功支持分布式训练和多图计算，在多种 GNN 工作负载中表现出良好的可扩展性。