[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等），假设其特征值为$\lambda$，特征向量为$U$，通过与频谱图对比：

• $U$就可以类比为Fourier变换的basis函数;
• $\lambda$就类比为频率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 Data

edge_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 = $\varphi$; aggregation = $\square$; Update = $\gamma$. 那么替换上面的消息传递公式，得到如下的新形式:


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

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

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

原edge_index

加self_loop后的edge_index

2. Linearly transform node feature matrix. x = $\Theta W$ (代码里对应self.matmul(weight, x))
   原输入x
   
   经过weight transform得到的x
   

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

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

4. Sum up neighboring node features. ${\sum }_{i\in 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)的输出:

5. 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=max(-0.0075,-0.0106)$, $0.7586=max(0.5364,0.7586)$

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

import torch
from torch_scatter import scatter_add
from torch_geometric.nn import MessagePassing
import math

def 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_index


def 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_index
        print("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.bias
        else:
            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_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)
        
# 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