torch_geometric(pyg)的介绍和简单使用

前言

最近做毕业设计，需要用到图神经网络（以下简称GNN）。由于刚入门GNN，不想看大段的公式和相关论文（然而事实证明该看的永远逃不了），所以怎么办？百度上找呗！因为自己平时用pytorch比较多，所以找到了基于pytorch的图神经网络库，pytorch_geometric（以下简称pyg）。在用这个库的过程中，由于这个库“约定大于配置”的一些特性，也遇到了许多坑，而中文资料中，大多都是直接翻译文档，对一些细节没有做解释。因此将整个过程记录下来，供大家未来参考。

图神经网络基本概念

考虑到有许多同学和我最开始一样，就是想知道GNN是干什么的，大致原理是什么，因此这一章将简单介绍GNN

GNN是干什么的

相信大家都知道“图”是什么东西，一堆的节点（每个节点都有各自的特征），中间用箭头连起来。那么我们就很自然地想，这些相连的节点，可能信息上比较相关，可以互补，那我们能不能对这些相关的节点做一定的特征聚合操作呢？比如求和（这就是GIN），取平均（这就是大家常说的GCN），加权求和（这就是GAN）？做特征聚合时，可以只对一阶相邻（也就是直接相邻）节点做聚合，也可以对二阶相邻（最多可以通过一个中间节点间接相连）节点做聚合，由此衍生出各种花里胡哨的GNN……

所以我们说，GNN的核心操作是节点的特征聚合，具体怎么聚合，各个GNN有自己的花样。但是这些花样一定都是依据邻接关系得到的。毕竟脱离了图结构，GNN也没有存在的意义了，对吧？
$$\begin{gathered} X^{l+1}=aggregate(X^l,A) \\ A表示邻接矩阵，X^l表示第l层的输出 \end{gathered}$$
看到这个，我们可以联想一下CNN，一个3x3的卷积核，实际上就是对8邻域做了加权和，对吧？如果把二维图像看成是一张8邻域连通的图，那实际上就是GNN了。

除了聚合特征之外，GNN们通常还带有一个单节点的特征变换运算，这个运算可以是单纯的线性变换，可以带有非线性激活函数等，可以认为是对单节点的特征做了增强。所以GNN的大致运算过程可以写为
$$\begin{gathered} X^{l+1}=aggregate(f(X^l),A) \\ {X}^{l+1}=f(aggregate(X^l,A)) \end{gathered}$$
这两种写法并不会影响输入输出的维度，重要的是明白有一次特征聚合，有一次特征增强即可。

pyg库介绍

pyg的功能比较强大，包括一些utils包下的图级别的工具函数，nn包下常见的GNN层，data包下封装好的图特征对象，loader包下封装好的图batch的loader，甚至还有端到端的GNN模型。自己搭GNN，主要会用到Data、DataLoader、nn、utils等相关工具。因此下面将从搭建一个最简单的GNN，并完成一次输入输出运算为主线，介绍这些工具的用法，以及一些暗箱约定（或者说，坑）

以下的内容主要根据pyg的官方教程梳理，对其中重要的坑会进行强调。

那么，想要使用三方库完成一次GNN运算，我们就必须了解以下的一些内容

• 输入数据是什么格式？尤其矩阵的维度是什么样的？
• DataLoader对数据进行了怎样的变换，模型又该如何处理一个mini-batch？
• 内置GNN层的输入输出又是什么？可以直接把mini-batch给它吗？

接下来从探究以上问题出发，我们力求把pyg的用法讲明白

输入数据的格式

Data对象

源教程

我们知道GNN的输入除了各顶点的特征之外，还有邻接矩阵，甚至还会有边的特征。pyg内置了Data对象，用于封装GNN的输入。Data对象中最常使用的几个属性包括以下三个。如果实际科研工作中需要使用更复杂的特征，可以回到上面找源教程。

• data.x: 节点特征矩阵，维度是 [num_nodes, num_node_features]
• data.edge_index: 图连接关系，也就是之前所说的邻接矩阵。只不过这里采用了稀疏格式的输入，维度是 [2, num_edges] 类型是 torch.long。也就是只存储每条边的出发点和终止点，而不是真正的邻接矩阵（这样的矩阵在顶点多边少的时候，非常占内存）。
• data.y: 模型的期望输出。如果是完成节点级别任务的GNN，维度一般为[num_nodes, *] ；如果是完成图级别任务的GNN，维度一般为 [1, *]

DataSet对象

源教程

DataSet对象通常被我们用于原始数据读取和加工，将数据转换成DataLoader所能接受的输入。

说的现实点，主要就是把我们的数据源转换成一系列的Data对象。这个对象本身也是pyg封装好的，需要我们削足适履，把我们的转换逻辑填进去。以下直接摘自官方教程

