[PyG] 1.如何使用GCN完成一个最基本的训练过程(含GCN实现)

0. 前言

为啥要学习Pytorch-Geometric呢？(下文统一简称为PyG) 简单来说，是目前做的项目有用到，还有1个特点，就是相比NYU的DeepGraphLibrary, DGL的问题是API比较棘手，而且目前没有迁移的必要性。

图卷积框架能做的事情比较多，提供了很多方便的数据集和各种GNN SOTA的实现，其实最吸引我的就是这个framework的API比较友好，再加之使用PyG做项目的人比较多，生态对我这种做3D mesh的人比较友好。

注意, 本教程完全基于官方最新 （2020.04.14） 的教程，在此基础上，完成了简化版本的GCN的实现，对GCN的官方实现感兴趣的童鞋可以康康[1]。

下面，我将完全按照[1]的步骤来，不同之处在于，我在这里将基于PyG的最新版本(1.4.3)来分析GCN的简化版实现，让大家更加理解GCN的实现原理, 以下是阐述顺序:

①图数据的Data Handling
②Common Benchmark Datasets
③Mini-batches
④Data Transforms
⑤Learning Methods on Graphs

此外，我所使用的环境是:

Ubuntu 18.04

Cuda10.0

pytorch 1.4.0 conda install pytorch=1.4.0 cudatoolkit=10.0

pytorch geometric 1.4.3

torch-scatter pip install torch-scatter==latest+cu100 -f https://pytorch-geometric.com/whl/torch-1.4.0.html

torch-spline-conv pip install torch-spline-conv==latest+cu100 -f https://pytorch-geometric.com/whl/torch-1.4.0.html

torch-cluster pip install torch-cluster==latest+cu100 -f https://pytorch-geometric.com/whl/torch-1.4.0.html

torch-sparse pip install torch-sparse==latest+cu100 -f https://pytorch-geometric.com/whl/torch-1.4.0.html

1. 图结构的数据处理

首先，图是什么？图是边和点的相关关系的组合。在PyG中，一个简单的graph可以被描述为torch_geometric.data.Data[2]的实例，其中有几个重要的属性需要说明，此外，如果你的图需要扩展，那么你可以对torch_geometric.data.Data这个类进行修改即可。

图1.1 torch_geometric.data.Data的常用成员变量

通常来讲，对分一般的任务来说，数据类只需要有x,edge_index,edge_attr,y等几个属性即可，而且，这些属性都是optional(可选)的，也就是说，Data类并不局限于这些属性。

举个栗子，可以扩展data.face(torch.LongTensor, [3, num_faces])来保存3D mesh的三角形的连接关系.

图1.2 torch_geometric.data.Data的官方说明

图1.3 Data实例(3个节点,4条边(双向), 每个节点有2个特征[-1, 2], [0, 3], [1, 1].)

需要注意的是，尽管图只有2条边，我们还是需要定义4个index tuple来考虑边的双向关系。
图1.3搭建的graph的示意图如下:

2. 常见Benchmark数据集

尽管最近Bengio团队是基于DGL开发的6个Benchmark数据集，但是在pyG上做这个也没问题呀～。所以也不必直接因此就转去DGL。

PyTorch Geometric包含了大量的基础数据集, 所有的Planetoid datasets (Cora, Citeseer, Pubmed), 来自多特蒙德工大的清洗过的图分类数据集, 一系列3D点云和mesh的数据集，比如FAUST，ShapeNet等。

PyG提供了这些数据的自动下载，并将其处理成之前说的Data形式，以ENZYMES数据集为例（包含600个图和6个类别）:

图2.1 ENZYMES数据集的解析

由图2.1可见，其中的每个样本都是Data的instance,有顶点特征x，连接关系edge_index以及类别y 3个属性. 可以看出，ENZYMES的每个数据都是1个图。

注意: 可以通过使用dataset=dataset.shuffle()来对数据集进行shuffle。

此外，教程上还提供了Planetoid的Cora数据集的说明（用于semi-supervised graph node classification）, 这里Cora数据集的数据有3个新的属性train_mask, test_mask, val_mask, 这3个属性用于表征需要训练、测试和验证的数据节点。

Cora与ENZYMES的区别是，Cora中的每个数据是整个图中的1个节点，而ENZYMES的每个数据都是1个独立的图。

