torch-geometric-graph-neural-networks

PyTorch Geometric (PyG) — グラフニューラルネットワーク

概要

PyTorch Geometric は、グラフニューラルネットワーク（GNN）の開発と学習のために構築された PyTorch ベースのライブラリです。40 種類以上の畳み込み層、ブロック対角隣接行列によるミニバッチ処理、大規模グラフ向けの近傍サンプリング、および複数のノード／エッジタイプを持つネットワーク向けの異種グラフ対応を提供します。

使用する場面

• 引用、ソーシャル、または生物学的ネットワークでのノード分類
• グラフレベルの分類（分子活性、タンパク質の機能）
• リンク予測（知識グラフ、レコメンデーションシステム）
• 分子特性の予測（創薬、量子化学）
• 3D 点群処理とメッシュ解析
• 近傍サンプリングによる大規模グラフ学習（ノード数 100K 超）
• 複数のノード／エッジタイプを持つ異種グラフ
• グラフ以外の深層学習には → PyTorch を直接使用してください
• 従来のグラフアルゴリズム（最短経路、中心性）には → NetworkX を使用してください

前提条件

pip install torch torch_geometric
# Optional sparse operations (recommended):
# pip install pyg_lib torch_scatter torch_sparse torch_cluster

import torch
import torch.nn.functional as F
from torch_geometric.data import Data
from torch_geometric.nn import GCNConv

クイックスタート

from torch_geometric.datasets import Planetoid
from torch_geometric.nn import GCNConv
import torch, torch.nn.functional as F

dataset = Planetoid(root='/tmp/Cora', name='Cora')
data = dataset[0]

class GCN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = GCNConv(dataset.num_features, 16)
        self.conv2 = GCNConv(16, dataset.num_classes)
    def forward(self, data):
        x = F.relu(self.conv1(data.x, data.edge_index))
        return self.conv2(x, data.edge_index)

model = GCN()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
for epoch in range(200):
    model.train(); optimizer.zero_grad()
    F.cross_entropy(model(data)[data.train_mask], data.y[data.train_mask]).backward()
    optimizer.step()

model.eval()
pred = model(data).argmax(dim=1)
acc = (pred[data.test_mask] == data.y[data.test_mask]).float().mean()
print(f'Test Accuracy: {acc:.4f}')  # ~0.81

コア API

1. データ表現

import torch
from torch_geometric.data import Data

# Create a graph: 3 nodes, 4 edges (undirected)
edge_index = torch.tensor([[0, 1, 1, 2],
                           [1, 0, 2, 1]], dtype=torch.long)
x = torch.randn(3, 16)  # Node features [num_nodes, features]
y = torch.tensor([0, 1, 0])  # Node labels

data = Data(x=x, edge_index=edge_index, y=y)
print(f'Nodes: {data.num_nodes}, Edges: {data.num_edges}')
print(f'Features: {data.num_node_features}')
print(f'Has self-loops: {data.has_self_loops()}')
print(f'Is undirected: {data.is_undirected()}')

# Optional attributes
data.edge_attr = torch.randn(4, 8)   # Edge features [num_edges, features]
data.pos = torch.randn(3, 3)          # Node positions (3D)
data.train_mask = torch.tensor([True, True, False])  # Custom masks

# Mini-batch processing — graphs concatenated as block-diagonal
from torch_geometric.loader import DataLoader

loader = DataLoader(dataset, batch_size=32, shuffle=True)
for batch in loader:
    print(f'Graphs: {batch.num_graphs}, Nodes: {batch.num_nodes}')
    # batch.batch maps each node → its source graph index
    # No padding needed — computationally efficient

2. 畳み込み層

from torch_geometric.nn import GCNConv, GATConv, SAGEConv, GINConv
import torch.nn as nn

# GCNConv — spectral graph convolution (baseline)
conv = GCNConv(in_channels=16, out_channels=32)
# Supports: edge_weight, SparseTensor, Bipartite, Lazy init

# GATConv — attention-based neighbor weighting
conv = GATConv(16, 32, heads=8, dropout=0.6)
# Output: [N, heads * out_channels] (concat) or [N, out_channels] (concat=False)

