torch-geometric

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 (基本的な使用には拡張ホイールは不要)
uv pip install torch_geometric

オプションの高速化演算 (pyg-lib、torch-scatter、torch-sparse、torch-cluster) は、基本的な PyG 使用には 必要ありません (PyG 2.3 以降)。PyTorch と CUDA のバージョンを確認した後、PyG ホイールインデックス からバージョンが一致したホイールをインストールします。

python -c "import torch; print(torch.__version__, torch.version.cuda)"
# 次に、torch+CUDA の組み合わせに合わせてホイールをインストール。例：
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 と上記のホイールインデックスを使用してください。

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 形式、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 (1回のみ、ディスクに保存) または 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	点群、動的グラフ	エッジ特徴量 (x_i, x_j - x_i) の MLP
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 をサブクラス化します。フレームワークは次のとおりです。

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"、または "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 にアクセスできます。

完全な GCN および EdgeConv 実装例については、references/message_passing.md を参照してください。

タスク固有のパターン

ノード分類

# 単一グラフでのフルバッチトレーニング (例：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) でモデルを推論モードに (ドロップアウト/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)

# トレーニングループ
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) はこれを使用してグラフごとに集約します。

リンク予測

エッジを train/val/test に分割し、ネガティブサンプリングを使用します。

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

# ノード特徴量 — ノードタイプ文字列でインデックス
データ['user'].x = torch.randn(1000, 64)
データ['movie'].x = torch.randn(500, 128)

# エッジインデックス — (src_type、edge_type、dst_type) の3つ組でインデックス
データ['user', 'rates', 'movie'].edge_index = torch.randint(0, 500, (2, 3000))
データ['user', 'follows', 'user'].edge_index = torch.randint(0, 1000, (2, 5000))

# アクセス便利辞書
データ.x_dict        # {'user': tensor, 'movie': tensor}
データ.edge_index_dict  # {('user','rates','movie'): tensor, ...}
データ.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')

# 辞書を受け入れるようになりました:
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),
    ('author', 'affiliated_with', 'institution'): GATConv((-1, -1), 64),
}, aggr='sum')

# フォワードパス
x_dict = conv(x_dict, edge_index_dict)

3. タイプ固有のモジュール — torch_geometric.nn には RGCNConv、HGTConv、HANConv などのレイヤーがあり、タイプ固有のパラメーター化をネイティブにサポートしています。

使い方と例については references/heterogeneous.md を参照してください。

エッジ特徴量を扱うためのその他のヒント

エッジ特徴量を GCNConv / SAGEConv に渡す

GCNConv と SAGEConv はデフォルトではエッジ特徴量を受け入れません。エッジ特徴量を含めるには TransformerConv またはカスタムの MessagePassing 実装を使用します。

エッジインデックスに属性を追加

エッジ属性がない場合、edge_index を torch_geometric.utils.to_undirected に渡すと、属性テンソルのサイズが保持されます (何もしません)。属性は個別に拡張する必要があります。

from torch_geometric.utils import to_undirected

edge_index, edge_attr = to_undirected(edge_index, edge_attr)

ミニバッチでのエッジ特徴量

DataLoader を使用する場合、バッチ内のサブグラフ間でエッジ特徴量の次元が一貫している必要があります。PyG はリストではなくテンソル形式を想定しているため、パディングは SliceDataset またはカスタムの collate 関数で処理する必要があります。

完全なガイドと例

• references/quick_start.md — Cora から始める完全な例
• references/link_prediction.md — GAE/VGAE オートエンコーダ、完全なトレーニングループ、LinkNeighborLoader
• references/heterogeneous.md — 異種グラフ変換、to_hetero()、RGCN、HGT の例
• references/scaling.md — NeighborLoader、ClusterLoader、GraphSAINT、ShaDowKHopSampler、マルチ GPU
• references/message_passing.md — カスタム MessagePassing レイヤーのチュートリアル