torch-geometric

PyTorch Geometric (PyG)

概要

PyTorch Geometric は PyTorch 上に構築されたライブラリで、グラフニューラルネットワーク（GNN）の開発と学習に使用します。ミニバッチ処理、マルチ GPU 学習、幾何学的深層学習の応用など、グラフや不規則な構造に対する深層学習にこのスキルを適用してください。

このスキルを使用する場面

このスキルは次のような作業で使用します:

• グラフベースの機械学習: ノード分類、グラフ分類、リンク予測
• 分子特性予測: 創薬、化学特性予測
• ソーシャルネットワーク分析: コミュニティ検出、影響力予測
• 引用ネットワーク: 論文分類、推薦システム
• 3D 幾何データ: 点群、メッシュ、分子構造
• 異種グラフ: 複数の種類のノードとエッジ（例: ナレッジグラフ）
• 大規模グラフ学習: 近傍サンプリング、分散学習

クイックスタート

インストール

uv pip install torch_geometric

追加の依存関係（スパース演算、クラスタリング）を使用する場合:

uv pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html

基本的なグラフの作成

import torch
from torch_geometric.data import Data

# Create a simple graph with 3 nodes
edge_index = torch.tensor([[0, 1, 1, 2],  # source nodes
                           [1, 0, 2, 1]], dtype=torch.long)  # target nodes
x = torch.tensor([[-1], [0], [1]], dtype=torch.float)  # node features

data = Data(x=x, edge_index=edge_index)
print(f"Nodes: {data.num_nodes}, Edges: {data.num_edges}")

ベンチマークデータセットの読み込み

from torch_geometric.datasets import Planetoid

# Load Cora citation network
dataset = Planetoid(root='/tmp/Cora', name='Cora')
data = dataset[0]  # Get the first (and only) graph

print(f"Dataset: {dataset}")
print(f"Nodes: {data.num_nodes}, Edges: {data.num_edges}")
print(f"Features: {data.num_node_features}, Classes: {dataset.num_classes}")

基本概念

データ構造

PyG は torch_geometric.data.Data クラスを使ってグラフを表現します。主な属性は次のとおりです:

• data.x: ノード特徴行列 [num_nodes, num_node_features]
• data.edge_index: COO 形式のグラフ接続情報 [2, num_edges]
• data.edge_attr: エッジ特徴行列 [num_edges, num_edge_features]（省略可能）
• data.y: ノードまたはグラフの目的ラベル
• data.pos: ノードの空間位置 [num_nodes, num_dimensions]（省略可能）
• カスタム属性: 任意の属性を追加できます（例: data.train_mask、data.batch）

重要: これらの属性は必須ではありません。必要に応じてカスタム属性で Data オブジェクトを拡張してください。

エッジインデックスの形式

エッジは [2, num_edges] のテンソルとして COO（座標）形式で格納されます:

• 1 行目: 始点ノードのインデックス
• 2 行目: 終点ノードのインデックス

# Edge list: (0→1), (1→0), (1→2), (2→1)
edge_index = torch.tensor([[0, 1, 1, 2],
                           [1, 0, 2, 1]], dtype=torch.long)

ミニバッチ処理

PyG はブロック対角の隣接行列を作成してバッチ処理を行い、複数のグラフを 1 つの大きな非連結グラフに連結します:

• 隣接行列は対角方向に積み上げられます
• ノード特徴はノード次元に沿って連結されます
• batch ベクトルが各ノードを元のグラフに対応付けます
• パディングは不要で、計算効率に優れています

from torch_geometric.loader import DataLoader

loader = DataLoader(dataset, batch_size=32, shuffle=True)
for batch in loader:
    print(f"Batch size: {batch.num_graphs}")
    print(f"Total nodes: {batch.num_nodes}")
    # batch.batch maps nodes to graphs

グラフニューラルネットワークの構築

メッセージパッシングのパラダイム

PyG の GNN は近傍集約の方式に従います:

1. ノード特徴を変換する
2. エッジに沿ってメッセージを伝播させる
3. 近傍からのメッセージを集約する
4. ノード表現を更新する

組み込みレイヤーの利用

PyG は 40 種類以上の畳み込みレイヤーを提供しています。代表的なものは次のとおりです:

GCNConv（Graph Convolutional Network）:

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

class GCN(torch.nn.Module):
    def __init__(self, num_features, num_classes):
        super().__init__()
        self.conv1 = GCNConv(num_features, 16)
        self.conv2 = GCNConv(16, num_classes)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index)
        return F.log_softmax(x, dim=1)

GATConv（Graph Attention Network）:

from torch_geometric.nn import GATConv

class GAT(torch.nn.Module):
    def __init__(self, num_features, num_classes):
        super().__init__()
        self.conv1 = GATConv(num_features, 8, heads=8, dropout=0.6)
        self.conv2 = GATConv(8 * 8, num_classes, heads=1, concat=False, dropout=0.6)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        x = F.dropout(x, p=0.6, training=self.training)
        x = F.elu(self.conv1(x, edge_index))
        x = F.dropout(x, p=0.6, training=self.training)
        x = self.conv2(x, edge_index)
        return F.log_softmax(x, dim=1)

GraphSAGE:

from torch_geometric.nn import SAGEConv

class GraphSAGE(torch.nn.Module):
    def __init__(self, num_features, num_classes):
        super().__init__()
        self.conv1 = SAGEConv(num_features, 64)
        self.conv2 = SAGEConv(64, num_classes)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index)
        return F.log_softmax(x, dim=1)

カスタムメッセージパッシングレイヤー

独自のレイヤーを作成する場合は MessagePassing を継承します:

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", or "max"
        self.lin = torch.nn.Linear(in_channels, out_channels)

    def forward(self, x, edge_index):
        # Add self-loops to adjacency matrix
        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))

        # Transform node features
        x = self.lin(x)

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

        # Propagate messages
        return self.propagate(edge_index, x=x, norm=norm)

    def message(self, x_j, norm):
        # x_j: features of source nodes
        return norm.view(-1, 1) * x_j

主なメソッド:

• forward(): メインのエントリポイント
• message(): 送信元から送信先ノードへのメッセージを構築する
• aggregate(): メッセージを集約する（通常はオーバーライドせず、aggr パラメータを設定します）
• update(): 集約後にノードの埋め込みを更新する

変数命名規則: テンソル名に _i または _j を付けると、それぞれ送信先ノードまたは送信元ノードに自動的に対応付けられます。

データセットの扱い方

組み込みデータセットの読み込み

PyG は豊富なベンチマークデータセットを提供しています:

# Citation networks (node classification)
from torch_geometric.datasets import Planetoid
dataset = Planetoid(root='/tmp/Cora', name='Cora')  # or 'CiteSeer', 'PubMed'

# Graph classification
from torch_geometric.datasets import TUDataset
dataset = TUDataset(root='/tmp/ENZYMES', name='ENZYMES')

# Molecular datasets
from torch_geometric.datasets import QM9
dataset = QM9(root='/tmp/QM9')

# Large-scale datasets
from torch_geometric.datasets import Reddit
dataset = Reddit(root='/tmp/Reddit')

網羅的な一覧は references/datasets_reference.md を参照してください。

カスタムデータセットの作成

メモリに収まるデータセットの場合は InMemoryDataset を継承します:

from torch_geometric.data import InMemoryDataset, Data
import torch

class MyOwnDataset(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 ['my_data.csv']  # Files needed in raw_dir

    @property
    def processed_file_names(self):
        return ['data.pt']  # Files in processed_dir

    def download(self):
        # Download raw data to self.raw_dir
        pass

    def process(self):
        # Read data, create Data objects
        data_list = []

        # Example: Create a simple graph
        edge_index = torch.tensor([[0, 1], [1, 0]], dtype=torch.long)
        x = torch.randn(2, 16)
        y = torch.tensor([0], dtype=torch.long)

        data = Data(x=x, edge_index=edge_index, y=y)
        data_list.append(data)

        # Apply pre_filter and pre_transform
        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]

        # Save processed data
        self.save(data_list, self.processed_paths[0])

メモリに収まらない大規模データセットの場合は、Dataset を継承して len() と get(idx) を実装してください。

CSV からのグラフの読み込み

import pandas as pd
import torch
from torch_geometric.data import HeteroData

# Load nodes
nodes_df = pd.read_csv('nodes.csv')
x = torch.tensor(nodes_df[['feat1', 'feat2']].values, dtype=torch.float)