# SAGEConv — inductive learning via sampling
conv = SAGEConv(16, 32, aggr='mean')  # 'mean', 'max', 'lstm'

# GINConv — maximally powerful for graph isomorphism
nn_module = nn.Sequential(nn.Linear(16, 32), nn.ReLU(), nn.Linear(32, 32))
conv = GINConv(nn_module)

# TransformerConv — graph transformer
from torch_geometric.nn import TransformerConv
conv = TransformerConv(16, 32, heads=8, beta=True)

# All layers: x_out = conv(x, edge_index)
x_out = conv(x, edge_index)
print(f'Output shape: {x_out.shape}')  # [num_nodes, out_channels]

3. カスタムメッセージパッシング

from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops, degree

class CustomConv(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):
        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))
        x = self.lin(x)

        # Degree-based normalization
        row, col = edge_index
        deg = degree(col, x.size(0), dtype=x.dtype)
        norm = deg.pow(-0.5)
        norm = norm[row] * norm[col]

        return self.propagate(edge_index, x=x, norm=norm)

    def message(self, x_j, norm):
        # x_j: source node features (automatic via _j suffix)
        return norm.view(-1, 1) * x_j

# Key methods: forward(), message(), aggregate(), update()
# _i suffix → target node, _j suffix → source node

4. プーリングとグラフレベルの読み出し

from torch_geometric.nn import (
    global_mean_pool, global_max_pool, global_add_pool,
    TopKPooling, SAGPooling
)

# Global pooling: node features → graph-level representation
x_graph = global_mean_pool(x, batch)  # [num_graphs, features]

# Hierarchical pooling: coarsen graph
pool = TopKPooling(64, ratio=0.8)  # Keep top 80% nodes
x, edge_index, _, batch, _, _ = pool(x, edge_index, None, batch)

# Graph classification model
class GraphClassifier(torch.nn.Module):
    def __init__(self, num_features, num_classes):
        super().__init__()
        self.conv1 = GCNConv(num_features, 64)
        self.conv2 = GCNConv(64, 64)
        self.pool = TopKPooling(64, ratio=0.8)
        self.lin = torch.nn.Linear(64, num_classes)

    def forward(self, data):
        x, edge_index, batch = data.x, data.edge_index, data.batch
        x = F.relu(self.conv1(x, edge_index))
        x, edge_index, _, batch, _, _ = self.pool(x, edge_index, None, batch)
        x = F.relu(self.conv2(x, edge_index))
        x = global_mean_pool(x, batch)
        return self.lin(x)

5. 異種グラフ

from torch_geometric.data import HeteroData
from torch_geometric.nn import HeteroConv, GCNConv, SAGEConv, to_hetero

# Create heterogeneous graph
data = HeteroData()
data['paper'].x = torch.randn(100, 128)
data['author'].x = torch.randn(200, 64)
data['author', 'writes', 'paper'].edge_index = torch.randint(0, 200, (2, 500))
data['paper', 'cites', 'paper'].edge_index = torch.randint(0, 100, (2, 300))
print(data)  # Shows all node/edge types

# Method 1: Auto-convert homogeneous model
model = GCN(...)
model = to_hetero(model, data.metadata(), aggr='sum')
out = model(data.x_dict, data.edge_index_dict)

# Method 2: Custom per-edge-type convolutions
class HeteroGNN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = HeteroConv({
            ('paper', 'cites', 'paper'): GCNConv(-1, 64),
            ('author', 'writes', 'paper'): SAGEConv((-1, -1), 64),
        }, aggr='sum')

    def forward(self, x_dict, edge_index_dict):
        x_dict = self.conv1(x_dict, edge_index_dict)
        return {k: F.relu(v) for k, v in x_dict.items()}

6. 変換と前処理

from torch_geometric.transforms import (
    NormalizeFeatures, AddSelfLoops, ToUndirected,
    RandomNodeSplit, RandomLinkSplit, Compose,
    KNNGraph, RadiusGraph, AddLaplacianEigenvectorPE
)

# Single transform
dataset = Planetoid(root='/tmp/Cora', name='Cora', transform=NormalizeFeatures())

