[pyg Node2Vec]グラフ情報を埋め込んで、機械学習の分類モデルを試してみる(グラフ分類)

1. Node2Vecを試してみる

node2vecを試しに使ってみたので、よければ参考にしてください。

今回はCoraデータセットを使用しました。

node2vec.py

import sys
import torch
import matplotlib.pyplot as plt

from torch_geometric.datasets import Planetoid
from torch_geometric.nn import Node2Vec
from sklearn.decomposition import PCA

# 埋め込みデータ保存先
MODEL_PATH = f'./embedded_data/node2vec_Cora.pickle'

# パラメータ設定
torch.manual_seed(1)
batch_size = 32
num_workers = 4 if sys.platform == 'linux' else 0

# データセットの取得
data = Planetoid('./data/Planetoid', name='Cora')[0]

# モデルの定義
model = Node2Vec(
    data.edge_index,
    embedding_dim=batch_size,
    walk_length=30,
    walks_per_node=20,
    context_size=10,
)

# データを分ける
loader = model.loader(batch_size=batch_size, shuffle=True, num_workers=num_workers)

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# 学習
for epoch in range(1, 31):
    model.train()
    for pos_rw, neg_rw in loader:
        optimizer.zero_grad()
        loss = model.loss(pos_rw, neg_rw)
        loss.backward()
        optimizer.step()
    print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}')

# 評価
model.eval()
z = model()
acc = model.test(
    train_z=z[data.train_mask],
    train_y=data.y[data.train_mask],
    test_z=z[data.test_mask],
    test_y=data.y[data.test_mask],
    max_iter=150,
)

print(f'Loss: {loss:.4f}, Acc: {acc:.4f}')

# 埋め込みデータの作成(32次元)
torch.save(z, MODEL_PATH)

emb32d = z.detach().numpy()

# 主成分分析
pca = PCA(n_components=2)
emb2d = pca.fit_transform(emb32d)

# 埋め込みデータ(2次元)をグラフで表示
plt.title("node embedding in 2D")
plt.scatter(emb2d[:,0],emb2d[:,1])
plt.show()

1-1. 実行してみる

コンソール結果

$ python node2vec.py

Epoch: 001, Loss: 1.8629
Epoch: 002, Loss: 1.1869
Epoch: 003, Loss: 1.0128
Epoch: 004, Loss: 0.9484
Epoch: 005, Loss: 0.9224
Epoch: 006, Loss: 0.8813
Epoch: 007, Loss: 0.8847
Epoch: 008, Loss: 0.8594
Epoch: 009, Loss: 0.8528
Epoch: 010, Loss: 0.8606
Epoch: 011, Loss: 0.8592
Epoch: 012, Loss: 0.8469
Epoch: 013, Loss: 0.8447
Epoch: 014, Loss: 0.8431
Epoch: 015, Loss: 0.8450
Epoch: 016, Loss: 0.8569
Epoch: 017, Loss: 0.8463
Epoch: 018, Loss: 0.8389
Epoch: 019, Loss: 0.8370
Epoch: 020, Loss: 0.8452
Epoch: 021, Loss: 0.8443
Epoch: 022, Loss: 0.8412
Epoch: 023, Loss: 0.8443
Epoch: 024, Loss: 0.8408
Epoch: 025, Loss: 0.8375
Epoch: 026, Loss: 0.8458
Epoch: 027, Loss: 0.8502
Epoch: 028, Loss: 0.8502
Epoch: 029, Loss: 0.8488
Epoch: 030, Loss: 0.8542

Loss: 0.8542, Acc: 0.7070

1-2. 埋め込みグラフの表示

2. SVMでグラフ分類

svm.py

import torch
import numpy as np

from sklearn.svm import SVC
from torch_geometric.datasets import Planetoid
from sklearn.model_selection import cross_val_score

MODEL_PATH = f'./embedded_data/node2vec_Cora.pickle'

# 保存した埋め込みデータを読み込む
x = torch.load(MODEL_PATH).detach().numpy()

# Coraデータセットのノードラベルを取得
y = Planetoid('./data/Planetoid', name='Cora')[0].y.detach().numpy()

# SVMモデル呼び出し
model = SVC(kernel='linear')

# 交差検証
scores = cross_val_score(model, x, y)
print("交差検証結果：", round(np.mean(scores), 5))

2-1. 実行してみる

コンソール結果

$ python svm.py

交差検証結果： 0.79985