神奇的PyG：快速实现Node2vector

公众号：异度侵入
“我最早接触数据挖掘，还没学会调参，只会用树模型时，比赛时便在群里听别人讨论图神经网络，谈起Deepwalk，Node2Vector，听起来唬人的算法，觉得这样的算法离我很远，好想使用啊，然后现在我遇到了PyG，一切都不一样了”

01 PyG

什么是PyG?PyG，全称Pytorch Geometric，听见这个名字就知道它是pytorch的生态圈的，囊括了许多图算法。比如Node2vector，GCN，GraphSage。https://pytorch-geometric.readthedocs.io/en/latest/。那先来安装一下吧。

pip install torch-scatter -f https://pytorch-geometric.com/whl/torch-1.8.0+cu111.html
pip install torch-sparse -f https://pytorch-geometric.com/whl/torch-1.8.0+cu111.html
pip install torch-cluster -f https://pytorch-geometric.com/whl/torch-1.8.0+cu111.html
pip install torch-spline-conv -f https://pytorch-geometric.com/whl/torch-1.8.0+cu111.html
pip install torch-geometric

支持的pytorch版本1.4——1.8。(https://pytorch-geometric.com/whl/）

02 Node2vector

说起Node2Vector，可能要说起randomwalk,然后再说起deepwalk，然后word2vector。

熟悉nlp的同学，都知道word2vector的思想很简单，通过一句话里几个单词预测一个单词，然后不断去训练一个模型让这个模型预测的越来越准，然后取其中隐藏层做embedding层。发一个word2vector的链接，可以详细了解一下(https://zhuanlan.zhihu.com/p/27234078)。

图是什么呢？图是一种不同于图片语音等的非欧式结构数据，比如社交网络。如下图所示。

那么现在对于图结构，我们是不是也可以利用word2vector的思想得到图节点embedding，那么首先要解决的问题就是，图结构中怎么找到nlp中“句子”这个对应的单位。好吧，为了解决这问题，我们就用到了randomwalk，随机游走，顾名思义，就是随便走，用在图结构上就是从一个节点出发随机选择下一个点，然后几步之后停止，这样就得到序列了，比如上图的1-2-3-4。得到序列之后就用word2vector类似的方法训练得到embedding，没错，这个方法就是deepwalk。那么node2vector就是在每一步走的时候加上概率，比如有多大概率从1——>2，本质和deepwalk一样。所以说很简单的算法，非要取一个复杂的名字。

03 PyG实现

我得稍微介绍一下PyG中的数据结构。

通过说明，发现PyG中传进去的不是邻接表，而是节点对。这样其实很符合比赛的用法，最直接的应用就是推荐系统里的比赛，数据都是节点对形式储存。那其实是可以直接用的。那么我直接把PyG代码直接贴过来不过分吧

import os.path as osp

import torch
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
from torch_geometric.nn import Node2Vec


dataset = 'Cora'

dataset = Planetoid('./data', dataset)
data = dataset[0]

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = Node2Vec(data.edge_index, embedding_dim=128, walk_length=20,
                 context_size=10, walks_per_node=10,
                 num_negative_samples=1, p=1, q=1, sparse=True).to(device)

loader = model.loader(batch_size=128, shuffle=True, num_workers=4)
optimizer = torch.optim.SparseAdam(list(model.parameters()), lr=0.01)

def train():
    model.train()
    total_loss = 0
    for pos_rw, neg_rw in tqdm(loader):
        optimizer.zero_grad()
        loss = model.loss(pos_rw.to(device), neg_rw.to(device))
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    return total_loss / len(loader)

@torch.no_grad()
def test():
    model.eval()
    z = model()
    acc = model.test(z[data.train_mask], data.y[data.train_mask],
                     z[data.test_mask], data.y[data.test_mask],
                     max_iter=150)
    return acc

for epoch in range(1, 101):
    loss = train()
    acc = test()
    print(f'Epoch: {epoch:02d}, Loss: {loss:.4f}, Acc: {acc:.4f}')

@torch.no_grad()
def plot_points(colors):
    model.eval()
    z = model(torch.arange(data.num_nodes, device=device))
    z = TSNE(n_components=2).fit_transform(z.cpu().numpy())
    y = data.y.cpu().numpy()

    plt.figure(figsize=(8, 8))
    for i in range(dataset.num_classes):
        plt.scatter(z[y == i, 0], z[y == i, 1], s=20, color=colors[i])
    plt.axis('off')
    plt.show()

colors = [
    '#ffc0cb', '#bada55', '#008080', '#420420', '#7fe5f0', '#065535',
    '#ffd700'
]
plot_points(colors)

04 结语
最近半年可能都在做图相关的东西，就想起来去年打kdd时候，听见好多人用图，也听见了deepwalk这些名词，但是对于刚入门的自己好像还挺难的，然后这几天看了看PyG发现了好多这两年里论文的算法，所以即使刚入门的同学，之后再听见别人再用图的时候也不会束手无策了，可以选择PyG，当然还有DGL，国产paddle的PLG。