class MyOwnDataset(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super().__init__(root, transform, pre_transform)
        # 读取已经转换好格式的数据
        self.data, self.slices = torch.load(self.processed_paths[0])
	
    # 未处理好的数据，如果pyg发现文件不存在，会进行下载
    @property
    def raw_file_names(self):
        return ['some_file_1', 'some_file_2', ...]
	
    # 处理好的数据，如果pyg发现文件不存在，会调用self.process()函数
    @property
    def processed_file_names(self):
        return ['data.pt']
	
    # 下载原始数据的函数，如果不需要就直接pass
    def download(self):
        download_url(url, self.raw_dir)

    # 将原始数据进行处理，转换为Data对象，并且保存下来
    def process(self):
        # 假设这里已经经过了一系列处理，得到了包含Data对象的List
        data_list = [...]

        # 雷打不动的两句话，处理List并存盘
        data, slices = self.collate(data_list)
        torch.save((data, slices), self.processed_paths[0])

DataLoader对象

源教程

DataLoader打包mini-batch

玩深度学习的同学们都知道，训练模型一般要把多个数据打包成一个mini-batch，再丢给模型训练（原因我就不解释了）。DataLoader就是完成这个工作的。如果你恰好用过pytorch，你肯定也知道pytorch默认的DataLoader会把batch_size个样本打包成[batch_size, d1, d2, ..., dn]维度的输入，其中[d1, d2, ..., dn]是样本本来的特征维度。

而pyg的DataLoader，会把数据打包成[batch_size*num_nodes, num_node_features]的维度（也就是batch_size不会单独成一维）。说实话，这一点非常坑（当然，从性能的角度，也可以说“妙”）。pyg官方的解释是“为了增加并行度”，那么，这一步操作是怎么增加并行度的？

如何提高并行度

稍加思考就可以明白。之前我们就说过，GNN最重要的操作之一，就是进行特征聚合。那进行特征聚合的代码怎么写？我们以求均值为例，最暴力的，当然是对着邻接矩阵，一个一个把邻居的特征加起来再取平均了
$$\begin{gathered} X_i^{l+1}=(\sum _j{a}_{ij}{x}_{ij})/n \\ {X}_i=X[i,:] \\ {x}_{ij}=X[i,j] \end{gathered}$$
当然，大家都知道拿循环来算加权和，效率非常低，因此应该用矩阵的形式来表示这一运算。假设我们已经算出了各个节点加权和的系数，形成一个系数矩阵，那上面的循环直接就可以用一个矩阵乘法表示了
$$\begin{gathered} X^{l+1}=A\cdot X^l \\ A\in [num\_nodes,num\_nodes] \\ X\in [num\_nodes,num\_node\_features] \end{gathered}$$
可以想象，A乘在左边，就是对X做了行变换，也就是对X的每一行进行了加权和。

图神经网络的计算效率是比较低的，多张图之间难以进行并行化。假如num_nodes不是很大，那进行一次上述的运算，也不会有太大的加速。因此pyg的DadaLoader将数据打包成了[batch_size*num_nodes, num_node_features]维度，相当于大大提升了参与一次图运算的顶点数，因此可以充分利用向量运算的优势。

mini-batch的内容

mini-batch内打包了节点特征、样本标签、连接关系、batch信息等内容

我们现在以最简单的两个图为例，说明打包后的数据长什么样

• 图1：三个顶点[0, 1, 2]，三个顶点两两双向连接，标签是节点维度的
• 图2：两个顶点[0,1]，两个顶点两两双向连接，标签是节点维度的

接下来，特征、标签都正常拼接；但是节点连接关系会进行一定的运算

• data.x: 将两个图的节点特征直接拼接成[5, num_node_features]的矩阵
• data.y: 将两个图的节点标签直接拼接成[5, *]的矩阵
• data.edge_index: 将两个图混合成一张大图，形成[2, 10]的矩阵。得到的邻接矩阵大概是[[0,1],[1,0],[0,2],[2,0],[1,2],[2,1],[3,4],[4,3]]。这里为了看着方便，我把稀疏邻接矩阵转置了一下，实际上它的维度是[2, num_edges]

欸，哪来的节点3和4呢？这就是这一并行化算法巧妙的地方，它将多个图融合成一张大图——其实图的编号没有太大的实际意义，它只是表达哪几个节点需要进行信息交换，需要把X的哪几行进行交换罢了。因此整个运算结果是非常正确的。

因此，我们再回头看一眼，Data对象中，各个元素的维度——知道为什么num_nodes一定作为第一维了吗？想要最大限度地利用pyg库带来的遍历，就必须削足适履，迎合它的编码方式。在下一篇文章中，我将讲述如何定制DataLoader从而增加一些灵活性。

除了以上所述的那些内容，DataLoader还会打包一个batch信息。这一信息主要是为了从batch中再区分出各个图所用，在进行一些图级别的全局运算，比如softmax，比如global_average，比如global_max，我们肯定希望是在一个图样本中进行（不然在这张用于运算的大图上进行全局计算，有什么实际意义吗）

一个简单的GNN例子

源教程

接下来我们搭一个最简单的两层GCN网络

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv

class GCN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = GCNConv(dataset.num_node_features, 16)
        self.conv2 = GCNConv(16, dataset.num_classes)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index

        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index)

        return F.log_softmax(x, dim=1)

接下来是训练过程

from torch_geometric.datasets import Planetoid

dataset = Planetoid(root='/tmp/Cora', name='Cora')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = GCN().to(device)
data = dataset[0].to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

model.train()
for epoch in range(200):
    optimizer.zero_grad()
    out = model(data)
    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()

最后是测试过程

model.eval()
pred = model(data).argmax(dim=1)
correct = (pred[data.test_mask] == data.y[data.test_mask]).sum()
acc = int(correct) / int(data.test_mask.sum())
print(f'Accuracy: {acc:.4f}')

好吧，这一章很水，毕竟这只是一个最简单的GNN例子，pyg也内置了很多的GNN模型供大家调用。需要注意的是，在刚刚的例程中，其实pyg内置GNN层的输入，并不是data，而是分立的data.x和data.edge_index，很多其它GNN层也是如此。这是因为通常大家先会用基本的GNN层搭建一些小模块，这些模块内可能带着卷积，可能带着池化，而对于卷积层来说，它并不需要知道batch信息。

后记

这一篇文章基本在翻译教程的过程中写完了，加上了自己在构建DataSet对象和DataLoader对象中踩的坑。但是这不是全部，后续，我将从一个交通领域的T-GCN模型出发，讲述如何使用pyg库复现这一时空图神经网络模型。