概要

PyTorch Geometricを利用してGCNを実装し、ノードラベリングを解くまでの基本的な流れを理解する。
基本的には公式のチュートリアルのものを少し変更して利用する。

~~「モデルの変更」の部分についてバグがあったことに投稿後に気がついたため、修正中です。~~ 修正済みです。

コード

データの準備

今回はThe Cora dataset - Graph Data Science Consultingからダウンロードしたデータを用いており、データの読み込み部分についてもリンク先のコードを少し変更して利用した。
Coraデータセットをはじめとして、論文などでベースラインとして使用されることの多いデータはPyTorch Geometricから利用できるようになっているものもあるが、勉強目的のため自分で実装した。

PyTorch Geometricではデータはtorch_geometric.data.Dataのデータインスタンスとして扱う。
DataクラスはPyTorch Geometric でデータを扱う際のクラスであり、以下のパラメータを持つ。

data.x: 各ノードの特徴量 [ノード数, ノードの特徴量の次元] (有向なエッジであることに注意が必要)
data.edge_index: エッジの接続 [2, エッジ数], 型は torch.long
data.y: 学習時に使用するラベル, e.g., ノードラベリングの場合は [ノード数, *]

import torch
from torch_geometric.data import Data

# ノードを連番のintに変換
node_list = df_node_data.index.to_list()
node_index = {node: index for index, node in enumerate(node_list)}

# エッジのリストをノードの連番で置き換える
edge_list = df_edge_list.replace(node_index).values.T
print(f"shape: {edge_list.shape}")
print(edge_list[:, 0:4])

edge_list = torch.tensor(edge_list, dtype=torch.long)
edge_list

shape: (2, 5429)
[[ 163  163  163  163]
 [ 402  659 1696 2295]]





tensor([[ 163,  163,  163,  ..., 1887, 1902,  837],
        [ 402,  659, 1696,  ..., 2258, 1887, 1686]])

# 各ノードの特徴量を作成
x = df_node_data.drop("subject", axis=1).values
x = torch.tensor(x, dtype=torch.float)
x.shape

torch.Size([2708, 1433])

# ラベルをエンコード
subject_list = df_node_data["subject"].unique()
subject_index = {subject: index for index, subject in enumerate(subject_list)}
y = df_node_data["subject"].replace(subject_index).values
y = torch.tensor(y, dtype=torch.long)
y

tensor([0, 1, 2,  ..., 5, 6, 0])

data = Data(x=x, edge_list=edge_list, y=y)
data

Data(x=[2708, 1433], y=[2708], edge_list=[2, 5429])

ノードラベリングのタスクでは、Data に学習とテスト用のマスクを付与することで、どのデータを利用するかを判断する。

# 学習とテストの分割
import numpy as np
from sklearn.model_selection import train_test_split

train_node_list, test_node_list = train_test_split(
    node_list, test_size=0.9, shuffle=True, random_state=42
)
train_mask = np.array([1 if node in train_node_list else 0 for node in node_list])
train_mask = torch.tensor(train_mask, dtype=torch.bool)
test_mask = np.array([1 if node in test_node_list else 0 for node in node_list])
test_mask = torch.tensor(test_mask, dtype=torch.bool)

Data.train_mask = train_mask
Data.test_mask = test_mask

モデルの定義

チュートリアルのものを少し変更して使用する。
基本的にはPyTorchのモデルの定義と同様で、pytorch geometoric からGCNConvレイヤーをインポートして使用する。

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


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

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

        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)

学習と評価

学習と評価については一般的なPyTorchと同様の利用法となる。
今回はデータが固定されている (Batchへの分割などを行わない) ため、DataLoaderなどは使用していない。

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = GCN(1433, 7, 16).to(device)
data = data.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

for epoch in range(30):
    # 学習
    model.train()
    optimizer.zero_grad()
    out = model(data)
    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)
        valid_loss = F.nll_loss(pred[data.test_mask], data.y[data.test_mask])
        pred = pred.argmax(dim=1)
        correct = (pred[data.test_mask] == data.y[data.test_mask]).sum()

    acc = int(correct) / int(data.test_mask.sum())

    print(
        "step {}, train_loss {:.3f}, valid_loss {:.3f}, acc {:.3f}".format(
            epoch + 1, loss, valid_loss, acc
        )
    )

