新的PyTorch圖神經(jīng)網(wǎng)絡庫，比前輩快14倍：LeCun盛贊，GitHub 2000星

skysun000001 2019-03-11

展開全文

“CNN已老，GNN當立！”

當科學家們發(fā)現(xiàn)，圖神經(jīng)網(wǎng)絡 (GNN) 能搞定傳統(tǒng)CNN處理不了的非歐數(shù)據(jù)，從前深度學習解不開的許多問題都找到了鑰匙。

如今，有個圖網(wǎng)絡PyTorch庫，已在GitHub摘下2000多星，還被CNN的爸爸Yann LeCun翻了牌：

它叫PyTorch Geometric，簡稱PyG，聚集了26項圖網(wǎng)絡研究的代碼實現(xiàn)。

這個庫還很快，比起前輩DGL圖網(wǎng)絡庫，PyG最高可以達到它的15倍速度。

應有盡有的庫

要跑結(jié)構(gòu)不規(guī)則的數(shù)據(jù)，就用PyG吧。不管是圖形 (Graphs)，點云 (Point Clouds) 還是流形 (Manifolds) 。

△ 右邊是不規(guī)則的，非歐空間

這是一個豐盛的庫：許多模型的PyTorch實現(xiàn)，各種有用的轉(zhuǎn)換 (Transforms) ，以及大量常見的benchmark數(shù)據(jù)集，應有盡有。

說到實現(xiàn)，包括Kipf等人的圖卷積網(wǎng)絡 (GCN) 和Bengio實驗室的圖注意力網(wǎng)絡 (GAT) 在內(nèi)，2017-2019年各大頂會的 (至少) 26項圖網(wǎng)絡研究，這里都能找到快速實現(xiàn)。

到底能多快？PyG的兩位作者用英偉達GTX 1080Ti做了實驗。

對手DGL，也是圖網(wǎng)絡庫：

在四個數(shù)據(jù)集里，PyG全部比DGL跑得快。最懸殊的一場比賽，是在Cora數(shù)據(jù)集上運行GAT模型：跑200個epoch，對手耗時33.4秒，PyG只要2.2秒，相當于對方速度的15倍。

每個算法的實現(xiàn)，都支持了CPU計算和GPU計算。

食用方法

庫的作者，是兩位德國少年，來自多特蒙德工業(yè)大學。

△ 其中一位

他們說，有了PyG，做起圖網(wǎng)絡就像一陣微風。

你看，實現(xiàn)一個邊緣卷積層 (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./pytorch_geometric/build/html/index.html

PyG論文傳送門：
https:///pdf/1903.02428.pdf

— 完 —