图2.2 Cora数据集说明

3. Mini-Batches

我们知道，神经网络通常是按Batch训练的，PyG通过创建稀疏的邻接矩阵(sparse block diagnol adjacency matrices)实现在mini-batch上的并行化。

图3.1 PyG mini-batch对不同的节点、边数量的图的批处理

并按照node dimension来拼接节点特征x和类别特征y。通过这种方式，PyG可以在一个Batch中塞进不同nodes和edges数的样本。

图3.2 ENZYMES数据集加载说明（未shuffle）

(注意，这里的DataLoader用的是PyG自己的，而不是pytorch的，此外，use_node_attr=False时, x为[nodes_num, 3]; use_node_attr=True时, x为[nodes_num, 21])

这里，torch_geometric.data.Batch继承自 torch_geometric.data.Data，多了一个名为batch的属性，其作用是标示每个节点属于哪个图(ENZYMES)/样本.

此外，torch_geometric.data.DataLoader也只是pytorch的Dataloader重写了collate函数的版本而已。

正常传递给pytorch的Dataloader的参数，如pin_memory,num_workers等都可以传给torch_geometric.data.DataLoader.

当然，用户可以通过使用torch-scatter[3]对节点数据特征x进行自定义的处理并使用自定义的Dataset和Dataloader来处理自己的特殊形式数据[4].

4. Data Transforms

同torchvision在pytorch中的使用类似，我们也需要对graph数据进行处理和变换。PyG提供了自己的transform方式和工具包，要求的输入为Data对象，并返回transformed的Data对象。

类似地，transform可以通过torch_geometric.transforms.Compose来进行一系列的拼接。

作者举得例子是ShapeNet数据集（包含17,000 3D shape point clouds and per point labels from 16 shape categories）的Airplane类，作者通过pre_transform = T.KNNGraph(k=6)将point cloud数据变为graph数据集。

图4.1 ShapeNet数据集处理（将点云数据变为graph数据）

如有其它需要，用户可以自己去torch_geometric.transforms进行查阅是否有符合自己目的的transform，没有的话自己写～

5. Learning Methods on Graphs

在搞定前4步后，现在让我们开始搞起第1个GNN～，这里，我们将会使用最基础的GCN层来复现Cora Citation数据集上的实验，若要理解GCN，需要从Fourier变换讲起，类比time domain --> frequency domain, 经过Hemlholtz公式，将vertex domain变到 spectral domain来分析，这样一来，vertex domain的卷积就变成了spectral domain的点乘，节省了计算量。

此外，变换的过程中，还涉及到Laplacian矩阵L的意义（每个vertex的散度Divergence：可以理解为每个vertex的信息的增益情况，出射为正，入射为负），因为L的性质（半正定，特征值大于等于0等），假设其特征值为 $λ$ ，特征向量为 $U$ ，通过与频谱图对比：

$U$ 就可以类比为Fourier变换的basis函数;
$λ$ 就类比为频率w

GCN[5]就是在此基础上，经由2步优化得到的，它既考虑了self-loop，也考虑了k-localize（局部性），还对度进行了renormalization，避免马太效应过于明显，使得模型不会很容易陷入local minima。

好了，就不再多提了，对GCN的推导和出现感兴趣的，可以看[6-7](先理解Laplacian矩阵和变换在图论中的一般含义，再去油管上看台湾大学姜成翰助教关于GNN的教程)进行学习，下面我们看代码。

5.1 GCN在PyG的实现

PyG提供了torch_geometric.nn.MessagePassing这个base class，通过继承这个类，我们可以实现各种基于消息传递的GNN，借由MessagePassing
用户不再需要关注message的progation的内部实现细节，MessagePassing主要关注其UPDATE, AGGREGATION, MESSAGE 这3个成员函数。

用户在实现自己的GNN时，一般只overwrite AGGREGATION, UPDATE这2个成员函数，MESSAGE/Propagate用MessagePassing自带的。（官方的GCN就是这样的～）

我们的目标是: 实现1个与官方一致的简化版的GCN，并通过实现它来掌握如何在PyG中定义图卷积。

首先，我们先定义一个图数据data(有向图, 4个节点，3条边, 每个节点的特征维度都是1, 值也都为1):

