PyTorch Geometric (PyG)
PyG 是基於 PyTorch 構建的圖神經網路標準函式庫。它提供了圖的資料結構、超過 60 種 GNN 層實作、可擴展的小批次訓練以及對異質圖的支援。
安裝
已在 torch-geometric 2.7.x（2025 年 10 月）上測試。需要 Python 3.10+ 和 PyTorch 2.6+。
# 1. 先安裝 PyTorch（配對你的 CUDA/CPU 設定 — 請參閱 https://pytorch.org/get-started/locally/）
uv pip install torch
# 2. 核心 PyG（基本使用不需要擴展 wheel 套件）
uv pip install torch_geometric
可選的加速運算（pyg-lib、torch-scatter、torch-sparse、torch-cluster）不需要用於基本 PyG 使用（自 PyG 2.3 起）。在檢查你的 PyTorch 和 CUDA 版本後，從 PyG wheel 索引 安裝版本匹配的 wheel 套件：
python -c "import torch; print(torch.__version__, torch.version.cuda)"
# 然後為你的 torch+CUDA 組合安裝 wheel 套件，例如：
uv pip install pyg-lib torch-scatter torch-sparse torch-cluster \
-f https://data.pyg.org/whl/torch-2.8.0+cu128.html
檢查你的版本：
import torch_geometric
print(torch_geometric.__version__)
Conda： pyg conda 頻道已不再維護 PyTorch >2.5 — 請改用 uv pip install 和上面的 wheel 索引。
PyG 2.7 注意事項
PyG 2.7 已棄用 Python 3.9 和 PyTorch ≤2.5。請參閱 2.7.0 版本發行說明 以了解 PyTorch 2.6–2.8 的相容性表格。torch_geometric.distributed 已棄用 — 請改用標準的 torch.distributed DDP（請參閱 references/scaling.md）。
核心概念
圖資料：Data 和 HeteroData
圖存在於 Data 物件中。關鍵屬性：
from torch_geometric.data import Data
data = Data(
x=node_features, # [num_nodes, num_node_features]
edge_index=edge_index, # [2, num_edges] — COO 格式，dtype=torch.long
edge_attr=edge_features, # [num_edges, num_edge_features]
y=labels, # 節點層級 [num_nodes, *] 或圖層級 [1, *]
pos=positions, # [num_nodes, num_dimensions]（用於點雲/空間資料）
)
edge_index 格式至關重要：它是一個 [2, num_edges] 張量，其中 edge_index[0] = 來源節點，edge_index[1] = 目標節點。它不是元組列表。如果你有邊對作為行，請轉置並呼叫 .contiguous()：
# 如果邊是 [[src1, dst1], [src2, dst2], ...] — 先轉置：
edge_index = edge_pairs.t().contiguous()
對於無向圖，包括兩個方向：邊 (0,1) 需要在 edge_index 中同時包含 [0,1] 和 [1,0]。
對於異質圖，請使用 HeteroData — 請參閱下面的異質圖章節。
資料集
PyG 捆綁了許多標準資料集，這些資料集會自動下載和預處理：
from torch_geometric.datasets import Planetoid, TUDataset
# 單圖節點分類（Cora、Citeseer、Pubmed）
dataset = Planetoid(root='./data', name='Cora')
data = dataset[0] # 帶有 train/val/test 遮罩的單個圖
# 多圖分類（ENZYMES、MUTAG、IMDB-BINARY 等）
dataset = TUDataset(root='./data', name='ENZYMES')
# dataset[0]、dataset[1] ... 是各個圖
按任務分類的常見資料集：
- 節點分類：Planetoid（Cora/Citeseer/Pubmed）、OGB（ogbn-arxiv、ogbn-products、ogbn-mag）
- 圖分類：TUDataset（MUTAG、ENZYMES、PROTEINS、IMDB-BINARY）、OGB（ogbg-molhiv）
- 連結預測：OGB（ogbl-collab、ogbl-citation2）
- 分子：QM7、QM9、MoleculeNet
- 點雲/網格：ShapeNet、ModelNet10/40、FAUST
轉換
轉換預處理或增強圖資料，類似於 torchvision 轉換：
import torch_geometric.transforms as T
# 常見轉換
T.NormalizeFeatures() # 按行正規化節點特徵，使其總和為 1
T.ToUndirected() # 添加反向邊以使圖成為無向圖
T.AddSelfLoops() # 添加自環邊
T.KNNGraph(k=6) # 從點雲位置構建 k-NN 圖
T.RandomJitter(0.01) # 在位置上進行隨機噪聲增強
T.Compose([...]) # 鏈式組合多個轉換
# 應用為 pre_transform（僅執行一次，保存到磁碟）或 transform（每次存取時執行）
dataset = ShapeNet(root='./data', pre_transform=T.KNNGraph(k=6),
transform=T.RandomJitter(0.01))
建構 GNN 模型
快速入門：使用內建層
建構 GNN 最快的方法 — 從 torch_geometric.nn 堆疊 conv 層：
import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv
class GCN(torch.nn.Module):
def __init__(self, in_channels, hidden_channels, out_channels):
super().__init__()
self.conv1 = GCNConv(in_channels, hidden_channels)
self.conv2 = GCNConv(hidden_channels, out_channels)
def forward(self, x, edge_index):
x = self.conv1(x, edge_index).relu()
x = F.dropout(x, p=0.5, training=self.training)
x = self.conv2(x, edge_index)
return x
重要提示：PyG conv 層不包含激活函數 — 請在每層之後自行應用。這是出於靈活性的設計。
選擇 Conv 層
根據你的任務和圖結構選擇：
| 層 | 最適合 | 關鍵理念 |
|---|---|---|
GCNConv | 同質、半監督節點分類 | 頻譜啟發、度歸一化聚合 |
GATConv / GATv2Conv | 當鄰居重要性不同時 | 注意力加權訊息 |
SAGEConv | 大型圖、歸納式設定 | 便於採樣、可學習的聚合 |
GINConv | 圖分類、最大化表達力 | 與 WL 測試同樣強大 |
TransformerConv | 豐富的邊特徵、複雜交互 | 具有邊特徵的多頭注意力 |
EdgeConv | 點雲、動態圖 | 對邊特徵進行 MLP (x_i, x_j - x_i) |
RGCNConv | 具有多種關係類型的異質圖 | 關係專用權重矩陣 |
HGTConv | 異質圖 | 類型專用注意力 |
所有 conv 層至少接受 (x, edge_index)。許多也接受 edge_attr 作為邊特徵。
懶惰初始化
使用 -1 作為輸入通道數，讓 PyG 自動推斷維度 — 這對於異質模型特別有用：
conv = SAGEConv((-1, -1), 64) # 輸入維度在第一次前向傳播時推斷
# 初始化懶惰模組：
with torch.no_grad():
out = model(data.x, data.edge_index)
高階模型 API
對於常見架構，PyG 提供現成的模型類別：
from torch_geometric.nn import GraphSAGE, GCN, GAT, GIN
model = GraphSAGE(
in_channels=dataset.num_features,
hidden_channels=64,
out_channels=dataset.num_classes,
num_layers=2,
)
透過 MessagePassing 自訂層
若要實作新穎的 GNN 層，請子類化 MessagePassing。框架是：
propagate()協調訊息傳遞message()定義沿每條邊流動的資訊（phi 函數）aggregate()在每個節點處組合訊息（sum/mean/max）update()轉換聚合結果（gamma 函數）
from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops, degree
class MyConv(MessagePassing):
def __init__(self, in_channels, out_channels):
super().__init__(aggr='add') # "add"、"mean" 或 "max"
self.lin = torch.nn.Linear(in_channels, out_channels)
def forward(self, x, edge_index):
# 訊息傳遞前的預處理
x = self.lin(x)
# 開始訊息傳遞
return self.propagate(edge_index, x=x)
def message(self, x_j):
# x_j：每條邊的來源節點特徵 [num_edges, features]
# _j 後綴會自動索引來源節點，_i 索引目標節點
return x_j
_i / _j 慣例：傳遞給 propagate() 的任何張量都可以透過在 message() 簽名中附加 _i（目標/中心節點）或 _j（來源/鄰居節點）來自動索引。因此，如果你傳遞 x=... 給 propagate，你可以在 message() 中存取 x_i 和 x_j。
閱讀 references/message_passing.md 以了解完整的 GCN 和 EdgeConv 實作範例。
任務特定模式
節點分類
# 在單個圖上進行全批次訓練（例如 Cora）
model.train()
for epoch in range(200):
optimizer.zero_grad()
out = model(data.x, data.edge_index)
loss = F.cross_entropy(out[data.train_mask], data.y[data.train_mask])
loss.backward()
optimizer.step()
# 評估 — train(False) 將模型設為推論模式（停用 dropout/BN）
model.train(False)
pred = model(data.x, data.edge_index).argmax(dim=1)
acc = (pred[data.test_mask] == data.y[data.test_mask]).float().mean()
圖分類
多個圖 — 使用 DataLoader 進行小批次處理，使用全域池化來獲取圖層級表示：
from torch_geometric.loader import DataLoader
from torch_geometric.nn import GCNConv, global_mean_pool
loader = DataLoader(dataset, batch_size=32, shuffle=True)
class GraphClassifier(torch.nn.Module):
def __init__(self, in_ch, hidden_ch, out_ch):
super().__init__()
self.conv1 = GCNConv(in_ch, hidden_ch)
self.conv2 = GCNConv(hidden_ch, hidden_ch)
self.lin = torch.nn.Linear(hidden_ch, out_ch)
def forward(self, x, edge_index, batch):
x = self.conv1(x, edge_index).relu()
x = self.conv2(x, edge_index).relu()
x = global_mean_pool(x, batch) # [batch_size_in_batch, hidden_ch]
return self.lin(x)
# 訓練迴圈
for data in loader:
out = model(data.x, data.edge_index, data.batch)
loss = F.cross_entropy(out, data.y)
PyG 的 DataLoader 透過建立區塊對角鄰接矩陣來批次處理多個圖。batch 張量將每個節點映射到其圖索引。池化運算（global_mean_pool、global_max_pool、global_add_pool）使用它來按圖聚合。
連結預測
將邊分割為訓練/驗證/測試集，使用負採樣：
from torch_geometric.transforms import RandomLinkSplit
transform = RandomLinkSplit(
num_val=0.1,
num_test=0.1,
is_undirected=True,
add_negative_train_samples=False,
)
train_data, val_data, test_data = transform(data)
# 編碼節點，然後對邊進行評分
z = model.encode(train_data.x, train_data.edge_index)
# 正邊
pos_score = (z[train_data.edge_label_index[0]] * z[train_data.edge_label_index[1]]).sum(dim=1)
閱讀 references/link_prediction.md 以了解完整的連結預測指南：GAE/VGAE 自動編碼器、完整訓練迴圈、用於大型圖的 LinkNeighborLoader、異質連結預測以及評估指標。
擴展到大型圖
對於無法放入 GPU 記憶體的圖，透過 NeighborLoader 使用鄰居採樣：
from torch_geometric.loader import NeighborLoader
train_loader = NeighborLoader(
data,
num_neighbors=[15, 10], # 在第 1 跳採樣 15 個鄰居，在第 2 跳採樣 10 個
batch_size=128, # 每批次的種子節點數
input_nodes=data.train_mask, # 要從哪些節點採樣
shuffle=True,
)
for batch in train_loader:
batch = batch.to(device)
out = model(batch.x, batch.edge_index)
# 只使用前 batch_size 個節點計算損失（這些是種子節點）
loss = F.cross_entropy(out[:batch.batch_size], batch.y[:batch.batch_size])
關於 NeighborLoader 的重點：
num_neighbors列表長度應與 GNN 深度（訊息傳遞層數）匹配- 種子節點始終是輸出中的前
batch.batch_size個節點 batch.n_id將重新標記的索引映射回原始節點 ID- 同時適用於
Data和HeteroData - 對於連結預測，請改用
LinkNeighborLoader - 採樣超過 2-3 跳通常是不可行的（指數級增長）
其他可擴展性選項：ClusterLoader（ClusterGCN）、GraphSAINTSampler、ShaDowKHopSampler。對於多 GPU 訓練、DDP、PyTorch Lightning 整合和 torch.compile 支援，請閱讀 references/scaling.md。
異質圖
對於具有多種節點和邊類型的圖（社交網路、知識圖譜、推薦）：
from torch_geometric.data import HeteroData
data = HeteroData()
# 節點特徵 — 按節點類型字串索引
data['user'].x = torch.randn(1000, 64)
data['movie'].x = torch.randn(500, 128)
# 邊索引 — 按 (src_type, edge_type, dst_type) 三元組索引
data['user', 'rates', 'movie'].edge_index = torch.randint(0, 500, (2, 3000))
data['user', 'follows', 'user'].edge_index = torch.randint(0, 1000, (2, 5000))
# 存取便捷字典
data.x_dict # {'user': tensor, 'movie': tensor}
data.edge_index_dict # {('user','rates','movie'): tensor, ...}
data.metadata() # ([node_types], [edge_types])
建立異質 GNN 的三種方法
1. 使用 to_hetero() 自動轉換 — 編寫同質模型，自動轉換：
from torch_geometric.nn import SAGEConv, to_hetero
class GNN(torch.nn.Module):
def __init__(self, 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(64, dataset.num_classes)
model = to_hetero(model, data.metadata(), aggr='sum')
# 現在接受字典：
out = model(data.x_dict, data.edge_index_dict)
對二分圖輸入通道使用 (-1, -1)（來源、目標可能不同）。懶惰初始化會處理其餘部分。
2. HeteroConv 包裝器 — 每種邊類型使用不同的 conv：
from torch_geometric.nn import HeteroConv, GCNConv, SAGEConv, GATConv
conv = HeteroConv({
('paper', 'cites', 'paper'): GCNConv(-1, 64),
('author', 'writes', 'paper'): SAGEConv((-1, -1), 64),
('paper', 'written_by', 'author'): SAGEConv((-1, -1), 64),
}, aggr='sum')
# 在 forward 中：
out_dict = conv(x_dict, edge_index_dict)
3. 手動 HGTConv — 用於異質的專用轉換器層：
from torch_geometric.nn import HGTConv
class HGT(torch.nn.Module):
def __init__(self, hidden_channels, out_channels, num_heads, num_node_types, num_edge_types):
super().__init__()
self.lin_dict = torch.nn.ModuleDict()
for node_type in num_node_types:
self.lin_dict[node_type] = torch.nn.Linear(-1, hidden_channels)
self.conv = HGTConv(hidden_channels, hidden_channels, num_heads, num_node_types, num_edge_types)
self.lin = torch.nn.Linear(hidden_channels, out_channels)
def forward(self, x_dict, edge_index_dict):
x_dict = {node_type: self.lin_dict[node_type](x).relu()
for node_type, x in x_dict.items()}
x_dict = self.conv(x_dict, edge_index_dict)
return {node_type: self.lin(x) for node_type, x in x_dict.items()}
閱讀 references/heterogeneous.md 以了解完整的異質指南：to_hetero 轉換器的詳細資訊、對異質圖的 NeighborLoader、連結預測和元路徑。
常見陷阱
| 陷阱 | 解決方案 |
|---|---|
edge_index 形狀錯誤 | 始終為 [2, num_edges]，dtype=torch.long。如果形狀是 [num_edges, 2]，請轉置。 |
在 forward() 中忘記激活函數 | PyG conv 層不會為你應用激活函數。始終 在 conv 之後添加 .relu() 或同等函數。 |
非連續的 edge_index | 建立後立即呼叫 .contiguous()：edge_index = edge_pairs.t().contiguous() |
| 使用 DataLoader 進行節點分類（單一圖） | 不要 — 除非使用采樣器，否則節點分類使用全批次。 |
| 異質模型的通道數不匹配 | 對二分圖層使用 (-1, -1) 並讓懶惰初始化工作。請勿硬編碼通道數。 |
在 hetero 模型上呼叫 to_hetero() 兩次 | to_hetero() 期望一個同質模型。如果你有異質模型，請手動編寫。 |
| 與現有 PyG 程式碼的 CUDA 不匹配 | 始終在安裝 PyG 之前安裝匹配 CUDA 版本的 PyTorch。 |
參考資料
以下 references/ 檔案包含完整主題的詳細程式碼範例：
references/message_passing.md— 自訂 GNN 層：GCN、EdgeConv 和基於注意力的層的逐步實作 |references/scaling.md— 可擴展性策略：ClusterLoader、SAINT、ShaDow、sampling 最佳實務、torch.compile、DDP、PyTorch Lightning |references/link_prediction.md— 完整連結預測：GAE/VGAE、負採樣、大規模圖的 LinkNeighborLoader、異質連結預測 |references/heterogeneous.md— 異質圖：to_hetero 轉換器的詳細資訊、異質 NeighborLoader、連結預測、元路徑 |references/transforms.md— 轉換：資料增強、預處理鏈、虛擬節點、大型資料集的大規模預處理 |references/datasets.md— 資料集：OGB、自訂資料集、InMemoryDataset 與 Dataset、下載器、多處理
官方資源
- PyG 文件
- PyG GitHub
- PyG 範例 — 用於常見任務的現成訓練腳本
- PyG 論文
- OGB: Open Graph Benchmark
- OGB 排行榜