比DGL快14倍：PyTorch图神经网络库PyG上线了

2019 年 3 月 15 日 机器学习算法与Python学习

机器之心编译

项目链接：https://github.com/rusty1s/pytorch_geometric

PyTorch Geometric 主要是现有模型快速重新实现的集合，其自定义稀疏/分散操作也非常方便。基于PyTorch构建，用于处理不规则结构化输入数据（如图、点云、流形）。除了一般的图形数据结构和处理方法外，它还包含从关系学习到3D数据处理等领域中最新发布的多种方法。通过利用稀疏 GPU 加速、提供专用的 CUDA 内核以及为不同大小的输入样本引入高效的小批量处理，PyTorch Geometric 实现了很高的数据吞吐量。

PyTorch Geometric（PyG）库包含易用的小批量加载器（mini-batch loader）、多GPU支持、大量常见基准数据集和有用的变换，适用于任意图像、三维网格（3D mesh）和点云。

简单易用

大大简化了实现图卷积网络的过程。比如，它可以用以下几行代码实现一个层（如edge convolution layer）：

速度快

下图展示了这一工具和其它图神经网络库的训练速度对比情况：

最高比 DGL 快 14 倍！

已实现方法多

目前已实现以下方法，所有实现方法均支持 CPU 和 GPU 计算：

PyG 概览

图神经网络（GNN）最近成为在图、点云和流形上进行表征学习的强大方法。与规则域中常用的卷积层和池化层概念类似，GNN通过传递、变换和聚合信息来（层级化地）提取局部嵌入。但是，实现GNN并不容易，因为它需要在不同大小的高度稀疏与不规则数据上实现较高的GPU吞吐量。PyTorch Geometric (PyG) 是基于Pytorch构建的几何深度学习扩展库。它可以利用专门的CUDA内核实现高性能。在简单的消息传递API之后，它将大多数近期提出的卷积层和池化层捆绑成一个统一的框架。所有的实现方法都支持 CPU 和 GPU 计算，并遵循不变的数据流范式，这种范式可以随着时间的推移动态改变图结构。PyG已在MIT许可证下开源，它具备完备的文档，且提供了附带教程和示例。

PyG为用户提供通用的MessagePassing接口，以便对新的研究想法进行快速干净的原型制作。此外，几乎所有近期提出的邻域聚合函数都适用于此接口，其中包括PyG已经集成的方法。

全局池化。PyG提供多种readout函数（如global add、mean 或 max pooling），从而支持图级别输出，而非节点级别输出。PyG还提供更加复杂的方法，如 set-to-set (Vinyals et al., 2016)、sort pooling (Zhang et al., 2018) 和全局软注意力层 (Li et al., 2016)。

层级池化（Hierarchical Pooling）。为进一步提取层级信息和使用更深层的GNN模型，需要以空间或数据依赖的方式使用多种池化方法。PyG目前提供Graclus、voxel grid pooling、迭代最远点采样算法（iterative farthest point sampling algorithm）的实现示例，以及可微池化机制（如DiffPool和top_k pooling）。

小批量处理。PyG 可自动创建单个（稀疏）分块对角邻接矩阵（block-diagonal adjacency matrix），并在节点维度中将特征矩阵级联起来，从而支持对多个（不同大小）图实例的小批量处理。正因如此，PyG可在不经修改的情况下应用邻域聚合方法，因为不相连的图之间不会出现信息交流。此外，自动生成的 assignment 向量可确保节点级信息不会跨图聚合，比如当执行全局聚合运算时。

处理数据集。PyG提供统一的数据格式和易用的接口，方便使用者创建和处理数据集，大型数据集和训练期间可保存在内存中的数据集皆可适用。要想创建新数据集，用户只需读取／下载数据，并转换为PyG数据格式即可。此外，用户可以使用变换（transform，即访问单独的多个图并对其进行变换）方法来修改数据集，比如数据增强、使用合成结构化图属性来增强节点特征等，从而基于点云自动生成图，或者从网格中自动采样点云。

推荐阅读

京东回应 995 工作制：不强求但鼓励全情投入

4个Linux小技巧，帮你提高机器学习生产效率

把14亿中国人都拉到一个微信群在技术上能实现吗？

2019QS世界大学学科排名重磅发布！