# Compose multiple transforms
transform = Compose([
    ToUndirected(),
    AddSelfLoops(),
    NormalizeFeatures(),
])

# Data splitting
node_split = RandomNodeSplit(num_val=0.1, num_test=0.2)
link_split = RandomLinkSplit(num_val=0.1, num_test=0.2, is_undirected=True)

# Point cloud → graph
pc_transform = Compose([KNNGraph(k=6), NormalizeFeatures()])

# Positional encodings (for Graph Transformers)
pe_transform = AddLaplacianEigenvectorPE(k=10)

主要な概念

層の選択ガイド

タスク	層	主な特徴
ベースライン／汎用	GCNConv	スペクトル、キャッシュ対応、edge_weight
近傍ごとに重要度が異なる場合	GATConv / GATv2Conv	マルチヘッドアテンション
大規模な帰納的学習	SAGEConv	サンプリングに適している、mean/max/lstm aggr
グラフ分類	GINConv	WL テストで最大限の表現力
長距離依存関係	TransformerConv	グラフトランスフォーマー
スペクトルフィルタリング	ChebConv	チェビシェフ多項式、K ホップ
豊富なエッジ特徴量	NNConv	エッジ NN で edge_attr を処理
分子／3D 構造	SchNet、DimeNet	連続フィルター、角度
異種／複数リレーション	RGCNConv、HGTConv	複数のエッジタイプ
点群	EdgeConv、PointNetConv	動的グラフ、局所特徴量
深い GNN（過平滑化を回避）	APPNP + PairNorm	伝播を分離

データフローのアーキテクチャ

• edge_index：[2, num_edges] COO 形式。行 0 = ソース、行 1 = ターゲット
• ミニバッチ：ブロック対角隣接行列 + ノードからグラフへの対応を示す batch ベクトル。パディングなし
• 近傍サンプリング：NeighborLoader は各シードノードについて K ホップの部分グラフをサンプリングします。出力は有向で、ラベルが付け直されます
• 異種グラフ：x_dict（タイプごとの特徴量）、edge_index_dict（リレーションごとのエッジ）、スキーマ用の metadata()

集約オプション

集約	クラス	使用例
合計	SumAggregation	カウントが重要なタスク
平均	MeanAggregation	次数に依存しない場合
最大	MaxAggregation	顕著な特徴量の検出
ソフトマックス	SoftmaxAggregation(learn=True)	学習可能なアテンション
複合	MultiAggregation(['mean','max','std'])	シグナルの組み合わせ

一般的なワークフロー

1. ノード分類（グラフ全体）

import torch
import torch.nn.functional as F
from torch_geometric.datasets import Planetoid
from torch_geometric.nn import GCNConv

dataset = Planetoid(root='/tmp/Cora', name='Cora')
data = dataset[0]

class GCN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = GCNConv(dataset.num_features, 16)
        self.conv2 = GCNConv(16, dataset.num_classes)
    def forward(self, data):
        x = F.dropout(F.relu(self.conv1(data.x, data.edge_index)), p=0.5, training=self.training)
        return self.conv2(x, data.edge_index)

model = GCN()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

# Training
for epoch in range(200):
    model.train(); optimizer.zero_grad()
    out = model(data)
    F.cross_entropy(out[data.train_mask], data.y[data.train_mask]).backward()
    optimizer.step()

# Evaluation
model.eval()
pred = model(data).argmax(dim=1)
acc = (pred[data.test_mask] == data.y[data.test_mask]).float().mean()
print(f'Test Accuracy: {acc:.4f}')

2. グラフ分類（ミニバッチ）

from torch_geometric.datasets import TUDataset
from torch_geometric.loader import DataLoader
from torch_geometric.nn import GCNConv, global_mean_pool

dataset = TUDataset(root='/tmp/ENZYMES', name='ENZYMES')
train_dataset = dataset[:int(0.8 * len(dataset))]
test_dataset = dataset[int(0.8 * len(dataset)):]
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

class GraphNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = GCNConv(dataset.num_features, 64)
        self.conv2 = GCNConv(64, 64)
        self.lin = torch.nn.Linear(64, dataset.num_classes)
    def forward(self, data):
        x = F.relu(self.conv1(data.x, data.edge_index))
        x = F.relu(self.conv2(x, data.edge_index))
        x = global_mean_pool(x, data.batch)
        return self.lin(x)

