首页 » 技术分享 » 比DGL快14倍:PyTorch图神经网络库PyG上线了

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

 

机器之心编译

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

640?wx_fmt=png

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

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


简单易用

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


640?wx_fmt=png


速度快

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

640?wx_fmt=png

最高比 DGL 快 14 倍!


已实现方法多

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


640?wx_fmt=png640?wx_fmt=png


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世界大学学科排名重磅发布!

640?wx_fmt=png

喜欢就点一下「好看」呗~

转载自原文链接, 如需删除请联系管理员。

原文链接:比DGL快14倍:PyTorch图神经网络库PyG上线了,转载请注明来源!

0