第3章：高度なGNNアーキテクチャ

1. Graph Attention Networks（GAT）

Graph Attention Network（GAT、グラフ注意機構ネットワーク）は、GCNのように全ての隣接ノードを同じ重みで平均化するのではなく、注意機構（Attention Mechanism）を使ってノードごとに重要度の異なる重みを学習するGNNアーキテクチャです。自然言語処理のTransformerで使われるSelf-Attentionと同じ発想を、グラフ構造に応用したものと考えると理解しやすいでしょう。

GATの基本原理

GATでは、ノード \(i\) とその隣接ノード \(j\) の間の注意係数（Attention Coefficient） \(\alpha_{ij}\) を、両ノードの特徴量から学習します。

まず、注意スコア \(e_{ij}\) を次のように計算します：

$$e_{ij} = \text{LeakyReLU}\left(\mathbf{a}^T [\mathbf{W}\mathbf{h}_i \, \| \, \mathbf{W}\mathbf{h}_j]\right)$$

ここで \(\mathbf{W}\) は学習可能な重み行列、\(\mathbf{a}\) は学習可能な注意ベクトル、\(\|\) はベクトルの連結（concatenation）を表します。次に、隣接ノード全体で正規化してソフトマックス関数を適用します：

$$\alpha_{ij} = \frac{\exp(e_{ij})}{\sum_{k \in \mathcal{N}(i)} \exp(e_{ik})}$$

最後に、注意係数で重み付けした隣接ノード特徴量の加重和を計算し、ノード \(i\) の新しい特徴量とします：

$$\mathbf{h}_i' = \sigma\left(\sum_{j \in \mathcal{N}(i)} \alpha_{ij} \mathbf{W} \mathbf{h}_j\right)$$

マルチヘッド注意機構（Multi-Head Attention）

GATは、Transformerと同様にマルチヘッド注意機構を使い、複数の独立した注意機構を並列に学習して結果を連結（または平均化）することで、学習を安定させ表現力を高めます。torch_geometric.nn.GATConvではheads引数でヘッド数を指定します。

GATConvの実装

Coraデータセット（引用ネットワーク）を使って、GATConvによるノード分類を実装してみましょう。

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

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

class GAT(torch.nn.Module):
    def __init__(self, num_features, num_classes, hidden_channels=8, heads=8):
        super().__init__()
        # 第1層: 8ヘッドの注意機構を使い、出力を連結（concat=True）
        self.conv1 = GATConv(num_features, hidden_channels, heads=heads, dropout=0.6)
        # 第2層: 1ヘッドで出力次元をクラス数に合わせる
        self.conv2 = GATConv(hidden_channels * heads, num_classes, heads=1,
                              concat=False, dropout=0.6)

    def forward(self, x, 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)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = GAT(dataset.num_features, dataset.num_classes).to(device)
data = data.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.005, weight_decay=5e-4)

model.train()
for epoch in range(100):
    optimizer.zero_grad()
    out = model(data.x, data.edge_index)
    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()
    if (epoch + 1) % 20 == 0:
        print(f'Epoch {epoch+1:03d}, Loss: {loss:.4f}')

model.eval()
with torch.no_grad():
    pred = model(data.x, data.edge_index).argmax(dim=1)
    test_acc = (pred[data.test_mask] == data.y[data.test_mask]).float().mean().item()
print(f'Test Accuracy: {test_acc:.4f}')

出力例：

Epoch 020, Loss: 1.2345
Epoch 040, Loss: 0.7823
Epoch 060, Loss: 0.5412
Epoch 080, Loss: 0.4501
Epoch 100, Loss: 0.3987
Test Accuracy: 0.8280

2. GraphSAGE（サンプリング型GNN）

GraphSAGE（Graph SAmple and aggreGatE）は、数百万〜数十億ノード規模の大規模グラフでも学習できるように設計されたGNNです。GCNやGATが全ての隣接ノードを使って一度にグラフ全体を処理するフルバッチ学習（Full-batch Learning）を前提とするのに対し、GraphSAGEは各ノードの隣接ノードを固定数だけサンプリング（Sampling）し、ミニバッチ単位で学習することでメモリ使用量と計算量を抑えます。