model = GraphNet()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
for epoch in range(100):
    model.train()
    for batch in train_loader:
        optimizer.zero_grad()
        F.cross_entropy(model(batch), batch.y).backward()
        optimizer.step()

3. 近傍サンプリングによる大規模学習

from torch_geometric.loader import NeighborLoader

# Sample 25 1-hop and 10 2-hop neighbors per seed node
train_loader = NeighborLoader(
    data,
    num_neighbors=[25, 10],
    batch_size=128,
    input_nodes=data.train_mask,
)

model.train()
for batch in train_loader:
    optimizer.zero_grad()
    out = model(batch)
    # Only compute loss on seed nodes (first batch_size nodes)
    loss = F.cross_entropy(out[:batch.batch_size], batch.y[:batch.batch_size])
    loss.backward()
    optimizer.step()
# Note: output subgraphs are directed, indices relabeled 0..N-1

主なパラメーター

パラメーター	モジュール	デフォルト	範囲	効果
in_channels	すべての Conv レイヤー	—	int	入力特徴量の次元
out_channels	すべての Conv レイヤー	—	int	出力特徴量の次元
heads	GATConv	1	1-16	アテンションヘッドの数
dropout	GATConv	0.0	0-0.8	アテンション重みのドロップアウト
aggr	MessagePassing	'add'	add/mean/max	近傍ノードの集約方法
K	ChebConv	—	2-5	チェビシェフ多項式の次数
num_neighbors	NeighborLoader	—	list[int]	各ホップの近傍ノード数（例：[25,10]）
batch_size	DataLoader	—	16-512	バッチあたりのグラフ数またはシードノード数
ratio	TopKPooling	0.5	0.1-0.9	保持するノードの割合
lr	Adam	—	1e-4 to 0.01	学習率
weight_decay	Adam	0	0 to 5e-3	L2 正則化

ベストプラクティス

1. GCNConv から始める：複雑なアーキテクチャを試す前に、2 層の GCN をベースラインとして使用します
2. 遅延初期化を使用する：in_channels に -1 を渡すと、次元を自動的に推論できます：GCNConv(-1, 64)
3. 特徴量を正規化する：引用ネットワークやソーシャルネットワークには、NormalizeFeatures() 変換を適用します
4. アンチパターン — 層を増やしすぎる：GNN に必要な層は通常 2～3 層です。層を深くすると過平滑化が発生します。深くする必要がある場合は、JumpingKnowledge または PairNorm を使用します
5. GPU への転送：モデルとデータの両方を GPU に移動します：model.to(device)、data.to(device)
6. アンチパターン — バッチベクトルを無視する：グラフ分類では、常に global_mean_pool(x, batch) を使用します。batch を忘れると、すべてのグラフにわたってプーリングされます

よくあるレシピ

レシピ：モデルの説明可能性（GNNExplainer）

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)
print(f'Important edges: {explanation.edge_mask.topk(5).indices}')
print(f'Important features: {explanation.node_mask[10].topk(5).indices}')

レシピ：カスタム InMemoryDataset

from torch_geometric.data import InMemoryDataset, Data