# Load edges
edges_df = pd.read_csv('edges.csv')
edge_index = torch.tensor([edges_df['source'].values,
                           edges_df['target'].values], dtype=torch.long)

data = Data(x=x, edge_index=edge_index)

学習のワークフロー

ノード分類（単一グラフ）

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

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

# Create model
model = GCN(dataset.num_features, dataset.num_classes)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

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

    if epoch % 10 == 0:
        print(f'Epoch {epoch}, Loss: {loss.item():.4f}')

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

グラフ分類（複数グラフ）

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

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.lin = torch.nn.Linear(64, num_classes)

    def forward(self, data):
        x, edge_index, batch = data.x, data.edge_index, data.batch

        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = self.conv2(x, edge_index)
        x = F.relu(x)

        # Global pooling (aggregate node features to graph-level)
        x = global_mean_pool(x, batch)

        x = self.lin(x)
        return F.log_softmax(x, dim=1)

# Load dataset
dataset = TUDataset(root='/tmp/ENZYMES', name='ENZYMES')
loader = DataLoader(dataset, batch_size=32, shuffle=True)

model = GraphClassifier(dataset.num_features, dataset.num_classes)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# Training
model.train()
for epoch in range(100):
    total_loss = 0
    for batch in loader:
        optimizer.zero_grad()
        out = model(batch)
        loss = F.nll_loss(out, batch.y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()

    if epoch % 10 == 0:
        print(f'Epoch {epoch}, Loss: {total_loss / len(loader):.4f}')

ニューベイトサンプリングによる大規模グラフ

大規模グラフの場合は、NeighborLoader を使用してサブグラフをサンプリングします:

from torch_geometric.loader import NeighborLoader

# Create a neighbor sampler
train_loader = NeighborLoader(
    data,
    num_neighbors=[25, 10],  # Sample 25 neighbors for 1st hop, 10 for 2nd hop
    batch_size=128,
    input_nodes=data.train_mask,
)

# Training
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.nll_loss(out[:batch.batch_size], batch.y[:batch.batch_size])
    loss.backward()
    optimizer.step()

重要:

• 出力されるサブグラフは有向です
• ノードのインデックスは再ラベル付けされます (0 から batch.num_nodes - 1)
• 損失計算には、シードノードの予測のみを使用してください
• 2〜3 ホップを超えるサンプリングは一般に現実的ではありません

高度な機能

ヘテロジニアスグラフ

複数のノードタイプとエッジタイプを持つグラフの場合は、HeteroData を使用します:

from torch_geometric.data import HeteroData

data = HeteroData()

# Add node features for different types
data['paper'].x = torch.randn(100, 128)  # 100 papers with 128 features
data['author'].x = torch.randn(200, 64)  # 200 authors with 64 features

# Add edges for different types (source_type, edge_type, target_type)
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)

ホモジニアスモデルをヘテロジニアスモデルに変換します:

from torch_geometric.nn import to_hetero

# Define homogeneous model
model = GNN(...)

# Convert to heterogeneous
model = to_hetero(model, data.metadata(), aggr='sum')

# Use as normal
out = model(data.x_dict, data.edge_index_dict)

または、エッジタイプごとのカスタム処理には HeteroConv を使用します:

from torch_geometric.nn import HeteroConv, GCNConv, SAGEConv

class HeteroGNN(torch.nn.Module):
    def __init__(self, metadata):
        super().__init__()
        self.conv1 = HeteroConv({
            ('paper', 'cites', 'paper'): GCNConv(-1, 64),
            ('author', 'writes', 'paper'): SAGEConv((-1, -1), 64),
        }, aggr='sum')

        self.conv2 = HeteroConv({
            ('paper', 'cites', 'paper'): GCNConv(64, 32),
            ('author', 'writes', 'paper'): SAGEConv((64, 64), 32),
        }, aggr='sum')

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

変換 (Transforms)

変換を適用して、グラフの構造や特徴量を変更します:

from torch_geometric.transforms import NormalizeFeatures, AddSelfLoops, Compose

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

# Compose multiple transforms
transform = Compose([
    AddSelfLoops(),
    NormalizeFeatures(),
])
dataset = Planetoid(root='/tmp/Cora', name='Cora', transform=transform)

