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

Node2Vecとは

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

先日論文を紹介した記事を執筆したので、詳しくはこちらの記事をご覧ください。

目標

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

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

実装の設定は、基本的に元論文(node2vec: Scalable Feature Learning for Networks)とPyTorch GeometricのNode2Vecクラス(torch_geometric.nn.models.Node2Vec)に準拠しています。

データセットの準備

今回は、グラフの扱いに特化したライブラリPyTorch Geometric(PyG)を主に使用し、PyGに用意された簡単なグラフデータのKarateClubを使用します。

データセットの内容把握

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つのコミュニティからなる小さいデータセットです。

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	エッジに付与された重みなど(ここでは設定されていない)

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形式に変換し、表示。

edge_attrの設定

data._edge_attr = [1 for _ in range(data.edge_index.shape[1])]

edge_attrが設定されていないと扱えないので、全てのエッジに重み1を設定します。

学習

モデルの設定

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クラスを返し、バッチ処置や、イテレータによるサンプルの提供を行います。

マスクの生成

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個生成します。
正直、これではサンプルの数と多様性が不十分すぎるので、性能は出ません。

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

学習とバリデーション

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によって簡略化された方法で埋め込みの更新を行います。

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

性能のテスト

# 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つずつしかないから当然)。

別のデータセットでの実験

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