文章目录

• 前言
• 一、导入相关库
• 二、加载Cora数据集
• 三、定义Node2Vec
• 四、定义模型
• 五、模型训练
• 六、可视化
• 完整代码

---

前言

大家好，我是阿光。

本专栏整理了《图神经网络代码实战》，内包含了不同图神经网络的相关代码实现（PyG以及自实现），理论与实践相结合，如GCN、GAT、GraphSAGE等经典图网络，每一个代码实例都附带有完整的代码。

正在更新中~ ✨

🚨 我的项目环境：

• 平台：Windows10
• 语言环境：python3.7
• 编译器：PyCharm
• PyTorch版本：1.11.0
• PyG版本：2.1.0

💥 项目专栏：【图神经网络代码实战目录】

---

本文我们将使用Pytorch + Pytorch Geometric来简易实现一个Node2Vec，让新手可以理解如何PyG来搭建一个简易的图网络实例demo。

一、导入相关库

本项目我们需要结合两个库，一个是Pytorch，因为还需要按照torch的网络搭建模型进行书写，第二个是PyG，因为在torch中并没有关于图网络层的定义，所以需要torch_geometric这个库来定义一些图层。

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

二、加载Cora数据集

本文使用的数据集是比较经典的Cora数据集，它是一个根据科学论文之间相互引用关系而构建的Graph数据集合，论文分为7类，共2708篇。

• Genetic_Algorithms
• Neural_Networks
• Probabilistic_Methods
• Reinforcement_Learning
• Rule_Learning
• Theory

这个数据集是一个用于图节点分类的任务，数据集中只有一张图，这张图中含有2708个节点，10556条边，每个节点的特征维度为1433。

# 1.加载Cora数据集
dataset = Planetoid(root='./data/Cora', name='Cora')

本项目是使用 Node2Vec 来生成每个节点的特征，所以对于原始节点特征是无用的，本项目只是单纯利用 Cora 数据集的节点空间关系，也就是 edge_index ，基于节点的空间关系来生成对应的节点特征，最终验证生成的节点特征效果如何。

三、定义Node2Vec

这里我们就不重点介绍Node2Vec了，相信大家能够掌握基本原理，本文我们使用的是PyG定义这个网络，在PyG中已经定义好了 Node2Vec 这个层，该层采用的就是 Node2Vec 机制。

对于Node2Vec的常用参数：

• edge_index：图数据中的边关系，就是 data 的 edge_index，形状为【2，num_edges】
• embedding_dim：每个节点形成的嵌入维度
• walk_length：游走形成的序列长度
• context_size：上下文大小
• walks_per_node：每个节点形成多少个游走序列
• p：在游走时重新访问某个节点的概率，默认为1
• q：在BFS策略和DFS策略之间的控制参数，默认为1
• num_negative_samples：每个正样本对应的负样本数，默认为1
• num_nodes：图的节点数目
• sparse：如果设置为True，权重矩阵的梯度是以稀疏矩阵方式存储，默认为False

# node2vec模型
model = Node2Vec(edge_index=data.edge_index,
                 embedding_dim=128, # 节点维度嵌入长度
                 walk_length=5, # 序列游走长度
                 context_size=4, # 上下文大小
                 walks_per_node=1, # 每个节点游走10个序列
                 p=1,
                 q=1,
                 sparse=True # 权重设置为稀疏矩阵
                ).to(device)

上面我们定义了一个 node2vec 模型，对于第一个参数就是 Cora 的边集，第二个参数就是经过训练每个节点的Embedding的维度，第三个参数是序列游走长度，这里设置的为5，也就是经过游走会生成长度为5的节点序列，例如【0，1284，345，94，2031】，对于 context_size 这个参数代表上下文大小，如果设置为4，针对上面数据我们会生成两个向量用于训练，分别是【0，1284，345，94】和【1284，345，94，2031】，可以看到就是滑动了一下。

感兴趣可以查看下源码：

对于 p 和 q 就是在游走时设定的概率超参，可以根据实际情况进行调整，如果都是1就是 DeepWalk 机制，walks_per_node 这个参数代表每个节点会生成多少个游走序列，值越大，模型更加稳定，生成的 Embedding 更加健壮。

四、定义模型

下面就是定义了一些模型需要的参数，像学习率、迭代次数这些超参数，然后是模型的定义以及优化器及损失函数的定义，和pytorch定义网络是一样的。

# 迭代器
loader = model.loader(batch_size=128, shuffle=True)
# 优化器
optimizer = torch.optim.SparseAdam(model.parameters(), lr=0.01)

五、模型训练

模型训练部分也是和pytorch定义网络一样，因为都是需要经过前向传播、反向传播这些过程，对于损失、精度这些指标可以自己添加。

# 3.开始训练
model.train()