よく使われる変換:

• 構造: ToUndirected, AddSelfLoops, RemoveSelfLoops, KNNGraph, RadiusGraph
• 特徴量: NormalizeFeatures, NormalizeScale, Center
• サンプリング: RandomNodeSplit, RandomLinkSplit
• 位置エンコーディング: AddLaplacianEigenvectorPE, AddRandomWalkPE

全一覧は references/transforms_reference.md を参照してください。

モデルの説明可能性

PyG には、モデルの予測を理解するための説明可能性ツールが用意されています:

from torch_geometric.explain import Explainer, GNNExplainer

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

# Generate explanation for a specific node
node_idx = 10
explanation = explainer(data.x, data.edge_index, index=node_idx)

# Visualize
print(f'Node {node_idx} explanation:')
print(f'Important edges: {explanation.edge_mask.topk(5).indices}')
print(f'Important features: {explanation.node_mask[node_idx].topk(5).indices}')

プーリング演算

階層的なグラフ表現のために使用します:

from torch_geometric.nn import TopKPooling, global_mean_pool

class HierarchicalGNN(torch.nn.Module):
    def __init__(self, num_features, num_classes):
        super().__init__()
        self.conv1 = GCNConv(num_features, 64)
        self.pool1 = TopKPooling(64, ratio=0.8)
        self.conv2 = GCNConv(64, 64)
        self.pool2 = 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.pool1(x, edge_index, None, batch)

        x = F.relu(self.conv2(x, edge_index))
        x, edge_index, _, batch, _, _ = self.pool2(x, edge_index, None, batch)

        x = global_mean_pool(x, batch)
        x = self.lin(x)
        return F.log_softmax(x, dim=1)

一般的なパターンとベストプラクティス

グラフの性質の確認

# Undirected check
from torch_geometric.utils import is_undirected
print(f"Is undirected: {is_undirected(data.edge_index)}")

# Connected components
from torch_geometric.utils import connected_components
print(f"Connected components: {connected_components(data.edge_index)}")

# Contains self-loops
from torch_geometric.utils import contains_self_loops
print(f"Has self-loops: {contains_self_loops(data.edge_index)}")

GPU トレーニング

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
data = data.to(device)

# For DataLoader
for batch in loader:
    batch = batch.to(device)
    # Train...

モデルの保存と読み込み

# Save
torch.save(model.state_dict(), 'model.pth')

# Load
model = GCN(num_features, num_classes)
model.load_state_dict(torch.load('model.pth'))
model.eval()

レイヤーの機能

レイヤーを選択する際は、次の機能を考慮してください:

• SparseTensor: 効率的な疎行列演算をサポートします
• edge_weight: 1 次元のエッジ重みを扱います
• edge_attr: 多次元のエッジ特徴量を処理します
• Bipartite: 二部グラフ (ソースとターゲットで次元が異なる) に対応します
• Lazy: 入力次元を指定せずに初期化できます

GNN チートシートは references/layer_capabilities.md を参照してください。

リソース

同梱のリファレンス

このスキルには詳細なリファレンスドキュメントが含まれています:

• references/layers_reference.md: 40 種類以上の GNN レイヤーの完全な一覧。説明と機能を掲載
• references/datasets_reference.md: カテゴリ別に整理された包括的なデータセットカタログ
• references/transforms_reference.md: 利用可能なすべての変換と、その用途
• references/api_patterns.md: 一般的な API パターンとコーディング例

スクリプト

ユーティリティスクリプトは scripts/ に用意されています:

• scripts/visualize_graph.py: networkx と matplotlib を使ってグラフ構造を可視化します
• scripts/create_gnn_template.py: 一般的な GNN アーキテクチャのボイラープレートコードを生成します
• scripts/benchmark_model.py: 標準的なデータセットでモデルの性能をベンチマークします

スクリプトは直接実行するか、実装パターンを確認するために読んでください。

公式リソース

• ドキュメント: https://pytorch-geometric.readthedocs.io/
• GitHub: https://github.com/pyg-team/pytorch_geometric
• チュートリアル: https://pytorch-geometric.readthedocs.io/en/latest/get_started/introduction.html
• サンプル: https://github.com/pyg-team/pytorch_geometric/tree/master/examples