import torch
from torch_geometric.data import Dataedge_index = torch.tensor([[1, 2, 3], [0, 0, 0]], dtype=torch.long)
x = torch.tensor([[1], [1], [1], [1]], dtype=torch.float)data = Data(edge_index=edge_index, x=x)
print(edge_index)
print(data)

MessagePassing消息传递机制

通过上面这个图[9]，很容易了解到基于消息传递的GCN的每1块对应的内容是什么:
message = $ϕ\phi$ ; aggregation = $□$ ; Update = $γ\gamma$ . 那么替换上面的消息传递公式，得到如下的新形式:

此图就表示了本例的数据的流转方式，需要注意: GCN默认的scatter方式是add（至于为啥，请看下图: 因为用mean和max的情况下，每个subgraph随着GNN网络层数的加大，其中各个node之前的特征区分度越来越小，这不符合我们的目标）

GCN的实现也可以分为5步:

Add self-loops to the adjacency matrix. edge_index = $A^=IN+A\hat A = I_{N} + A$ (代码里通过修改edge_index实现), PyG源代码中通过add_remaining_self_loops函数来实现[10]

原edge_index

加self_loop后的edge_index

Linearly transform node feature matrix. x = $ΘW\Theta W$ (代码里对应self.matmul(weight, x))
原输入x

经过weight transform得到的x
Normalize node features. norm = $D^−0.5A^D^−0.5\hat D^{-0.5}\hat A \hat D^{-0.5}$ (在源码中 $A^\hat A$ 用以表示边的权重，默认情况下都是1.)

norm的值，其长度同edge_index的一致:

Sum up neighboring node features. $∑i∈N(p)(D^−0.5A^D^−0.5ΘW)\sum_{i ∈ N(p)}(\hat D^{-0.5}\hat A \hat D^{-0.5} \Theta W)$
(第4步在MessagePassing里面实现，即上面的图中的scatter_add/sum/mean, 用户无需操心) 在def message(self, x_j, norm)中的x_j就是第2步x的扩展到self_loop的结果.

x_j的值:

message(self, x_j, norm)的输出:

Return new node embeddings. 返回得到的结果 $X_{new}$ .

因为是scatter_add的方式，所以将[1, 0], [2, 0], [3, 0]的连接关系相加，得到最终输出结果:

显然
$- 0.0330 = - 0.0075 - 0.0075 - 0.0075 - 0.0106$
$2.3680 = 0.5364 + 0.5364 + 0.5364 + 0.7586$

同样，若改成scatter_max的话，结果为如下，因为 $- 0.0075 = m a x (- 0.0075, - 0.0106)$ , $0.7586 = m a x (0.5364, 0.7586)$

这五步的实现通过如下代码完整实现:

import torch
from torch_scatter import scatter_add
from torch_geometric.nn import MessagePassing
import mathdef glorot(tensor):if tensor is not None:stdv = math.sqrt(6.0 / (tensor.size(-2) + tensor.size(-1)))tensor.data.uniform_(-stdv, stdv)def zeros(tensor):if tensor is not None:tensor.data.fill_(0)def add_self_loops(edge_index, num_nodes=None):print("进入self_loops")loop_index = torch.arange(0, num_nodes, dtype=torch.long,device=edge_index.device)print(loop_index)loop_index = loop_index.unsqueeze(0).repeat(2, 1)print(loop_index)edge_index = torch.cat([edge_index, loop_index], dim=1)print(edge_index)print("出self_loops")# 原来的edge_index为[[1, 2, 3],#                   [0, 0, 0]]#  这样一来，就在原来的边连接关系edge_index的基础上增加了self_loop的关系.#  torch.cat([edge_index, loop_index], dim=1)#      tensor([[1, 2, 3, 0, 1, 2, 3],#              [0, 0, 0, 0, 1, 2, 3]])return edge_indexdef degree(index, num_nodes=None, dtype=None):out = torch.zeros((num_nodes), dtype=dtype, device=index.device)print(out.scatter_add_(0, index, out.new_ones((index.size(0)))))return out.scatter_add_(0, index, out.new_ones((index.size(0))))class GCNConv(MessagePassing):def __init__(self, in_channels, out_channels, bias=True):super(GCNConv, self).__init__(aggr='add')  # "Add" aggregation.# super(GCNConv, self).__init__(aggr='max')  # "Max" aggregation.self.weight = torch.nn.Parameter(torch.Tensor(in_channels, out_channels))if bias:self.bias = torch.nn.Parameter(torch.Tensor(out_channels))else:self.register_parameter('bias', None)self.reset_parameters()def reset_parameters(self):glorot(self.weight)zeros(self.bias)def forward(self, x, edge_index):# x has shape [N, in_channels]# edge_index has shape [2, E]# Step 1: 为adjacency matrix添加self_loop(通过对edge_index拼接连向自己的边[1, 1], [2, 2]等)# 原来的edge_index = tensor([[1, 2, 3],#                           [0, 0, 0]])# 加上self_loop的index = tensor([[1, 2, 3, 0, 1, 2, 3],#                               [0, 0, 0, 0, 1, 2, 3]])edge_index = add_self_loops(edge_index, x.size(0))# Step 2: 对输入的node feature matrix进行weight transform.x = torch.matmul(x, self.weight)# Step 3-5: 开始消息传递.edge_weight = torch.ones((edge_index.size(1),), dtype=x.dtype,device=edge_index.device)row, col = edge_indexprint("row", row)  # row tensor([1, 2, 3, 0, 1, 2, 3])print("col", col)  # col tensor([0, 0, 0, 0, 1, 2, 3])deg = scatter_add(edge_weight, row, dim=0, dim_size=x.size(0))print("deg", deg)  # deg是[1, 2, 2, 2], 这是啥?# 因为# row = [1, 2, 3, 0, 1, 2, 3]# edge_weight = [1, 1, 1, 1, 1, 1, 1]# 所以，主对角上，第0个对应1，第1个对应2个，同理，得到degree矩阵. 这里只返回主对角的元素, 避免稀疏乘.deg_inv_sqrt = deg.pow(-0.5)deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0# 读edge_weight为None的情况,# deg_inv_sqrt[row] * edge_weight *  deg_inv_sqrt[col] == deg_inv_sqrt[row] *  deg_inv_sqrt[col]norm = deg_inv_sqrt[row] * edge_weight *  deg_inv_sqrt[col]print(norm)# norm = tensor([0.7071, 0.7071, 0.7071, 1.0000, 0.5000, 0.5000, 0.5000])return self.propagate(edge_index, x=x, norm=norm)           def message(self, x_j, norm):# x_j has shape [E, out_channels]# norm: 规则化后的权重.return norm.view(-1, 1) * x_j if norm is not None else x_j                  def update(self, aggr_out):# aggr_out has shape [N, out_channels]# Step 5: 返回新的node embeddings.if self.bias is not None:return aggr_out + self.biaselse:return aggr_out

进行实验，得到与官方实现一样的效果:

5.2 在Cora Citation数据集上进行训练

这里用回官方的GCN来做1个2层的GNN网络对Cora Citation数据集进行训练，如果一切ok，下面代码直接复制到你的本地就可以跑起来～

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv
from torch_geometric.datasets import Planetoid# 5.1) 加载Cora数据集.(自动帮你下载)
dataset = Planetoid(root='/home/pyG/Cora', name='Cora')# 5.2) 定义2层GCN的网络.
class Net(torch.nn.Module):def __init__(self):super(Net, self).__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_indexx = 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)# 5.3) 训练 & 测试.
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Net().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).max(dim=1)
correct = float (pred[data.test_mask].eq(data.y[data.test_mask]).sum().item())
acc = correct / data.test_mask.sum().item()
print('Accuracy: {:.4f}'.format(acc))
# >>> Accuracy: 0.8150

到这步，一个完整的基于GCN的GNN就搞定了，至于训练的数据处理和很多细节，需要大家hack源码啦，祝大家学习愉快～

参考资料

[1] PyG官方Tutorial
[2] torch_geometric.data.Data
[3] torch-scatter
[4] advanced mini-batching of PyG
[5] GCN: Semi-supervised Classfication with Graph Convolutional Networks
[6] [其实贼简单] 拉普拉斯算子和拉普拉斯矩
[7] GNN介绍台湾大学姜成翰
[8] Torch geometric GCNConv 源码分析
[9] MessagePassing
[10] add_remaining_self_loops