PyTorch Geometricを利用してノードラベリングを解く

# 概要

[PyTorch Geometric](https://github.com/pyg-team/pytorch_geometric)を利用してGCNを実装し、ノードラベリングを解くまでの基本的な流れを理解する。
基本的には[公式のチュートリアル](https://pytorch-geometric.readthedocs.io/en/latest/notes/introduction.html)のものを少し変更して利用する。

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

# コード

## データの準備

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

PyTorch Geometricではデータは[`torch_geometric.data.Data`](https://pytorch-geometric.readthedocs.io/en/latest/modules/data.html#torch_geometric.data.Data)のデータインスタンスとして扱う。
DataクラスはPyTorch Geometric でデータを扱う際のクラスであり、以下のパラメータを持つ。

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

```python
import torch
from torch_geometric.data import Data
```


```python
# ノードを連番の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]])




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




    torch.Size([2708, 1433])




```python
# ラベルをエンコード
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])




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




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



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


```python
# 学習とテストの分割
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
```

## モデルの定義

[チュートリアル](https://pytorch-geometric.readthedocs.io/en/latest/notes/introduction.html)のものを少し変更して使用する。
基本的にはPyTorchのモデルの定義と同様で、pytorch geometoric からGCNConvレイヤーをインポートして使用する。


```python
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などは使用していない。


```python
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などに変更してみる。


```python
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)
```


```python
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
    - [GitHub](https://github.com/pyg-team/pytorch_geometric)
    - [ドキュメント](https://pytorch-geometric.readthedocs.io/en/latest/index.html) 
    - [チュートリアル](https://pytorch-geometric.readthedocs.io/en/latest/notes/introduction.html)

- Cora
    - [The Cora dataset - Graph Data Science Consulting](https://graphsandnetworks.com/the-cora-dataset/)