for epoch in range(1, 101):
    total_loss = 0 # 每个epoch的总损失
    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()
    
    # 使用逻辑回归任务进行测试生成的embedding效果
    with torch.no_grad():
        model.eval() # 开启测试模式
        z = model() # 获取权重系数，也就是embedding向量表
        
        # z[data.train_mask] 获取训练集节点的embedding向量
        acc = model.test(z[data.train_mask], data.y[data.train_mask],
                         z[data.test_mask], data.y[data.test_mask],
                         max_iter=150) # 内部使用LogisticRegression进行分类测试
    
    # 打印指标
    print(f'Epoch: {epoch:02d}, Loss: {total_loss:.4f}, Acc: {acc:.4f}')

在训练的过程中，我们会进行验证生成的节点Embedding效果如何，可以利用 Node2Vec.test() 这个函数，这个函数内部实现了 LogisticRegression() 来实现分类任务，它就是为了验证生成的 Embedding 在下游的任务中效果如何，如果有其它下游任务（就是利用生成的节点Embedding做什么任务，常见就是节点分类），可以改成对应代码。

六、可视化

上面我们以经训练好了 Node2Vec 这个模型，通过调用 model() 即可获得内部的权重矩阵，也就是我们要的Embedding向量表（lookup table）。

生成好每个节点的 Embedding，我们可以通过可视化的方式更加直观的看到效果如何，对于可视化操作我们利用的是 TSNE 这个模块来进行降维，因为绘制二维图形需要x轴和y轴坐标（即二维），降到两个维度后，就获得了每个节点的坐标信息，然后利用 matplotlib 这个库来绘制不同类别的节点信息。

# 可视化节点的embedding
with torch.no_grad():
    # 不同类别节点对应的颜色信息
    colors = [
            '#ffc0cb', '#bada55', '#008080', '#420420', '#7fe5f0', '#065535',
            '#ffd700'
        ]

    model.eval() # 开启测试模式
    # 获取节点的embedding向量，形状为[num_nodes, embedding_dim]
    z = model(torch.arange(data.num_nodes, device=device))
    # 使用TSNE先进行数据降维，形状为[num_nodes, 2]
    z = TSNE(n_components=2).fit_transform(z.detach().numpy())
    y = data.y.detach().numpy()

    plt.figure(figsize=(8, 8))
    
    # 绘制不同类别的节点
    for i in range(dataset.num_classes):
        # z[y==0, 0] 和 z[y==0, 1] 分别代表第一个类的节点的x轴和y轴的坐标
        plt.scatter(z[y == i, 0], z[y == i, 1], s=20, color=colors[i])
    plt.axis('off')
    plt.show()

完整代码

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

# 1.加载Cora数据集
dataset = Planetoid(root='../data/Cora', name='Cora')
data = dataset[0]

# 2.定义模型
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # 设备

# node2vec模型
model = Node2Vec(edge_index=data.edge_index,
                 embedding_dim=128, # 节点维度嵌入长度
                 walk_length=5, # 序列游走长度
                 context_size=4, # 上下文大小
                 walks_per_node=1, # 每个节点游走1个序列
                 p=1,
                 q=1,
                 sparse=True # 权重设置为稀疏矩阵
                ).to(device)

# 迭代器
loader = model.loader(batch_size=128, shuffle=True)
# 优化器
optimizer = torch.optim.SparseAdam(model.parameters(), lr=0.01)

# 3.开始训练
model.train()

for epoch in range(1, 101):
    total_loss = 0 # 每个epoch的总损失
    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()
    
    # 使用逻辑回归任务进行测试生成的embedding效果
    with torch.no_grad():
        model.eval() # 开启测试模式
        z = model() # 获取权重系数，也就是embedding向量表
        
        # z[data.train_mask] 获取训练集节点的embedding向量
        acc = model.test(z[data.train_mask], data.y[data.train_mask],
                         z[data.test_mask], data.y[data.test_mask],
                         max_iter=150) # 内部使用LogisticRegression进行分类测试
    
    # 打印指标
    print(f'Epoch: {epoch:02d}, Loss: {total_loss:.4f}, Acc: {acc:.4f}')

# 可视化节点的embedding
with torch.no_grad():
    # 不同类别节点对应的颜色信息
    colors = [
            '#ffc0cb', '#bada55', '#008080', '#420420', '#7fe5f0', '#065535',
            '#ffd700'
        ]

    model.eval() # 开启测试模式
    # 获取节点的embedding向量，形状为[num_nodes, embedding_dim]
    z = model(torch.arange(data.num_nodes, device=device))
    # 使用TSNE先进行数据降维，形状为[num_nodes, 2]
    z = TSNE(n_components=2).fit_transform(z.detach().numpy())
    y = data.y.detach().numpy()

    plt.figure(figsize=(8, 8))
    
    # 绘制不同类别的节点
    for i in range(dataset.num_classes):
        # z[y==0, 0] 和 z[y==0, 1] 分别代表第一个类的节点的x轴和y轴的坐标
        plt.scatter(z[y == i, 0], z[y == i, 1], s=20, color=colors[i])
    plt.axis('off')
    plt.show()