GraphSAGEのアルゴリズム

GraphSAGEの各層では、次の2ステップを行います。

1. サンプリング: 各ノードについて、隣接ノードから固定数（例: 10個）をランダムに抽出する
2. 集約（Aggregation）: サンプリングした隣接ノードの特徴量を集約し、自ノードの特徴量と結合して更新する

$$\mathbf{h}_i^{(k)} = \sigma\left(\mathbf{W}^{(k)} \cdot \text{CONCAT}\left(\mathbf{h}_i^{(k-1)}, \; \text{AGG}_k\left(\{\mathbf{h}_j^{(k-1)} : j \in \mathcal{N}_s(i)\}\right)\right)\right)$$

ここで \(\mathcal{N}_s(i)\) はノード \(i\) のサンプリングされた隣接ノード集合、\(\text{AGG}_k\) は集約関数です。

SAGEConvモデルの実装

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

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

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

    def forward(self, x, edge_index):
        x = F.relu(self.conv1(x, edge_index))
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.conv2(x, edge_index)
        return x

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = GraphSAGE(dataset.num_features, dataset.num_classes).to(device)
print(model)

NeighborLoaderによるミニバッチ学習

大規模グラフを想定し、torch_geometric.loader.NeighborLoaderを使って隣接ノードのサンプリングとミニバッチ生成を行います。num_neighbors=[10, 5]は「1層目で最大10個、2層目で最大5個の隣接ノードをサンプリングする」という意味です。

from torch_geometric.loader import NeighborLoader

# 訓練ノードを起点に、2ホップ分の隣接ノードをサンプリングしながらミニバッチを生成
train_loader = NeighborLoader(
    data,
    num_neighbors=[10, 5],
    batch_size=128,
    input_nodes=data.train_mask,
    shuffle=True,
)

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

model.train()
for epoch in range(20):
    total_loss = 0
    for batch in train_loader:
        batch = batch.to(device)
        optimizer.zero_grad()
        out = model(batch.x, batch.edge_index)
        # バッチ内の先頭 batch_size 件が「起点ノード（seed nodes）」であり、
        # それ以外はサンプリングで追加された隣接ノード（損失計算には使わない）
        loss = F.cross_entropy(out[:batch.batch_size], batch.y[:batch.batch_size])
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    if (epoch + 1) % 5 == 0:
        print(f'Epoch {epoch+1:03d}, Loss: {total_loss / len(train_loader):.4f}')

評価

model.eval()
data_device = data.to(device)
with torch.no_grad():
    out = model(data_device.x, data_device.edge_index)
    pred = out.argmax(dim=1)
    test_acc = (pred[data_device.test_mask] == data_device.y[data_device.test_mask]).float().mean().item()
print(f'Test Accuracy: {test_acc:.4f}')

出力例：

Epoch 005, Loss: 1.5432
Epoch 010, Loss: 0.9821
Epoch 015, Loss: 0.6543
Epoch 020, Loss: 0.4987
Test Accuracy: 0.7920

3. グラフプーリング手法

ここまで見てきたGCN・GAT・GraphSAGEは、いずれも各ノードの特徴量を更新するノードレベル（Node-level）の処理でした。しかし、分子全体の毒性予測やタンパク質の機能分類のように「グラフ全体」に対して1つの予測を行いたい場合は、ノードごとの特徴量をグラフ全体の1つのベクトルに集約する必要があります。この集約処理をグラフプーリング（Graph Pooling）または読み出し（Readout）と呼びます。

グローバルプーリング

最も基本的なプーリングは、全ノードの特徴量を単純な統計量で集約するグローバルプーリング（Global Pooling）です。

import torch
from torch_geometric.nn import global_mean_pool, global_max_pool, global_add_pool
from torch_geometric.datasets import TUDataset
from torch_geometric.loader import DataLoader

