torch-geometric

PyTorch Geometric (PyG)

PyG は、PyTorch 上に構築されたグラフニューラルネットワークの標準ライブラリです。グラフのデータ構造、60 種類以上の GNN レイヤーの実装、スケーラブルなミニバッチ学習、異種グラフへの対応を提供します。

インストール

torch-geometric 2.7.x（2025 年 10 月）でテスト済みです。Python 3.10 以上とPyTorch 2.6 以上が必要です。

# 1. Install PyTorch first (match your CUDA/CPU setup — see https://pytorch.org/get-started/locally/)
uv pip install torch

# 2. Core PyG (no extension wheels required for basic usage)
uv pip install torch_geometric

オプションの高速化演算子（pyg-lib、torch-scatter、torch-sparse、torch-cluster）は、基本的な PyG の使用には不要です（PyG 2.3 以降）。PyTorch と CUDA のバージョンを確認してから、PyG wheel index からバージョンに合った wheel をインストールしてください。

python -c "import torch; print(torch.__version__, torch.version.cuda)"
# Then install wheels for your torch+CUDA combo, e.g.:
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： PyTorch >2.5 向けの pyg conda チャンネルは保守されなくなりました。代わりに uv pip install と上記の wheel index を使用してください。

PyG 2.7 の注意事項

PyG 2.7 では Python 3.9 と PyTorch ≤2.5 のサポートが終了しました。PyTorch 2.6～2.8 の互換性表については、2.7.0 のリリースノートを参照してください。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 format, dtype=torch.long
    edge_attr=edge_features,   # [num_edges, num_edge_features]
    y=labels,                  # node-level [num_nodes, *] or graph-level [1, *]
    pos=positions,             # [num_nodes, num_dimensions] (for point clouds/spatial)
)

edge_index の形式は重要です。これは [2, num_edges] 形式のテンソルで、edge_index[0] はソースノード、edge_index[1] はターゲットノードです。タプルのリストではありません。エッジペアが行として並んでいる場合は、転置して .contiguous() を呼び出してください。

# If edges are [[src1, dst1], [src2, dst2], ...] — transpose first:
edge_index = edge_pairs.t().contiguous()

無向グラフでは両方向を含めます。エッジ (0,1) には、edge_index に [0,1] と [1,0] の両方が必要です。

異種グラフには HeteroData を使用してください。詳しくは以下の「異種グラフ」セクションを参照してください。

データセット

PyG には、ダウンロードと前処理が自動で行われる、多くの標準データセットが含まれています。

from torch_geometric.datasets import Planetoid, TUDataset

# Single-graph node classification (Cora, Citeseer, Pubmed)
dataset = Planetoid(root='./data', name='Cora')
data = dataset[0]  # single graph with train/val/test masks

# Multi-graph classification (ENZYMES, MUTAG, IMDB-BINARY, etc.)
dataset = TUDataset(root='./data', name='ENZYMES')
# dataset[0], dataset[1], ... are individual graphs

タスク別の一般的なデータセット：

• ノード分類：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

# Common transforms
T.NormalizeFeatures()    # Row-normalize node features to sum to 1
T.ToUndirected()         # Add reverse edges to make graph undirected
T.AddSelfLoops()         # Add self-loop edges
T.KNNGraph(k=6)          # Build k-NN graph from point cloud positions
T.RandomJitter(0.01)     # Random noise augmentation on positions
T.Compose([...])         # Chain multiple transforms

# Apply as pre_transform (once, saved to disk) or transform (every access)
dataset = ShapeNet(root='./data', pre_transform=T.KNNGraph(k=6),
                   transform=T.RandomJitter(0.01))

GNN モデルの構築

クイックスタート：組み込みレイヤーの使用

GNN を最も手早く構築する方法は、torch_geometric.nn の畳み込みレイヤーを積み重ねることです。

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 の畳み込みレイヤーには活性化関数が含まれていません。柔軟性を確保するための設計なので、各レイヤーの後に自分で適用してください。

畳み込みレイヤーの選択

タスクとグラフ構造に応じて選択してください。

レイヤー	最適な用途	主な考え方
GCNConv	同種グラフ、半教師ありノード分類	スペクトル理論に着想を得た、次数で正規化された集約
GATConv / GATv2Conv	近傍ノードの重要度が異なる場合	注意機構で重み付けされたメッセージ
SAGEConv	大規模グラフ、帰納的設定	サンプリングに適した学習可能な集約
GINConv	グラフ分類、表現力の最大化	WL テストと同等の能力
TransformerConv	豊富なエッジ特徴、複雑な相互作用	エッジ特徴を用いたマルチヘッド注意機構
EdgeConv	点群、動的グラフ	エッジ特徴 (x_i, x_j - x_i) に対する MLP
RGCNConv	多数の関係タイプを持つ異種グラフ	関係ごとの重み行列
HGTConv	異種グラフ	タイプごとの注意機構

