49

新的PyTorch图神经网络库,快了14倍:LeCun盛赞,GitHub 2000星

 5 years ago
source link: http://news.51cto.com/art/201903/593411.htm?amp%3Butm_medium=referral
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.

QvEjErZ.jpg!web

本文经AI新媒体量子位(公众号ID:QbitAI)授权转载,转载请联系出处。

“CNN已老,GNN当立!”

当科学家们发现, 图神经网络 (GNN) 能搞定传统CNN处理不了的非欧数据,从前深度学习解不开的许多问题都找到了钥匙。

如今,有个图网络PyTorch库,已在GitHub摘下 2000多星 ,还被CNN的爸爸 Yann LeCun 翻了牌:

67VVneA.jpg!web

它叫 PyTorch Geometric ,简称PyG,聚集了 26项 图网络研究的代码实现。

这个库还很 ,比起前辈DGL图网络库,PyG最高可以达到它的15倍速度。

应有尽有的库

要跑结构不规则的数据,就用PyG吧。不管是 图形 (Graphs), 点云 (Point Clouds) 还是 流形 (Manifolds) 。

vUriA3B.jpg!web

△ 右边是不规则的,非欧空间

这是一个丰盛的库:许多模型的 PyTorch实现 ,各种有用的 转换 (Transforms) ,以及大量常见的 benchmark数据集 ,应有尽有。

说到实现,包括Kipf等人的图卷积网络 ( GCN ) 和Bengio实验室的图注意力网络 ( GAT ) 在内,2017-2019年各大顶会的 (至少) 26项图网络研究,这里都能找到快速实现。

到底能多快?PyG的两位作者用英伟达GTX 1080Ti做了实验。

对手 DGL ,也是图网络库:

ZBJVZnv.jpg!web

在四个数据集里,PyG全部比DGL跑得快。最悬殊的一场比赛,是在Cora数据集上运行GAT模型:跑200个epoch,对手耗时 33.4秒 ,PyG只要 2.2秒 ,相当于对方速度的15倍。

每个算法的实现,都支持了CPU计算和GPU计算。

食用方法

库的作者,是两位德国少年,来自多特蒙德工业大学。

uAF7naA.jpg!web

△ 其中一位

他们说,有了PyG,做起图网络就像一阵微风。

你看,实现一个边缘卷积层 (Edge Convolution Layer) 只要这样而已:

 1import torch 
 2from torch.nn import Sequential as Seq, Linear as Lin, ReLU 
 3from torch_geometric.nn import MessagePassing 
 4 
 5class EdgeConv(MessagePassing): 
 6 def __init__(self, F_in, F_out): 
 7 super(EdgeConv, self).__init__() 
 8 self.mlp = Seq(Lin(2 * F_in, F_out), ReLU(), Lin(F_out, F_out)) 
 9 
10 def forward(self, x, edge_index): 
11 # x has shape [N, F_in] 
12 # edge_index has shape [2, E] 
13 return self.propagate(aggr='max', edge_index=edge_index, x=x) # shape [N, F_out] 
14 
15 def message(self, x_i, x_j): 
16 # x_i has shape [E, F_in] 
17 # x_j has shape [E, F_in] 
18 edge_features = torch.cat([x_i, x_j - x_i], dim=1) # shape [E, 2 * F_in] 
19 return self.mlp(edge_features) # shape [E, F_out] 

安装之前确认一下,至少要有PyTorch 1.0.0;再确认一下cuda/bin在$PATH里,cuda/include在$CPATH里:

1$ python -c "import torch; print(torch.__version__)" 
2>>> 1.0.0 
3 
4$ echo $PATH 
5>>> /usr/local/cuda/bin:... 
6 
7$ echo $CPATH 
8>>> /usr/local/cuda/include:... 

然后,就开始各种pip install吧。

PyG项目传送门:

https://github.com/rusty1s/pytorch_geometric

PyG主页传送门:

https://rusty1s.github.io/pytorch_geometric/build/html/index.html

PyG论文传送门:

https://arxiv.org/pdf/1903.02428.pdf


About Joyk


Aggregate valuable and interesting links.
Joyk means Joy of geeK