dataset = TUDataset(root='/tmp/MUTAG', name='MUTAG')
loader = DataLoader(dataset, batch_size=4, shuffle=True)

batch = next(iter(loader))
print(f'Batch: {batch.num_graphs} graphs, {batch.num_nodes} nodes total')

# ノード特徴量をそのままグラフ表現とみなし、各プーリング手法の出力形状を比較
mean_pooled = global_mean_pool(batch.x, batch.batch)
max_pooled = global_max_pool(batch.x, batch.batch)
sum_pooled = global_add_pool(batch.x, batch.batch)

print(f'Mean pooling shape: {mean_pooled.shape}')
print(f'Max pooling shape: {max_pooled.shape}')
print(f'Sum pooling shape: {sum_pooled.shape}')

出力例：

Batch: 4 graphs, 71 nodes total
Mean pooling shape: torch.Size([4, 7])
Max pooling shape: torch.Size([4, 7])
Sum pooling shape: torch.Size([4, 7])

出力形状が [グラフ数, 特徴次元] になっている点に注目してください。batch属性（各ノードがどのグラフに属するかを表すインデックス）を使って、ノード数の異なる複数のグラフを1つのバッチにまとめたまま、グラフごとに正しく集約できています。

階層的プーリング（Hierarchical Pooling）

グローバルプーリングは全ノードを一度に集約しますが、階層的プーリングではCNNのプーリング層のように、GNN層の間でノード数を段階的に削減していきます。TopKPoolingは、学習可能なスコア関数でノードの重要度を評価し、上位 \(k\) 個（比率ratioで指定）のノードだけを残す代表的な手法です。

import torch
import torch.nn.functional as F
from torch_geometric.nn import GraphConv, TopKPooling, global_mean_pool

class HierarchicalPoolNet(torch.nn.Module):
    def __init__(self, num_features, num_classes, hidden_channels=64):
        super().__init__()
        self.conv1 = GraphConv(num_features, hidden_channels)
        # 重要度スコア上位80%のノードのみを残す
        self.pool1 = TopKPooling(hidden_channels, ratio=0.8)
        self.conv2 = GraphConv(hidden_channels, hidden_channels)
        self.lin = torch.nn.Linear(hidden_channels, num_classes)

    def forward(self, x, edge_index, batch):
        x = F.relu(self.conv1(x, edge_index))
        # プーリングでノード数を削減しつつ、edge_indexとbatchも整合させて更新
        x, edge_index, _, batch, _, _ = self.pool1(x, edge_index, None, batch)
        x = F.relu(self.conv2(x, edge_index))
        x = global_mean_pool(x, batch)
        return self.lin(x)

model = HierarchicalPoolNet(num_features=7, num_classes=2)
print(model)

4. グラフ分類タスク

グラフ分類（Graph Classification）は、グラフ全体に1つのラベルを予測するタスクです。分子の毒性・活性予測、タンパク質の機能分類などが代表例です。ここではMUTATデータセットならぬMUTAGデータセット（変異原性を持つ化合物の分子グラフ、TUDatasetの1つ）を使って、GATConvとSAGEConvを組み合わせたグラフ分類モデルを最初から最後まで実装します。

データセットの準備

import torch
import torch.nn.functional as F
from torch_geometric.datasets import TUDataset
from torch_geometric.loader import DataLoader
from torch_geometric.nn import GATConv, SAGEConv, global_mean_pool

torch.manual_seed(42)

# MUTAG: 188個の分子グラフ、2クラス分類（変異原性の有無）
dataset = TUDataset(root='/tmp/MUTAG', name='MUTAG').shuffle()

train_dataset = dataset[:150]
test_dataset = dataset[150:]

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32)

print(f'Number of training graphs: {len(train_dataset)}')
print(f'Number of test graphs: {len(test_dataset)}')
print(f'Number of node features: {dataset.num_node_features}')
print(f'Number of classes: {dataset.num_classes}')

出力例：

Number of training graphs: 150
Number of test graphs: 38
Number of node features: 7
Number of classes: 2

モデルの定義