すべての畳み込みレイヤーは、最低限 (x, edge_index) を受け取ります。多くのレイヤーでは、エッジ特徴を表す edge_attr も受け取ります。

遅延初期化

入力チャンネルに -1 を指定すると、PyG が次元を自動的に推測します。異種モデルで特に便利です。

conv = SAGEConv((-1, -1), 64)  # Input dims inferred on first forward pass
# Initialize lazy modules:
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 をサブクラス化します。フレームワークは次のとおりです。

1. propagate() がメッセージパッシングを調整します
2. message() が各エッジを流れる情報を定義します（phi 関数）
3. aggregate() が各ノードのメッセージを集約します（sum／mean／max）
4. 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", or "max"
        self.lin = torch.nn.Linear(in_channels, out_channels)

    def forward(self, x, edge_index):
        # Pre-processing before message passing
        x = self.lin(x)
        # Start message passing
        return self.propagate(edge_index, x=x)

    def message(self, x_j):
        # x_j: features of source nodes for each edge [num_edges, features]
        # The _j suffix auto-indexes source nodes, _i indexes target nodes
        return x_j

_i／_j の規則：propagate() に渡すテンソルには、message() のシグネチャで _i（対象／中心ノード）または _j（ソース／近傍ノード）を付けることで、自動的にインデックスを付与できます。そのため、propagate に x=... を渡すと、message() で x_i と x_j にアクセスできます。

GCN と EdgeConv の実装例の全容については、references/message_passing.md を参照してください。

タスク別のパターン

ノード分類

# Full-batch training on a single graph (e.g., 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()

# Evaluation — train(False) puts the model in inference mode (disables 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)  # [num_graphs_in_batch, hidden_ch]
        return self.lin(x)

# Training loop
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)

# Encode nodes, then score edges
z = model.encode(train_data.x, train_data.edge_index)
# Positive edges
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],     # Sample 15 neighbors in hop 1, 10 in hop 2
    batch_size=128,              # Number of seed nodes per batch
    input_nodes=data.train_mask, # Which nodes to sample from
    shuffle=True,
)

for batch in train_loader:
    batch = batch.to(device)
    out = model(batch.x, batch.edge_index)
    # Only use first batch_size nodes for loss (these are the seed nodes)
    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()

# Node features — indexed by node type string
data['user'].x = torch.randn(1000, 64)
data['movie'].x = torch.randn(500, 128)

# Edge indices — indexed by (src_type, edge_type, dst_type) triplet
data['user', 'rates', 'movie'].edge_index = torch.randint(0, 500, (2, 3000))
data['user', 'follows', 'user'].edge_index = torch.randint(0, 1000, (2, 5000))

# Access convenience dicts
data.x_dict        # {'user': tensor, 'movie': tensor}
data.edge_index_dict  # {('user','rates','movie'): tensor, ...}
data.metadata()    # ([node_types], [edge_types])

異種 GNN を構築する 3 つの方法

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')

# Now accepts dicts:
out = model(data.x_dict, data.edge_index_dict)

二部グラフの入力チャネルには (-1, -1) を使用します（ソースとターゲットの次元が異なる場合があります）。残りの処理は遅延初期化が行います。

2. HeteroConv ラッパー — エッジ型ごとに異なる畳み込みを使用します。

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', 'rev_writes', 'author'): GATConv((-1, -1), 64, add_self_loops=False),
}, aggr='sum')

3. HGTConv などのネイティブ異種演算子

from torch_geometric.nn import HGTConv
conv = HGTConv(hidden_channels, hidden_channels, data.metadata(), num_heads=4)

異種グラフに関する重要事項：

• 双方向のメッセージフローを実現するため、T.ToUndirected() を使用して逆向きのエッジ型を追加します
• 二部グラフの畳み込み層（ソース型と宛先型が異なる層）では、add_self_loops を無効にします。代わりにスキップ接続を使用します：conv(x, edge_index) + lin(x)
• HeteroData で NeighborLoader を使用する場合は、input_nodes に ('node_type', mask) のタプルを指定します
• num_neighbors には、細かく制御できるよう、エッジ型をキーとする辞書を指定できます

訓練ループや異種グラフでの NeighborLoader の使用方法を含む完全な例については、references/heterogeneous.md を参照してください。

カスタムデータセット

独自のデータを PyG に読み込む方法：

• 簡単（クラス不要）：Data オブジェクトを直接作成し、リストを DataLoader に渡します
• 再利用可能（RAM に収まる場合）：InMemoryDataset をサブクラス化し、raw_file_names、processed_file_names、download()、process() をオーバーライドします
• 大規模（ディスクベース）：Dataset をサブクラス化し、さらに len() と get() をオーバーライドします
• CSV から読み込む場合：pandas でノード／エッジのテーブルを読み込み、連続したインデックスへのマッピングを作成し、Data または HeteroData に組み立てます
• NetworkX から読み込む場合：from_networkx(G) で NetworkX グラフを直接変換します
• scipy の疎行列から読み込む場合：from_scipy_sparse_matrix(adj) で edge_index を抽出します

