新的PyTorch圖神經網絡庫，快了14倍：LeCun盛贊，GitHub 2000星
本文經AI新媒體量子位（公眾號ID:QbitAI）授權轉載，轉載請聯系出處。
“CNN已老，GNN當立！”
當科學家們發現，圖神經網絡 (GNN) 能搞定傳統CNN處理不了的非歐數據，從前深度學習解不開的許多問題都找到了鑰匙。
如今，有個圖網絡PyTorch庫，已在GitHub摘下2000多星，還被CNN的爸爸Yann LeCun翻了牌：
它叫PyTorch Geometric，簡稱PyG，聚集了26項圖網絡研究的代碼實現。
這個庫還很快，比起前輩DGL圖網絡庫，PyG***可以達到它的15倍速度。
應有盡有的庫
要跑結構不規則的數據，就用PyG吧。不管是圖形 (Graphs)，點云 (Point Clouds) 還是流形(Manifolds) 。
△ 右邊是不規則的，非歐空間
這是一個豐盛的庫：許多模型的PyTorch實現，各種有用的轉換 (Transforms) ，以及大量常見的benchmark數據集，應有盡有。
說到實現，包括Kipf等人的圖卷積網絡 (GCN) 和Bengio實驗室的圖注意力網絡 (GAT) 在內，2017-2019年各大頂會的 (至少) 26項圖網絡研究，這里都能找到快速實現。
到底能多快？PyG的兩位作者用英偉達GTX 1080Ti做了實驗。
對手DGL，也是圖網絡庫：
在四個數據集里，PyG全部比DGL跑得快。最懸殊的一場比賽，是在Cora數據集上運行GAT模型：跑200個epoch，對手耗時33.4秒，PyG只要2.2秒，相當于對方速度的15倍。
每個算法的實現，都支持了CPU計算和GPU計算。
食用方法
庫的作者，是兩位德國少年，來自多特蒙德工業大學。
△ 其中一位
他們說，有了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
