使用 PyG 进行图神经网络训练
使用 PyG 进行图神经网络训练
使用 PyG 进行图神经网络训练
前言
最近一直在想创新点，搭模型，想尝试一下图神经网络，想着自己实现一个，但是之前也没有尝试过写 GNN 模型，对其中的实现细节也没有实际尝试过，最后找到了 PyG ，尝试一下之后发现还是挺简单的，也比较好拿到现有模型里面，于是开始挖坑。
PyG (PyTorch Geometric) 是一个基于 PyTorch 的库，可轻松编写和训练图形神经网络 (GNN)，用于与结构化数据相关的广泛应用。
目前网上对 PyG 的相关文档并不多，大本部也都是比较重复的内容，因此我主要参考的还是官方文档。具体安装方式参考 PyG Installation。
图结构
建图
首先，我们需要根据数据集进行建图，在 PyG 中，一个 Graph 的通过torch_geometric.data.Data进行实例化，它包括下面两个最主要的属性：
data.x: 节点的特征矩阵，形状为[num_nodes, num_node_features];data.edge_index: 图的边索引，用 COO 稀疏矩阵格式保存。形状为[2, num_edges];
稀疏矩阵是数值计算中普遍存在的一类矩阵，主要特点是绝大部分的矩阵元为零。COO (Coordinate) 格式是将矩阵中的非零元素存储，每一个元素用一个三元组来表示，分别是（行号，列号，数值）
上面两个属性也就是整个图中最重要的部分。同时你也可以定义下面的一些其他信息：
data.edge_attr: 边的特征矩阵，形状为[num_edges, num_edge_features];data.y: 计算损失所需的目标数据，比如节点级别的形状为[num_nodes, *]，或者图级别的形状为[1, *];
现在，我们来简单创建一张图：
import torch
from torch_geometric.data import Data
edge_index = torch.tensor([[0, 1, 1, 2],
[1, 0, 2, 1]], dtype=torch.long)
x = torch.tensor([[-1], [0], [1]], dtype=torch.float)
data = Data(x=x, edge_index=edge_index)
>>> Data(edge_index=[2, 4], x=[3, 1])需要注意的是：
- 第一行
edge_index[0]表示起点，第二行edge_index[1]表示终点； - 虽然只有两条边，但在 PyG 中处理无向图时实际上是互为头尾节点；
- 矩阵中的值表示索引，指代
x中的节点。
实际上，你不需要找对专门的属性去存取某些信息，你完全可以「自定义属性名」，并放入你想要的内容，事实上就像一个字典。例如，你想要为图再添加一个 freq 属性以表示访问频率：
edge_index = torch.tensor([[0, 1, 1, 2],
[1, 0, 2, 1]], dtype=torch.long)
x = torch.tensor([[-1], [0], [1]], dtype=torch.float)
data = Data(x=x, edge_index=edge_index)
data.freq = torch.tensor([2, 5, 1], dtype=torch.long)
>>> Data(x=[3, 1], edge_index=[2, 4], freq=[3])此外，这里再列举一些个人觉得比较常用的操作，感觉官方文档有些地方不清不楚，完整的类说明可以看这里：
coalesce(): 对edge_index中的边排序并去重clone(): 创建副本to(cuda:0): 把数据放到 GPUnum_edges: 返回边数量num_nodes: 返回节点数量
关于 train_mask
在官方的数据集中，划分「训练集」和「测试集」的方式是创建一张大图，然后指定训练节点以及测试节点，通过 train_mask 和 test_mask 来实现。
from torch_geometric.datasets import Planetoid
dataset = Planetoid(root='/tmp/Cora', name='Cora')
data = dataset[0]
>>> Data(edge_index=[2, 10556], test_mask=[2708],
train_mask=[2708], val_mask=[2708], x=[2708, 1433], y=[2708])
data.is_undirected()
>>> True
data.train_mask.sum().item()
>>> 140
data.val_mask.sum().item()
>>> 500
data.test_mask.sum().item()
>>> 1000上面是官方提供的数据集中的一个例子，但个人觉得并不是所有的数据集都适合这种方法，我仍习惯于训练集和测试集分开，也就是创建两张图，因为大多数情况下都需要自定义一个 MyDataset，你可以把建图的操作放在里面，所以其实感觉这样还更方便。
data_train, data_test = split_train_test(data)
# Create dataset
dataset_train = MyDataset(data_train)
dataset_test = MyDataset(data_test)当然，这还是要看具体做什么任务的，对某些任务来说可能只需要一张大图，对我来说可能就需要分成两张。
关于 Embedding
最开始的时候我有想过，为什么要在一开始创建 x 的时候就让我把节点的维度给定下来，这不应该是我后面模型里面 Embedding 的时候再做的事情吗，难不成建图的时候就要 Embedding？当然，对于有些情况，建图的时候完全就可以直接 Embedding。
事实上，其实就是我想多了，其实把它当做一个字典就好。
import torch
from torch import nn
from torch_geometric.data import Data
# create graph
edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]], dtype=torch.long)
x = torch.tensor([1, 0, 2], dtype=torch.long)
data = Data(x=x, edge_index=edge_index)
data.freq = torch.tensor([2, 5, 1], dtype=torch.long)
...
# model.py
# init embedding layer
emb_layer = nn.Embedding(3, 8)
data.x = emb_layer(data.x)
# data.x_emb = emb_layer(data.x)
print(data)
>>> Data(x=[3, 8], edge_index=[2, 4], freq=[3])就像上面展示的这样，我的做法是先把节点的 ID 填进去，接着在模型里面进行 Embedding，当然你可以直接使用 data.x = emb_layer(data.x) 把原来的 ID 给替换掉；也可能你需要保留 ID，那么就可以把它放到一个新的属性中，比如 data.x_emb。总的来说其实就是不需要把 data 的属性看得那么死，完全可以按照你的需要来。
关于 Batch 和 DataLoader
Batch
在 PyG 中，你可以通过手动的方式进行批量化操作：
import torch
from torch import nn
from torch_geometric.data import Data, Batch
edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]], dtype=torch.long)
x = torch.tensor([1, 0, 2], dtype=torch.long)
g1 = Data(x=x, edge_index=edge_index)
g2 = Data(x=x, edge_index=edge_index)
g3 = Data(x=x, edge_index=edge_index)
data = [g1, g2, g3]
>>> Data(x=[3], edge_index=[2, 4])
batch = Batch().from_data_list(data)
>>> DataDataBatch(x=[9], edge_index=[2, 12], batch=[9], ptr=[4])在上面的例子中，为了方便我创建了 3 张相同的图。现在我们来看看 batch 具体做了些什么：
print(batch.x)
>>> tensor([1, 0, 2, 1, 0, 2, 1, 0, 2])
print(batch.edge_index)
>>> tensor([[0, 1, 1, 2, 3, 4, 4, 5, 6, 7, 7, 8],
[1, 0, 2, 1, 4, 3, 5, 4, 7, 6, 8, 7]])
print(batch.batch)
>>> tensor([0, 0, 0, 1, 1, 1, 2, 2, 2])
print(batch.ptr)
>>> tensor([0, 3, 6, 9])实际上，Batch 的本质就是将多个图组合起来，创建一张大图。可以看出，首先在 x 中将所有节点拼接在一起，batch 指明了每个节点所属的批次，ptr 指明了每个 batch 的节点的起始索引号，然后在 edge_index 中对边进行拼接，同时为了避免冲突对索引加上了 ptr。
用公式来表示也就是：
DataLoader
PyTorch 原生的 DataLoader 实际上对 Data 并不支持，虽然可以创建成功，但在遍历取数据的时候，你会发现如下错误：
default_collate: batch must contain tensors, numpy arrays, numbers, dicts or lists; found <class ‘torch_geometric.data.data.Data’>
PyG 有一个自己的 DataLoader，实际上只需要用它替换 PyTorch 原生的 DataLoader 就可以，个人觉得使用体验上和 PyTorch 差别不大。
from torch.utils.data import Dataset
# from torch.utils.data import DataLoader
from torch_geometric.loader import DataLoader
class MyDataset(Dataset):
def __init__(self, data):
super(MyDataset, self).__init__()
self.data = data
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
sub = self.data[idx]
return sub
dataset = MyDataset(data)
dataloader = DataLoader(dataset)图神经网络
讲完了图结构，以及数据集之后，现在正式进入到了模型训练阶段
Convolutional Layers
PyG 其实定义了非常多可供直接使用的 Convolutional Layers，具体你可以看这里
… 待施工 …