すべてのパターン、エンコーダーを使った CSV の読み込み、MovieLens の手順を含む詳しい例については、references/custom_datasets.md をお読みください。

説明可能性

PyG は GNN の予測を解釈するための torch_geometric.explain を提供しています。

from torch_geometric.explain import Explainer, GNNExplainer

explainer = Explainer(
    model=model,
    algorithm=GNNExplainer(epochs=200),
    explanation_type='model',
    node_mask_type='attributes',
    edge_mask_type='object',
    model_config=dict(
        mode='multiclass_classification',
        task_level='node',
        return_type='log_probs',
    ),
)

explanation = explainer(data.x, data.edge_index, index=10)
explanation.visualize_graph()           # Important subgraph
explanation.visualize_feature_importance(top_k=10)  # Feature importance

使用可能なアルゴリズム：GNNExplainer（最適化ベース）、PGExplainer（パラメトリック、学習可能）、CaptumExplainer（Captum による勾配ベース）、AttentionExplainer（アテンション重み）。同種グラフと異種グラフの両方で動作します。

すべてのアルゴリズム、異種グラフの説明、評価指標、PGExplainer のトレーニングについては、references/explainability.md をお読みください。

よくある落とし穴

1. edge_index の形状：[num_edges, 2] ではなく [2, num_edges] である必要があります。必要に応じて転置してください。
2. 活性化関数の付け忘れ：畳み込み層には ReLU などが含まれていないため、手動で追加してください。
3. 異種二部グラフの自己ループ：ソースと宛先のノードタイプが異なる場合は、add_self_loops=True を使わないでください。代わりにスキップ接続を使います。
4. NeighborLoader のスライス：最初の batch.batch_size 個のノードだけがシードノードです。予測とラベルをそれに応じてスライスしてください。
5. 無向グラフ：グラフが無向の場合は、edge_index に両方向のエッジを含めるか、T.ToUndirected() を使ってください。
6. 遅延初期化：入力チャネル数が -1 のモデルでは、パラメーターを初期化するために、トレーニング前に torch.no_grad() を使ってフォワードパスを 1 回実行する必要があります。
7. グラフタスクでのグローバルプーリング：ノード特徴量をグラフレベルに集約するには、手動で形状を変更せず global_mean_pool(x, batch) を使ってください。
8. num_neighbors の整合性：len(num_neighbors) を GNN 層の数と一致させてください。層数よりホップ数が多いと計算が無駄になり、少ないとモデルの能力を活用しきれません。

エージェントの作業手順

1. 環境を確認します。 Python 環境とライブラリのバージョン（python -c "import pkg; print(pkg.__version__)"）を確認し、データの形状、型、欠損値を調べます。
2. 入力を特定します。 データまたはユーザーに確認して、形式、識別子、パラメーターを確定します。結果を左右する値は推測せず、質問してください。
3. まず小規模に実行します。 まずサンプルまたは単一のフォールドで実行し、実行時間とメモリを確認します。
4. 上記の指示と参考資料に従って、タスク全体を実行します。
5. 結果を検証します。 ホールドアウトデータ、固定した乱数シード、適切な評価指標を使用し、データリークがないか確認します。不確実性（信頼区間、シード間の標準偏差）を報告します。
6. 報告します。 実行した内容（バージョン、コマンド、パラメーター）、確認した内容、未確定の事項を明記します。

発生した問題	対応
メモリ不足または実行が非常に遅い	サブサンプリング、チャンク処理または遅延計算を使用するか、モデルサイズを小さくし、変更内容をユーザーに伝えます。
これらの指示に記載された関数、フラグ、エンドポイントがインストール済みのバージョンに存在しない	インストール済みバージョンのドキュメント（help()、--help、公式ドキュメント）を確認して適宜対応し、ユーザーに伝えます。API を捏造してはいけません。
必須の入力、識別子、パラメーターが曖昧	ユーザーに質問するか、実行前に仮定を明示します。

誠実性に関するルール

• 結果、パラメーター、識別子、引用、統計値を捏造してはいけません。実行または検証できないことがあれば、はっきり伝えてください。
• このセッションで計算していない評価指標を報告してはいけません。すべての数値について、その算出に使ったコードの経路を示してください。
• ここに記載されているバージョン固有の情報は古くなっている可能性があるものとして扱い、インストール済みバージョンの公式ドキュメントで確認してください。
• 費用が発生する操作、共有 GPU やクラウドのクォータを消費する操作、個人データや患者データにアクセスする操作、取り消せない操作を行う前に、必ず確認してください。

関連スキル

• networkx：NetworkX を使って、Python で複雑なネットワークやグラフを作成、分析、可視化します。
• torchdrug：分子グラフ、特性予測、自己教師あり事前学習、分子生成のための TorchDrug 0.2.1 ワークフローを構築し、トラブルシューティングします…