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

机器之心编译

项目链接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,即访问单独的多个图并对其进行变换)方法来修改数据集,比如数据增强、使用合成结构化图属性来增强节点特征等,从而基于点云自动生成图,或者从网格中自动采样点云。

原文发布于微信公众号 - 机器学习算法与Python学习(guodongwei1991)

原文发表时间:2019-03-16

本文参与腾讯云自媒体分享计划,欢迎正在阅读的你也加入,一起分享。

发表于

我来说两句

0 条评论
登录 后参与评论

扫码关注云+社区

领取腾讯云代金券