step 1, train_loss 1.970, valid_loss 1.856, acc 0.367
step 2, train_loss 1.816, valid_loss 1.762, acc 0.475
step 3, train_loss 1.693, valid_loss 1.663, acc 0.521
step 4, train_loss 1.550, valid_loss 1.562, acc 0.566
step 5, train_loss 1.433, valid_loss 1.461, acc 0.612
step 6, train_loss 1.271, valid_loss 1.362, acc 0.661
step 7, train_loss 1.165, valid_loss 1.267, acc 0.695
step 8, train_loss 1.042, valid_loss 1.179, acc 0.710
step 9, train_loss 0.948, valid_loss 1.103, acc 0.720
step 10, train_loss 0.889, valid_loss 1.039, acc 0.731
step 11, train_loss 0.795, valid_loss 0.985, acc 0.735
step 12, train_loss 0.727, valid_loss 0.940, acc 0.737
step 13, train_loss 0.659, valid_loss 0.902, acc 0.738
step 14, train_loss 0.599, valid_loss 0.871, acc 0.743
step 15, train_loss 0.558, valid_loss 0.846, acc 0.750
step 16, train_loss 0.526, valid_loss 0.825, acc 0.760
step 17, train_loss 0.467, valid_loss 0.809, acc 0.764
step 18, train_loss 0.501, valid_loss 0.794, acc 0.768
step 19, train_loss 0.404, valid_loss 0.781, acc 0.773
step 20, train_loss 0.372, valid_loss 0.770, acc 0.777
step 21, train_loss 0.374, valid_loss 0.762, acc 0.780
step 22, train_loss 0.319, valid_loss 0.758, acc 0.779
step 23, train_loss 0.298, valid_loss 0.755, acc 0.779
step 24, train_loss 0.309, valid_loss 0.755, acc 0.780
step 25, train_loss 0.280, valid_loss 0.756, acc 0.779
step 26, train_loss 0.245, valid_loss 0.757, acc 0.777
step 27, train_loss 0.283, valid_loss 0.761, acc 0.775
step 28, train_loss 0.228, valid_loss 0.764, acc 0.774
step 29, train_loss 0.254, valid_loss 0.767, acc 0.774
step 30, train_loss 0.255, valid_loss 0.773, acc 0.774

モデルの変更

モデルの変更

ここまでの例ではGCNを利用していたが、GATなどに変更してみる。

from torch_geometric.nn import GATConv


class GAT(torch.nn.Module):
    def __init__(self, num_node_features, num_classes, hidden_size):
        super().__init__()
        # レイヤーをGATに変更する
        self.conv1 = GATConv(num_node_features, hidden_size)
        self.conv2 = GATConv(hidden_size, num_classes)

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

        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)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = GAT(1433, 7, 16).to(device)
data = data.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

for epoch in range(30):
    # 学習
    model.train()
    optimizer.zero_grad()
    out = model(data)
    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)
        valid_loss = F.nll_loss(pred[data.test_mask], data.y[data.test_mask])
        pred = pred.argmax(dim=1)
        correct = (pred[data.test_mask] == data.y[data.test_mask]).sum()

    acc = int(correct) / int(data.test_mask.sum())

    print(
        "step {}, train_loss {:.3f}, valid_loss {:.3f}, acc {:.3f}".format(
            epoch + 1, loss, valid_loss, acc
        )
    )

step 1, train_loss 1.970, valid_loss 1.880, acc 0.468
step 2, train_loss 1.852, valid_loss 1.806, acc 0.529
step 3, train_loss 1.753, valid_loss 1.715, acc 0.589
step 4, train_loss 1.622, valid_loss 1.613, acc 0.637
step 5, train_loss 1.488, valid_loss 1.506, acc 0.684
step 6, train_loss 1.363, valid_loss 1.403, acc 0.710
step 7, train_loss 1.194, valid_loss 1.306, acc 0.730
step 8, train_loss 1.114, valid_loss 1.216, acc 0.740
step 9, train_loss 1.040, valid_loss 1.136, acc 0.745
step 10, train_loss 0.944, valid_loss 1.066, acc 0.746
step 11, train_loss 0.781, valid_loss 1.008, acc 0.749
step 12, train_loss 0.687, valid_loss 0.961, acc 0.748
step 13, train_loss 0.700, valid_loss 0.925, acc 0.749
step 14, train_loss 0.651, valid_loss 0.899, acc 0.750
step 15, train_loss 0.589, valid_loss 0.882, acc 0.750
step 16, train_loss 0.539, valid_loss 0.873, acc 0.749
step 17, train_loss 0.542, valid_loss 0.868, acc 0.748
step 18, train_loss 0.471, valid_loss 0.867, acc 0.747
step 19, train_loss 0.431, valid_loss 0.868, acc 0.746
step 20, train_loss 0.411, valid_loss 0.871, acc 0.748
step 21, train_loss 0.372, valid_loss 0.875, acc 0.750
step 22, train_loss 0.376, valid_loss 0.877, acc 0.752
step 23, train_loss 0.354, valid_loss 0.877, acc 0.756
step 24, train_loss 0.285, valid_loss 0.882, acc 0.759
step 25, train_loss 0.293, valid_loss 0.888, acc 0.760
step 26, train_loss 0.254, valid_loss 0.894, acc 0.762
step 27, train_loss 0.270, valid_loss 0.903, acc 0.763
step 28, train_loss 0.285, valid_loss 0.914, acc 0.759
step 29, train_loss 0.243, valid_loss 0.922, acc 0.758
step 30, train_loss 0.240, valid_loss 0.928, acc 0.758

まとめ

公式のチュートリアルを参考に、PyTorch Geometricを用いてGCNを実装しノードラベリングのタスクを解くまでの流れをまとめた。
モデルの変更なども容易に実装できるためPyTorchやTensorflowをベタ書きするよりも短時間で実装できる。
今回は試していないが、Network Embeddingの手法なども実装されており、グラフデータに対してベースラインとして試してみたい手法はおおよそ実装されているため、使いこなせると非常に便利だと思われる。

参考文献

PyTorch Geometric
Cora
- The Cora dataset - Graph Data Science Consulting