【PyTorch】Node2Vecで簡単な多クラス分類を実装してみる。

## Node2Vecとは
簡単に言うと、グラフ上のランダムウォークによって特徴を読み取り、ノードやエッジの埋め込み(特徴ベクトル)を学習する方法のことです。

先日論文を紹介した記事を執筆したので、詳しくは[こちらの記事](https://qiita.com/pale_tomato/items/8a036dbd232de0743237)をご覧ください。

## 目標
ノードをネットワーク上の個人として、その人の属するコミュニティをラベルとします。Node2Vecによって適切な分散表現を学習し、属するコミュニティが未知のノードに対する予測を行うことを目標とします。

<font color="LightGray">データセットが小さすぎて精度が出ていないので、コードの書き方だけ参考にしてください。</font>

実装の設定は、基本的に元論文([node2vec: Scalable Feature Learning for Networks](https://arxiv.org/abs/1607.00653))とPyTorch GeometricのNode2Vecクラス([torch_geometric.nn.models.Node2Vec](https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.models.Node2Vec.html))に準拠しています。

## データセットの準備
今回は、グラフの扱いに特化したライブラリ[PyTorch Geometric(PyG)](https://pytorch-geometric.readthedocs.io/en/latest/)を主に使用し、PyGに用意された簡単なグラフデータのKarateClubを使用します。

### データセットの内容把握
```python
import torch_geometric
from torch_geometric.datasets import KarateClub
import matplotlib.pyplot as plt
import networkx as nx


dataset = KarateClub()
data = dataset[0]
```
PyGのデータセットは基本的に`InMemoryDataset`クラスを継承しており、インスタンス変数として`Data`クラスのオブジェクトを持っています。
`KarateClub`は34個のノードと4つのコミュニティからなる小さいデータセットです。

```python
print(f"shape of x: {data.x.shape}")
print(f"x: {data.x}")

# shape of x: torch.Size([34, 34])
# x: tensor([[1., 0., 0.,  ..., 0., 0., 0.],
#         [0., 1., 0.,  ..., 0., 0., 0.],
#         [0., 0., 1.,  ..., 0., 0., 0.],
#         ...,
#         [0., 0., 0.,  ..., 1., 0., 0.],
#         [0., 0., 0.,  ..., 0., 1., 0.],
#         [0., 0., 0.,  ..., 0., 0., 1.]])

print(f"shape of y: {data.y.shape}")
print(f"y: {data.y}")

# shape of y: torch.Size([34])
# y: tensor([1, 1, 1, 1, 3, 3, 3, 1, 0, 1, 3, 1, 1, 1, 0, 0, 3, 1, 0, 1, 0, # 1, 0, 0, 2, 2, 0, 0, 2, 0, 0, 2, 0, 0])

print(f"shape of edge index: {data.edge_index.shape}")
print(f"edge index: {data.edge_index}")

# shape of edge index: torch.Size([2, 156])
# edge index: tensor([[ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  1,  1,
#           1,  1,  1,  1,  1,  1,  1,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  3,
#           3,  3,  3,  3,  3,  4,  4,  4,  5,  5,  5,  5,  6,  6,  6,  6,  7,  7,
#           7,  7,  8,  8,  8,  8,  8,  9,  9, 10, 10, 10, 11, 12, 12, 13, 13, 13,
#          13, 13, 14, 14, 15, 15, 16, 16, 17, 17, 18, 18, 19, 19, 19, 20, 20, 21,
#          21, 22, 22, 23, 23, 23, 23, 23, 24, 24, 24, 25, 25, 25, 26, 26, 27, 27,
#          27, 27, 28, 28, 28, 29, 29, 29, 29, 30, 30, 30, 30, 31, 31, 31, 31, 31,
#          31, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 33, 33, 33, 33, 33,
#          33, 33, 33, 33, 33, 33, 33, 33, 33, 33, 33, 33],
#         [ 1,  2,  3,  4,  5,  6,  7,  8, 10, 11, 12, 13, 17, 19, 21, 31,  0,  2,
#           3,  7, 13, 17, 19, 21, 30,  0,  1,  3,  7,  8,  9, 13, 27, 28, 32,  0,
#           1,  2,  7, 12, 13,  0,  6, 10,  0,  6, 10, 16,  0,  4,  5, 16,  0,  1,
#           2,  3,  0,  2, 30, 32, 33,  2, 33,  0,  4,  5,  0,  0,  3,  0,  1,  2,
#           3, 33, 32, 33, 32, 33,  5,  6,  0,  1, 32, 33,  0,  1, 33, 32, 33,  0,
#           1, 32, 33, 25, 27, 29, 32, 33, 25, 27, 31, 23, 24, 31, 29, 33,  2, 23,
#          24, 33,  2, 31, 33, 23, 26, 32, 33,  1,  8, 32, 33,  0, 24, 25, 28, 32,
#          33,  2,  8, 14, 15, 18, 20, 22, 23, 29, 30, 31, 33,  8,  9, 13, 14, 15,
#          18, 19, 20, 22, 23, 26, 27, 28, 29, 30, 31, 32]])

print(f"edge attr: {data.edge_attr}")

# edge attr: None
```

`Data`クラスの主なアトリビュートは以下の通りです。

| アトリビュート名 | 役割 |
| ---- | ---- |
| `x` | ノードごとの特徴量(ここではone-hotベクトル) |
| `y` | ノードごとのラベル(ここでは属するコミュニティ) |
| `edge_index` | エッジの始点の集合と終点の集合, インデックスで1:1に対応する |
| `edge_attr` | エッジに付与された重みなど(ここでは設定されていない) |


```python
G = torch_geometric.utils.convert.to_networkx(data)

plt.figure(figsize=(8, 8))
nx.draw_networkx(G, with_labels=False, alpha=0.5, node_color=data.y)
plt.show()
```

NetworkX形式に変換し、表示。

![KarateClub.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3588797/9661ee7d-f46e-58f4-095c-12c704f373af.png)

### edge_attrの設定
```python
data._edge_attr = [1 for _ in range(data.edge_index.shape[1])]
```
`edge_attr`が設定されていないと扱えないので、全てのエッジに重み1を設定します。

## 学習

### モデルの設定
```python
model = torch_geometric.nn.Node2Vec(
    data.edge_index,
    embedding_dim=16,
    walk_length=10,
    context_size=5,
    walks_per_node=10,
    num_negative_samples=1,
    p=0.25,
    q=4.0,
    sparse=True,
).to(device)

loader = model.loader(batch_size=16, shuffle=True, num_workers=8)
optimizer = torch.optim.SGD(list(model.parameters()), lr=0.1)
```
| パラメータ | 意味 |
| ---- | ---- |
| `embedding_dim` | 特徴ベクトルの次元 |
| `walk_length` | ランダムウォークの長さ |
| `context_size` | 文脈のサイズ(ウォークの何ノードずつまとめて扱うか) |
| `walks_per_node` | ノードごとのウォーク回数 |
| `num_negative_samples` | 正例1つに対する負例の個数　|
| `p` | Return Parameter |
| `q` | In-out Parameter |

パラメータの詳細については前回の記事で詳しく説明しています。
`loader`は`DataLoader`クラスを返し、バッチ処置や、イテレータによるサンプルの提供を行います。

### マスクの生成

```python
def generate_masks(data, k):
    """
    各クラスのノードを1つずつtestに割り当て、残りをtrainに割り当てる。
    train, valはk-fold交差検証用に分割する。
    書き方が汚い気がする。
    """
    y = data.y
    test_mask = torch.zeros(y.size(0), dtype=torch.bool)

    for i in range(int(y.max()) + 1):
        test_mask[(y == i).nonzero(as_tuple=False)[0]] = True

    num_val_nodes = (y.size(0) - int(sum(test_mask))) // k

    train_masks = []
    val_masks = []

    tmp_val_index = 0
    for _ in range(k):
        train_mask = torch.zeros(y.size(0), dtype=torch.bool)
        val_mask = torch.zeros(y.size(0), dtype=torch.bool)
        cnt = 0
        for i in range(y.size(0)):
            if test_mask[i] and i >= tmp_val_index and cnt < num_val_nodes:
                tmp_val_index += 1
                pass
            elif i == tmp_val_index and cnt < num_val_nodes:
                val_mask[i] = True
                cnt += 1
                tmp_val_index += 1
            else:
                train_mask[i] = True

        train_masks.append(train_mask)
        val_masks.append(val_mask)

    return train_masks, val_masks, test_mask


train_masks, val_masks, test_mask = generate_masks(data, 10)
data.test_mask = test_mask
```
ノード数が非常に少ないので、ラベルごとにテスト用ノードを1つずつ切り出し、残りのノードの1/10を検証用に当てたセットを10個生成します。
正直、これではサンプルの数と多様性が不十分すぎるので、性能は出ません。

**このマスクは、分散表現の学習自体とは関係がないので注意してください。**

### 学習とバリデーション
```python
z = model()
loss_hist = []

for epoch in range(1, NUM_EPOCHS + 1):
    # training
    model.train()
    total_loss = 0

    # loaderをイテレートするとsampleメソッドが呼ばれ、正例と負例が得られる。
    for pos_rw, neg_rw in loader:
        optimizer.zero_grad()
        loss = model.loss(pos_rw.to(device), neg_rw.to(device))
        loss.backward()
        optimizer.step()
        total_loss += loss.item()

    total_loss = total_loss / len(loader)
    loss_hist.append(loss)

    # set masks
    data.train_mask = train_masks[epoch - 1]
    data.val_mask = val_masks[epoch - 1]

    # validation
    with torch.no_grad():
        model.eval()
        acc = model.test(
            z[data.train_mask],
            data.y[data.train_mask],
            z[data.val_mask],
            data.y[data.val_mask],
            max_iter=150,
        )

    acc_hist.append(acc)
    print(f"Epoch {epoch}, Loss: {total_loss:.4f}, Acc: {acc:.4f}")
```

trainingセクションでランダムウォークを生成し、Negative Samplingによって簡略化された方法で埋め込みの更新を行います。

学習した分散表現を用いて、ロジスティック回帰によりクラス分類の性能を評価します。

### 性能のテスト
```python
# test
model.eval()
rest_mask = [not x for x in data.test_mask]
test_acc = model.test(
    z[rest_mask],
    data.y[rest_mask],
    z[data.test_mask],
    data.y[data.test_mask],
    max_iter=150,
)

fig, ax = plt.subplots(2, 1)
ax[0].plot(loss_hist)
ax[0].set_xlabel("Epoch")
ax[0].set_ylabel("Loss")
ax[1].plot(acc_hist)
ax[1].set_xlabel("Epoch")
ax[1].set_ylabel("Accuracy")
plt.suptitle(f"test acc: {test_acc}")
plt.subplots_adjust(hspace=0.3)
plt.show()
```
最初にテスト用に除外しておいたノードを用いて、性能の検証を行います。
ノード数が少ないので正答率は低く安定していないです(検証データが3つずつしかないから当然)。

![Figure_1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3588797/5ce3c8bf-a6bd-ff73-2ffe-e77f2eec5a27.png)

## 別のデータセットでの実験
![magpen_d_64_r_10_l_30_k_10_neg_1_p_4_q_0_25_batch_32_workers_8_lr_1_iter_10_plateau.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3588797/dd53f15d-82ff-d1f7-6dff-3ed953898290.png)

もう少し大きいデータセットで同じように実験したところ、もう少しそれらしいグラフが出ました。しかし正答率は低いです。