[论文解读] Fast Graph Representation Learning with PyTorch Geometric

Matthias Fey, Jan Eric Lenssen|arXiv (Cornell University)|Mar 6, 2019

Advanced Graph Neural Networks参考文献 45被引用 1,259

一句话总结

PyTorch Geometric 是一个基于 PyTorch、通过 CUDA 加速的快速、灵活的图表示学习库，提供统一的 MessagePassing API 和广泛的运算符支持。

ABSTRACT

We introduce PyTorch Geometric, a library for deep learning on irregularly structured input data such as graphs, point clouds and manifolds, built upon PyTorch. In addition to general graph data structures and processing methods, it contains a variety of recently published methods from the domains of relational learning and 3D data processing. PyTorch Geometric achieves high data throughput by leveraging sparse GPU acceleration, by providing dedicated CUDA kernels and by introducing efficient mini-batch handling for input examples of different size. In this work, we present the library in detail and perform a comprehensive comparative study of the implemented methods in homogeneous evaluation scenarios.

研究动机与目标

• 解决在不规则图数据（图、点云、流形）上实现高吞吐量 GNN 的挑战。
• 提供统一的、类似 PyTorch 的 API，以便快速原型化图神经网络。
• 提供由高效 CUDA 内核实现的丰富的邻域聚合、池化和数据集处理工具。

提出的方法

• 采用简单的消息传递形式，结合 gather/scatter 操作实现邻域聚合。
• 提供统一的 MessagePassing 接口，用户在其中实现 phi（message）和 gamma/update（aggregation）以及所选的 Dottedsquare（aggregation）函数。
• 实现广泛的 GNN 运算符（GCN、SGC、ARMA、GraphSAGE、GAT、AGNN、GIN、APPNP、DNA、signed operators）以及点云/流形变体（Relational GCN、PointNet++/PointCNN、MPNN、MoNet、SplineCNN、EdgeCNN）。
• 支持全局池化和分层池化，提供多种读出与池化选项（add/mean/max、Set2Set、SortPooling、全局注意力、DiffPool、top_k 池化）。
• 通过稀疏块对角邻接矩阵和连接的特征矩阵实现对不同大小图的迷你批处理，自动分配图以防止跨图泄漏。
• 提供易用的数据集处理，包括变换、自动数据集下载，以及对常见图数据集的基准测试。

Figure 1: Computation scheme of a GNN layer by leveraging gather and scatter methods based on edge indices $\bm{I}$ , hence alternating between node parallel space and edge parallel space.

实验结果

研究问题

• RQ1 在标准基准测试上，PyG 在广泛的 GNN 运算符上表现如何？
• RQ2 基于 gather/scatter 的实现是否能在稀疏、非规则图上提供比传统稀疏矩阵乘法更高的吞吐量？
• RQ3 与现有框架相比，PyG 在半监督节点分类、图分类和点云任务中的性能的实用性与可重复性如何？
• RQ4 各种池化和 Jumping Knowledge 风格的读出对图级任务的性能有何影响？
• RQ5 与其他图学习库相比，PyG 在运行时有哪些优势？

主要发现

• PyG 通过利用稀疏 GPU 加速和专用 CUDA 内核实现高数据吞吐量。
• 该库提供统一的、类似 PyTorch 的 API，通过 MessagePassing 接口实现对新研究想法的快速原型开发。
• 实现了广泛的图、点云和流形算子，覆盖了许多有影响力的 GNN 架构。
• 通过构造块对角邻接矩阵，支持不同大小的图的 mini-batch 处理，从而实现可扩展的训练。
• 实证评估表明，在半监督节点分类、图分类和三维数据处理任务中具有很强的可重复性和竞争力。
• 与 DGL 相比，在若干基准测试中 PyG 能实现显著更快的训练运行时间，展示了实际的速度优势。

Figure 2: Forward and backward runtimes of $1\,000$ runs of gather and scatter operations (GS) in comparison to sparse-matrix multiplications (SpMM) on Erdős Rényi graphs with $10\,000$ nodes and varying average node degrees. Runtimes are capped at two seconds due to visualization. We report runtime