GATConvでノード間の重要度を学習した後、SAGEConvで近傍情報を集約し、最後にglobal_mean_poolでグラフ全体のベクトル表現を得る、というアーキテクチャです。

class GraphClassifier(torch.nn.Module):
    def __init__(self, num_features, num_classes, hidden_channels=64):
        super().__init__()
        # GAT層: マルチヘッド注意機構でノード間の重要度を考慮した特徴抽出
        self.gat = GATConv(num_features, hidden_channels, heads=4, concat=True, dropout=0.2)
        # GraphSAGE層: 近傍情報を集約して特徴量を洗練
        self.sage = SAGEConv(hidden_channels * 4, hidden_channels)
        self.lin = torch.nn.Linear(hidden_channels, num_classes)

    def forward(self, x, edge_index, batch):
        x = F.elu(self.gat(x, edge_index))
        x = F.dropout(x, p=0.2, training=self.training)
        x = F.relu(self.sage(x, edge_index))
        # グラフ全体を1つのベクトルに集約（readout）
        x = global_mean_pool(x, batch)
        x = F.dropout(x, p=0.5, training=self.training)
        return self.lin(x)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = GraphClassifier(dataset.num_node_features, dataset.num_classes).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
print(model)

学習と評価

def train():
    model.train()
    total_loss = 0
    for data in train_loader:
        data = data.to(device)
        optimizer.zero_grad()
        out = model(data.x, data.edge_index, data.batch)
        loss = F.cross_entropy(out, data.y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * data.num_graphs
    return total_loss / len(train_loader.dataset)

@torch.no_grad()
def test(loader):
    model.eval()
    correct = 0
    for data in loader:
        data = data.to(device)
        out = model(data.x, data.edge_index, data.batch)
        pred = out.argmax(dim=1)
        correct += (pred == data.y).sum().item()
    return correct / len(loader.dataset)

for epoch in range(1, 101):
    loss = train()
    if epoch % 20 == 0:
        train_acc = test(train_loader)
        test_acc = test(test_loader)
        print(f'Epoch {epoch:03d}, Loss: {loss:.4f}, '
              f'Train Acc: {train_acc:.4f}, Test Acc: {test_acc:.4f}')

出力例：

Epoch 020, Loss: 0.5893, Train Acc: 0.7133, Test Acc: 0.6842
Epoch 040, Loss: 0.4721, Train Acc: 0.7867, Test Acc: 0.7368
Epoch 060, Loss: 0.4102, Train Acc: 0.8200, Test Acc: 0.7632
Epoch 080, Loss: 0.3654, Train Acc: 0.8467, Test Acc: 0.7895
Epoch 100, Loss: 0.3389, Train Acc: 0.8600, Test Acc: 0.8158

5. リンク予測タスク

リンク予測（Link Prediction）は、グラフ中の2つのノード間にエッジ（リンク）が存在するかどうかを予測するタスクです。ソーシャルネットワークの友人推薦、レコメンデーションシステムの商品推薦、知識グラフの欠損関係補完など、幅広い応用があります。

タスクの設計

リンク予測では、既存のエッジを正例（Positive Sample）、存在しないノードペアを負例（Negative Sample）としてサンプリングし、二値分類問題として学習します。エッジを訓練・検証・テスト用に分割する際は、情報の漏洩（Leakage）を防ぐため、グラフ構造の分割とラベルの分割を同時に行う必要があります。PyTorch GeometricのRandomLinkSplitがこれを自動化してくれます。

import torch
import torch.nn.functional as F
from torch_geometric.datasets import Planetoid
from torch_geometric.transforms import RandomLinkSplit, NormalizeFeatures
from torch_geometric.nn import GCNConv
from torch_geometric.utils import negative_sampling

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

# エッジを訓練/検証/テスト用に分割（負例エッジも自動生成）
transform = RandomLinkSplit(
    num_val=0.05,
    num_test=0.1,
    is_undirected=True,
    add_negative_train_samples=False,
)
train_data, val_data, test_data = transform(data)

print(f'Train message-passing edges: {train_data.edge_index.shape}')
print(f'Train supervision edges: {train_data.edge_label_index.shape}')
print(f'Val supervision edges: {val_data.edge_label_index.shape}')
print(f'Test supervision edges: {test_data.edge_label_index.shape}')

エンコーダ・デコーダモデルの実装

GCNConvでノード埋め込みを計算するエンコーダ（Encoder）と、2つのノード埋め込みの内積でエッジの存在確率を算出するデコーダ（Decoder）から成るモデルを構築します。

class LinkPredGCN(torch.nn.Module):
    def __init__(self, num_features, hidden_channels=128, out_channels=64):
        super().__init__()
        self.conv1 = GCNConv(num_features, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, out_channels)

    def encode(self, x, edge_index):
        x = F.relu(self.conv1(x, edge_index))
        return self.conv2(x, edge_index)

    def decode(self, z, edge_label_index):
        # 2つのノード埋め込みの内積を、エッジが存在するスコアとする
        src, dst = edge_label_index
        return (z[src] * z[dst]).sum(dim=-1)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = LinkPredGCN(dataset.num_features).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

train_data = train_data.to(device)
val_data = val_data.to(device)
test_data = test_data.to(device)

学習と評価

訓練時は毎エポック新しい負例エッジをサンプリングすることで、モデルが特定の負例に過学習することを防ぎます。評価指標には二値分類でよく使われるROC-AUC（Receiver Operating Characteristic - Area Under the Curve）を用います。

from sklearn.metrics import roc_auc_score

def train():
    model.train()
    optimizer.zero_grad()
    z = model.encode(train_data.x, train_data.edge_index)

    # 各エポックで負例エッジ（存在しないエッジ）をランダムに再サンプリング
    neg_edge_index = negative_sampling(
        edge_index=train_data.edge_index,
        num_nodes=train_data.num_nodes,
        num_neg_samples=train_data.edge_label_index.size(1),
        method='sparse',
    )

    edge_label_index = torch.cat([train_data.edge_label_index, neg_edge_index], dim=1)
    edge_label = torch.cat([
        train_data.edge_label,
        train_data.edge_label.new_zeros(neg_edge_index.size(1)),
    ], dim=0)

    out = model.decode(z, edge_label_index).view(-1)
    loss = F.binary_cross_entropy_with_logits(out, edge_label)
    loss.backward()
    optimizer.step()
    return loss.item()

@torch.no_grad()
def test(data):
    model.eval()
    z = model.encode(data.x, data.edge_index)
    out = model.decode(z, data.edge_label_index).view(-1).sigmoid()
    return roc_auc_score(data.edge_label.cpu().numpy(), out.cpu().numpy())

for epoch in range(1, 101):
    loss = train()
    if epoch % 20 == 0:
        val_auc = test(val_data)
        test_auc = test(test_data)
        print(f'Epoch {epoch:03d}, Loss: {loss:.4f}, '
              f'Val AUC: {val_auc:.4f}, Test AUC: {test_auc:.4f}')

出力例：

Epoch 020, Loss: 0.6123, Val AUC: 0.7845, Test AUC: 0.7791
Epoch 040, Loss: 0.5234, Val AUC: 0.8412, Test AUC: 0.8356
Epoch 060, Loss: 0.4756, Val AUC: 0.8723, Test AUC: 0.8681
Epoch 080, Loss: 0.4489, Val AUC: 0.8891, Test AUC: 0.8834
Epoch 100, Loss: 0.4321, Val AUC: 0.8956, Test AUC: 0.8902

6. 異種グラフの扱い

ここまで扱ってきたグラフは、全て同じ種類のノード（例: 論文ノードのみ）から成る同種グラフ（Homogeneous Graph）でした。しかし実世界の多くのグラフは、ユーザーと商品、著者と論文、原子と結合種別のように、複数の種類のノードやエッジを含む異種グラフ（Heterogeneous Graph）です。

HeteroDataオブジェクト

PyTorch Geometricでは、異種グラフをHeteroDataオブジェクトで表現します。ノード特徴量はノード種類ごとに、エッジは(始点ノード種類, 関係名, 終点ノード種類)という3つ組（トリプル）で管理されます。ユーザーが映画を評価するレコメンデーションシステムを例に、二部グラフを構築してみましょう。

import torch
from torch_geometric.data import HeteroData

data = HeteroData()

# ノード特徴量（ノード種類ごとに次元が異なってもよい）
data['user'].x = torch.randn(5, 16)    # 5人のユーザー、16次元特徴
data['movie'].x = torch.randn(8, 32)   # 8本の映画、32次元特徴

# エッジ: (始点ノード種類, 関係名, 終点ノード種類) の3つ組で定義
data['user', 'rates', 'movie'].edge_index = torch.tensor([
    [0, 0, 1, 2, 3, 4],
    [0, 1, 1, 2, 5, 7],
], dtype=torch.long)

# 逆方向のエッジも明示的に追加する（メッセージパッシングは有向的に扱われるため）
data['movie', 'rated_by', 'user'].edge_index = \
    data['user', 'rates', 'movie'].edge_index.flip(0)

print(data)
print(f'Node types: {data.node_types}')
print(f'Edge types: {data.edge_types}')

出力例：

HeteroData(
  user={ x=[5, 16] },
  movie={ x=[8, 32] },
  (user, rates, movie)={ edge_index=[2, 6] },
  (movie, rated_by, user)={ edge_index=[2, 6] }
)
Node types: ['user', 'movie']
Edge types: [('user', 'rates', 'movie'), ('movie', 'rated_by', 'user')]

HeteroConvによる異種グラフGNN

torch_geometric.nn.HeteroConvを使うと、エッジ種類ごとに個別のGNN層（ここではSAGEConv）を適用し、同じノード種類に集まった結果を集約するモデルを簡潔に構築できます。ノード種類ごとに入力次元が異なる場合は、SAGEConv((-1, -1), hidden_channels)のように-1を指定すると、初回の順伝播時に入力次元が自動推論されます。

from torch_geometric.nn import HeteroConv, SAGEConv, Linear

class HeteroGNN(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        # エッジ種類ごとに個別のSAGEConvを適用し、ノード種類ごとに結果を集約
        self.conv1 = HeteroConv({
            ('user', 'rates', 'movie'): SAGEConv((-1, -1), hidden_channels),
            ('movie', 'rated_by', 'user'): SAGEConv((-1, -1), hidden_channels),
        }, aggr='sum')
        self.lin_user = Linear(hidden_channels, out_channels)
        self.lin_movie = Linear(hidden_channels, out_channels)

    def forward(self, x_dict, edge_index_dict):
        x_dict = self.conv1(x_dict, edge_index_dict)
        x_dict = {key: x.relu() for key, x in x_dict.items()}
        return {
            'user': self.lin_user(x_dict['user']),
            'movie': self.lin_movie(x_dict['movie']),
        }

model = HeteroGNN(hidden_channels=16, out_channels=4)
out_dict = model(data.x_dict, data.edge_index_dict)

print(f"User embeddings shape: {out_dict['user'].shape}")
print(f"Movie embeddings shape: {out_dict['movie'].shape}")

出力例：

User embeddings shape: torch.Size([5, 4])
Movie embeddings shape: torch.Size([8, 4])

演習問題

for heads in [2, 4, 8]:
    model = GAT(dataset.num_features, dataset.num_classes, heads=heads).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.005, weight_decay=5e-4)

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

    model.eval()
    with torch.no_grad():
        pred = model(data.x, data.edge_index).argmax(dim=1)
        acc = (pred[data.test_mask] == data.y[data.test_mask]).float().mean().item()
    print(f'heads={heads}: Test Accuracy = {acc:.4f}')

from torch_geometric.nn import global_mean_pool, global_max_pool

class GraphClassifierV2(torch.nn.Module):
    def __init__(self, num_features, num_classes, hidden_channels=64):
        super().__init__()
        self.gat = GATConv(num_features, hidden_channels, heads=4, concat=True, dropout=0.2)
        self.sage = SAGEConv(hidden_channels * 4, hidden_channels)
        # mean と max の連結により2倍の次元になる
        self.lin = torch.nn.Linear(hidden_channels * 2, num_classes)

    def forward(self, x, edge_index, batch):
        x = F.elu(self.gat(x, edge_index))
        x = F.dropout(x, p=0.2, training=self.training)
        x = F.relu(self.sage(x, edge_index))

        x_mean = global_mean_pool(x, batch)
        x_max = global_max_pool(x, batch)
        x = torch.cat([x_mean, x_max], dim=1)

        x = F.dropout(x, p=0.5, training=self.training)
        return self.lin(x)

import torch
from torch_geometric.data import HeteroData
from torch_geometric.nn import HeteroConv, SAGEConv, Linear

data = HeteroData()
data['user'].x = torch.randn(5, 16)
data['movie'].x = torch.randn(8, 32)
data['genre'].x = torch.randn(3, 8)   # 新しいノード種類: ジャンル

data['user', 'rates', 'movie'].edge_index = torch.tensor([
    [0, 0, 1, 2, 3, 4],
    [0, 1, 1, 2, 5, 7],
], dtype=torch.long)
data['movie', 'rated_by', 'user'].edge_index = \
    data['user', 'rates', 'movie'].edge_index.flip(0)

# 各映画（0-7）を対応するジャンル（0-2）に接続
data['movie', 'has_genre', 'genre'].edge_index = torch.tensor([
    [0, 1, 2, 3, 4, 5, 6, 7],
    [0, 0, 1, 1, 2, 2, 0, 1],
], dtype=torch.long)

class HeteroGNNv2(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = HeteroConv({
            ('user', 'rates', 'movie'): SAGEConv((-1, -1), hidden_channels),
            ('movie', 'rated_by', 'user'): SAGEConv((-1, -1), hidden_channels),
            ('movie', 'has_genre', 'genre'): SAGEConv((-1, -1), hidden_channels),
        }, aggr='sum')
        self.lin_user = Linear(hidden_channels, out_channels)
        self.lin_movie = Linear(hidden_channels, out_channels)

    def forward(self, x_dict, edge_index_dict):
        x_dict = self.conv1(x_dict, edge_index_dict)
        x_dict = {key: x.relu() for key, x in x_dict.items()}
        return {
            'user': self.lin_user(x_dict['user']),
            'movie': self.lin_movie(x_dict['movie']),
        }

model = HeteroGNNv2(hidden_channels=16, out_channels=4)
out_dict = model(data.x_dict, data.edge_index_dict)
print(f"User embeddings shape: {out_dict['user'].shape}")
print(f"Movie embeddings shape: {out_dict['movie'].shape}")

まとめ

この章では、より高度なGNNアーキテクチャとグラフレベルのタスクを学びました。冒頭の学習目標を振り返ってみましょう。

• ✅ GAT: 注意機構によりノードごとに異なる重要度を学習し、GCNより柔軟な特徴集約を実現する仕組みを理解しました
• ✅ GraphSAGE: 隣接ノードのサンプリングとミニバッチ学習により、大規模グラフでも学習可能にするサンプリング戦略を学びました
• ✅ グラフプーリング: global_mean_pool等のグローバルプーリングと、TopKPoolingによる階層的プーリングでグラフ全体の表現を得る方法を実装しました
• ✅ グラフ分類: MUTAGデータセットに対し、GATConvとSAGEConvを組み合わせたモデルで分子の変異原性を予測するパイプラインを構築しました
• ✅ リンク予測: エンコーダ・デコーダ構造と負例サンプリングにより、エッジの存在確率を予測する基本を理解しました
• ✅ 異種グラフ: HeteroDataとHeteroConvを使い、複数種類のノード・エッジを持つグラフを扱う方法を学びました

参考リソース

• PyTorch Geometric公式ドキュメント
• PyTorch Geometric GitHub
• GAT論文: Graph Attention Networks
• GraphSAGE論文: Inductive Representation Learning on Large Graphs