class MyDataset(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super().__init__(root, transform, pre_transform)
        self.load(self.processed_paths[0])

    @property
    def raw_file_names(self):
        return ['data.csv']

    @property
    def processed_file_names(self):
        return ['data.pt']

    def process(self):
        data_list = []
        # Build Data objects from raw files
        edge_index = torch.tensor([[0, 1], [1, 0]], dtype=torch.long)
        x = torch.randn(2, 16)
        data_list.append(Data(x=x, edge_index=edge_index, y=torch.tensor([0])))

        if self.pre_filter is not None:
            data_list = [d for d in data_list if self.pre_filter(d)]
        if self.pre_transform is not None:
            data_list = [self.pre_transform(d) for d in data_list]
        self.save(data_list, self.processed_paths[0])

レシピ：JumpingKnowledge を使用したディープ GNN

from torch_geometric.nn import GCNConv, JumpingKnowledge, LayerNorm

class DeepGNN(torch.nn.Module):
    def __init__(self, in_ch, hidden, num_layers, out_ch):
        super().__init__()
        self.convs = torch.nn.ModuleList()
        self.norms = torch.nn.ModuleList()
        self.convs.append(GCNConv(in_ch, hidden))
        self.norms.append(LayerNorm(hidden))
        for _ in range(num_layers - 2):
            self.convs.append(GCNConv(hidden, hidden))
            self.norms.append(LayerNorm(hidden))
        self.convs.append(GCNConv(hidden, hidden))
        self.jk = JumpingKnowledge(mode='cat')
        self.lin = torch.nn.Linear(hidden * num_layers, out_ch)

    def forward(self, x, edge_index, batch):
        xs = []
        for conv, norm in zip(self.convs[:-1], self.norms):
            x = F.relu(norm(conv(x, edge_index)))
            xs.append(x)
        xs.append(self.convs[-1](x, edge_index))
        return self.lin(global_mean_pool(self.jk(xs), batch))

トラブルシューティング

問題	原因	解決策
edge_index の形状エラー	形式が正しくありません（[2, E] である必要があります）	COO 形式を使用します：torch.tensor([[src...],[dst...]], dtype=torch.long)
大規模グラフで OOM が発生する	グラフ全体を対象としたフォワードパス	ミニバッチ学習には NeighborLoader を使用します
精度が低い	過平滑化（層が多すぎる）	2～3 層に減らし、JumpingKnowledge または PairNorm を追加します
学習中に NaN が発生する	勾配の爆発	勾配クリッピングを追加し、学習率を下げ、特徴量のスケールを確認します
グラフレベルの出力が正しくない	プーリング時に batch がありません	batch テンソルを global_mean_pool(x, batch) に渡します
異種グラフの型エラー	ノード型またはエッジ型の不一致	data.metadata() がモデル定義と一致していることを確認します
DataLoader が遅い	グラフが大きく、サンプリングしていない	適切な num_neighbors（例：[25,10]）を指定して NeighborLoader を使用します
x の次元が一致しない	マルチヘッドアテンションの出力	GATConv では、concat=False でない限り、出力は heads*out_channels です
スパース演算のインポートエラー	オプションの依存関係が不足しています	PyG のホイールから torch_scatter、torch_sparse をインストールします
事前変換が適用されない	データセットはすでに処理済みです	processed/ ディレクトリを削除して再読み込みします

同梱リソース

• references/layers_transforms_reference.md — 40 種類以上の畳み込みレイヤー（GCN、GAT、SAGE、GIN、分子レイヤー、ハイパーグラフ）、集約演算子、プーリング（グローバルおよび階層型）、正規化レイヤー、事前構築済みモデル、オートエンコーダー、知識グラフの埋め込み、ユーティリティレイヤーの完全なカタログです。変換カタログ：構造、特徴、空間、拡張、メッシュ、特殊化。元の layers_reference.md（486 行）と transforms_reference.md（680 行）を統合しています。スクリプトの機能（benchmark_model.py、create_gnn_template.py、visualize_graph.py）は、Core API のコードブロックと一般的なレシピで扱っています
• references/datasets_catalog.md — ドメイン別に整理された包括的なデータセットカタログです。引用ネットワーク（Planetoid、Coauthor、Amazon）、グラフ分類（TUDataset の 120 種類以上のベンチマーク）、分子（QM9、ZINC、MoleculeNet）、ソーシャル（Reddit、Twitch）、知識グラフ（WordNet、FB15k）、異種グラフ（OGB_MAG、MovieLens、DBLP）、時系列（JODIE）、3D メッシュ（ShapeNet、ModelNet）、OGB 連携を収録しています。元の datasets_reference.md（575 行）を統合しています

関連スキル

• matplotlib-scientific-plotting — グラフ構造、学習曲線、アテンションの重みを可視化します

参考資料

• PyTorch Geometric ドキュメント：https://pytorch-geometric.readthedocs.io/
• PyG GitHub：https://github.com/pyg-team/pytorch_geometric
• Fey & Lenssen（2019）「Fast Graph Representation Learning with PyTorch Geometric」、ICLR Workshop