支持異構圖、集成GraphGym,超好用的圖神經網絡庫PyG更新2.0版本
當前最流行和廣泛使用的 GNN 庫 PyG(PyTorch Geometric)現在出 2.0 版本了,新版本提供了全面的異構圖支持、GraphGam 以及很多其他特性,這一系列改進,為使用者帶來了更好的用戶體驗。
PyTorch Geometric(PyG)是一個構建于 PyTorch 之上的庫,用來為一系列與結構化數據相關的應用編寫和訓練圖神經網絡(GNN)。PyG 對機器學習研究者以及機器學習工具包的首次使用者都很友好。
就其特征和功能而言,PyG 具有易用和統一的 API,用戶花費更少的時間在實現和運行圖神經網絡的底層機制上,只需要 10 至 20 行代碼就可以訓練自定義 GNN 模型。PyG 涵蓋了大量的 SOTA GNN 架構以及訓練和可擴展流程,并且易于擴展以匹配用戶自己的特定用例或進行自己的 GNN 研究。
此外,PyG 中的全新 GraphGym 使得用戶可以輕松地復現 GNN 實驗,執行和分析數千個 GNN 實驗,并將自定義模塊嵌入到 GNN 學習 pipeline 中。
近日,PyG 核心領導者(Core lead)Matthias Fey 宣布 PyG 2.0 版本來了!新版本提供了全面的異構圖支持、GraphGam 以及很多其他特性。

PyG 2.0 詳情地址:https://github.com/pyg-team/pytorch_geometric/releases/tag/2.0.0
PyG 2.0 新特性
PyG 2.0 是一個新的版本,提供了復雜的異構圖支持、GraphGym 集成以及其他新特性。
PyG 2.0 中提供了完全的異構圖支持。異構圖支持包括數據存儲層的完整重寫(同時保持向后兼容性)、異構圖轉換、通過鄰采樣的關系型數據加載例程,以及一整套異構 GNN 模型 / 示例。
突出亮點
異構圖存儲:異構圖現在可以存儲在它們自己的專用 data.HeteroData 類中。
- from torch_geometric.data import HeteroData
- data = HeteroData()
- # Create two node types "paper" and "author" holding a single feature matrix:
- data['paper'].x = torch.randn(num_papers, num_paper_features)
- data['author'].x = torch.randn(num_authors, num_authors_features)
- # Create an edge type ("paper", "written_by", "author") holding its graph connectivity:
- data['paper', 'written_by', 'author'].edge_index = ... # [2, num_edges]
異構 Mini-Batch 加載:異構圖可以分別通過 loader.DataLoader 和 loader.NextorLoader 對許多小的、單個的巨大圖進行轉化,轉換成 mini-batches。這些 loaders 現在可以處理同構圖和異構圖:
- from torch_geometric.loader import DataLoader
- loader = DataLoader(heterogeneous_graph_dataset, batch_size=32, shuffle=True)
- from torch_geometric.loader import NeighborLoader
- loader = NeighborLoader(heterogeneous_graph, num_neighbors=[30, 30], batch_size=128,
- input_nodes=('paper', data['paper'].train_mask), shuffle=True)
異構圖神經網絡:異構 GNN 現在可以通過 nn.to_hetero、nn.to_hetero_with_bases 從同構 GNN 輕松創建。這些進程采用現有的 GNN 模型并復制其消息功能,以考慮不同的節點和邊緣類型:
- from torch_geometric.nn import SAGEConv, to_hetero
- class GNN(torch.nn.Module):
- def __init__(hidden_channels, out_channels):
- super().__init__()
- self.conv1 = SAGEConv((-1, -1), hidden_channels)
- self.conv2 = SAGEConv((-1, -1), out_channels)
- def forward(self, x, edge_index):
- x = self.conv1(x, edge_index).relu()
- x = self.conv2(x, edge_index)
- return x
- model = GNN(hidden_channels=64, out_channels=dataset.num_classes)
- model = to_hetero(model, data.metadata(), aggr='sum')

使用 GraphGym 管理實驗
PyG 2.0 現在通過 torch_geometric.graphgym 正式支持 GraphGym。總的來說,GraphGym 是一個平臺,用于通過高度模塊化的 pipeline 從配置文件中設計和評估圖神經網絡:
- GraphGym 是開始學習標準化 GNN 實現和評估的最佳平臺;
- GraphGym 提供了一個簡單的接口來并行嘗試數千個 GNN 架構,以找到適合特定任務的最佳設計;
- GraphGym 可輕松進行超參數搜索并可視化哪些設計選擇更好。

重大改變
datasets.AMiner 數據集現在返回 data.HeteroData 對象;transforms.AddTrainValTestMask 已被替換為 transforms.RandomNodeSplit;由于 data.Data 的存儲布局發生了顯著變化以支持異構圖,因此需要通過刪除 root/processed 文件夾來重新處理已處理的數據集。
此外,data.Data.__cat_dim__ 、 data.Data.__inc__現在需要額外的參數輸入:
- def __cat_dim__(self, key, value, *args, **kwargs):
- pass
- def __inc__(self, key, value, *args, **kwargs):
- pass
如果你修改了自定義 data.Data 對象中的__cat_dim__或__inc__ ,請確保應用上述更改。
了解更多變化,請參考原項目。
核心團隊成員
PyG 核心團隊共有 6 位成員,包括多特蒙德工業大學四年級圖機器學習博士生 Matthias Fey、多特蒙德工業大學博士生 Jan Eric Lenssen(核心貢獻者)和圖網絡領域大牛 & 斯坦福計算機科學副教授 Jure Leskovec(擔任顧問)。
從左至右依次為 Matthias Fey、Jan Eric Lenssen 和 Jure Leskovec。
此外還有三位華人成員,分別是斯坦福大學計算機科學四年級博士生尤佳軒(Jiaxuan You,核心領導者)、斯坦福大學計算機科學博士生 Rex Ying(核心貢獻者,他將于 2022 年加入耶魯大學擔任助理教授)和卡內基梅隆大學 - 漢斯學院三年級博士生趙越(Yue Zhao,核心貢獻者)。
從左至右依次為尤佳軒、Rex Ying 和趙越。
日前,趙越在知乎分享了自己參與 PyG 2.0 版本設計與升級過程中的心得體會,感興趣的讀者可以參考閱讀。
知乎鏈接:https://www.zhihu.com/pin/1420674585365196800