pytorch gather函数有什么用 pytorch geometric gat

开始使用Markdown编辑器写博客笔记。

1. PyTorch Geometric配置

PyG的配置比预期要麻烦一点。PyG只支持两种Cuda版本，分别是Cuda9.2和Cuda10.1。而我的笔记本配置是Cuda10.0，考虑到我Pytorch版本是1.2.0+cu92，不是最新的，因此选择使用Cuda9.2的PyG 1.2.0（Cuda向下兼容）。按照PyG官网的安装教程，需要安装torch-scatter，torch-sparse（稀疏图需要稀疏矩阵），和torch-cluster（FPS采样和KNN图构建等）。安装这些库出了问题，会报Microsoft Visual Studio相关的错误。试遍网上安装教程都没有效果，最后发现安装低版本的torch-sparse（0.4.0）和torch-cluster（1.3.0）后，就可以顺利安装PyG 1.2.0。使用的时候下载TU数据集会报错ssl.CertificateError: hostname 'ls11-www.cs.uni-dortmund.de' doesn't match 'ls11-www.cs.tu-dortmund.de'，

dataset = TUDataset(root='ENZYMES', name='ENZYMES')

最后发现tu_dataset.py文件中url有问题（使用tu-dortmund后就能正常下载TU数据集）：

# 原下载源网址
    #url = 'https://ls11-www.cs.uni-dortmund.de/people/morris/' \
    #      'graphkerneldatasets'
    # 修改后下载源网址
    url = 'https://ls11-www.cs.tu-dortmund.de/people/morris/' \
          'graphkerneldatasets'

使用的时候发现，即便url是对的，但是下载其他数据集还是会报错，可以直接进入url网址进行手动下载，然后把数据放入root目录下的raw文件夹里。因为笔记本配置陈旧，使用GPU的时候，torch会报警告：E:\Anaconda\envs\tensorflow_gpu\lib\site-packages\torch\cuda\__init__.py:132: UserWarning: Found GPU0 Quadro K2100M which is of cuda capability 3.0. PyTorch no longer supports this GPU because it is too old. The minimum cuda capability that we support is 3.5.。笔记本上就是用来学习的。训练网络还是在服务器上（笑哭）。笔记本上还是用CPU做训练，不然会报CUDA和GPU算力不匹配的错RuntimeError: CUDA error: no kernel image is available for execution on the device。PyG的入门操作可以按照PyG官方简易教程学习。因为官网的PyG版本是1.4.0，而电脑上的版本是1.2.0，所以有些接口会有一点小问题，甚至变量名有出入，这些都容易解决，不一一叙述。

2. 搭建并训练一个简单图神经网络

根据官网教程，搭建一个简单的用于多分类的图卷积网络。输入是一个无向图数据（Graph data），图数据由节点数据（Node）和边数据（Edge）构成。节点数据是一个$N\times C$的矩阵，其中$N$表示节点个数，而$C$表示节点特征长度。边数据是一个$2 \times E$的矩阵。$E$表示有向边个个数。在PyG中，由于边数据被定义为有向边，如果表示无向边图的话，那么$E/2$才是无向边的个数。代码如下所示。node_features代表$C$，而node_class代表预测类别，例子中是7类，即node_class=7。

class simpleNet(torch.nn.Module): # simpleNet继承自torch.nn.Module
    def __init__(self, node_features, node_classes): 
        super(simpleNet, self).__init__() # simpleNet继承自torch.nn.Module，做父类初始化
        self.conv1 = GCNConv(node_features, 16) # 构建一个通道数为16的标准图卷积
        self.conv2 = GCNConv(16, node_classes) # 构建一个通道数为node_classes的标准图卷积，用作预测
        
    # 卷积+激活函数和dropout是深度学习常见操作，对图卷积神经网络来说也不例外
    # 分类问题最后使用softmax做归一化
    def forward(self, data):
        x, edge_index = data.x, data.edge_index # 分解节点和边数据，x是N*C的张量
        x = self.conv1(x, edge_index) # 它是N*16张量
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index) # 它是N*7的张量
        return F.log_softmax(x,dim=1) # 输出N*7的张量，在dim=1的维度做softmax

简易的训练和测试过程如下所示：

def example_5():
    # load model and dataset
    dataset = Planetoid(root='Cora', name='Cora')
    data = dataset[0]
    model = simpleNet(dataset.num_features, dataset.num_classes)
    device = torch.device('cuda')

    # report
    print("dataset.num_features:", dataset.num_features)
    print("dataset.num_classes:", dataset.num_classes)

    # transfer them to GPU 笔记本陈旧，只能使用CPU训练（笑哭）
    # model = model.to(device)
    # data = data.to(device)

    # define optimizer
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

    # begin training with evaluation
    # train_mask和test_mask表示训练数据和测试数据的位置
    model.train()
    num_epoch = 5
    for epoch in range(num_epoch):
        # train
        optimizer.zero_grad()
        out = model(data)
        loss = F.nll_loss(out[data.train_mask.bool()], data.y[data.train_mask.bool()])
        loss.backward()
        optimizer.step()

    # evaluation
    model.eval（)
    temp = model(data) # 模型输出N*7的张量
    _, pred = model(data).max(dim=1) # pred是N*1的张量，max(dim=1)表示选取7个类别概率最大的位置的索引
    correct = float(pred[data.test_mask.bool()].eq(data.y[data.test_mask.bool()]).sum().item())
    acc = correct/data.test_mask.sum().item()
    print('Accuracy: {:.4f}'.format(acc))
    print("finished...")
    print('data x type:', data.x.size())
    print('data edge_index type:', data.edge_index.size())
    print('temp type:', temp.size())
    print('pred type:', pred.size())

if __name__ == "__main__":
    example_5()

跟官网稍不一样的是，我在mask变量后面加上强制转换为bool类型的操作，即.bool()。不然运行时候会报警告UserWarning: indexing with dtype torch.uint8 is now deprecated, please use a dtype torch.bool instead.。输出结果是：

dataset.num_features: 1433
dataset.num_classes: 7
Accuracy: 0.6450
finished...
data x type: torch.Size([2708, 1433])
data edge_index type: torch.Size([2, 10556])
temp type: torch.Size([2708, 7])
pred type: torch.Size([2708])

在这个例子中，图的关系实现已经构建好了。对于点云数据，还需要提取构建KNN近邻图。在后续会做讨论。

3. PyTorch Geometric机制浅析

3.1 图卷积层

图卷积层定义我参考了PyG论文【1】。在PyG中，设计者们构建一个图$G=(X, (I, \mathbf{E}))$。$X \in \mathbf{R}^{N \times F}$表示图中的节点。$N$表示节点的个数。$F$表示节点特征的长度。$(I, \mathbf{E})$表示节点之间的连接关系。$I \in \mathbf{R}^{2 \times E}$，其中$E$表示有向边的个数。$\mathbf{E} \in \mathbf{R}^{E \times D}$表示有向边的特征。$D$表示边特征的长度。并不是所有的图都有$\mathbf{E}$（在第二节的代码应用中，就没有它）。所以边特征是一个可选项。在我的讨论中，请注意两个符号$E$和$\mathbf{E}$的含义哈。

PyG设计了一个通用的图卷积层。符号$x_i^L$表示第$i$个节点在第$L$层的特征。$x_i^{L+1}$则表示它在第$L+1$层的特征。集合$X^L$表示第$L$层所有节点特征的集合。通用的图卷积层输入是图$G^L$，输出$X^{L+1}$。$x_i^{L+1}$的计算方式如下所示：

$x_i^{L+1}=\gamma_{\theta}^L(x_i^L, \Box_{j\in N^L(i)} \phi_{\theta}^L(x_i^L,x_j^L, e_{ij}^L))$

其中$N^L(i)$表示第$L$层的KNN近邻图。$\Box$是对称聚合操作符，例如求和，最大，乘积，平均等。$e_{ij}^L\in \mathbf{E}^L$表示第$L$层边的特征。函数$\gamma^L$和$\phi^L$都是带超参数的可微函数，比如MLP。$\phi^L$是消息传递函数（跟概率图中的消息传递机制很相似），而$\gamma^L$是信息更新函数。

考虑到一些图的近邻图以及边特征不会因为图卷积层层数变化而变化，那么$x_i^{L+1}$的计算方式可以简单地表示为：

$x_i^{L+1}=\gamma_{\theta}^L(x_i^L, \Box_{j\in N(i)} \phi_{\theta}^L(x_i^L,x_j^L, e_{ij}))$

3.2 图卷积层基类

在PyG中，消息传递层（message passing layers）是实现任何图卷积类的基类，计算方式如3.1节所示。官方教程讲解的比较细。此外这篇博客比较清晰地讲解了消息传递层。用下行语句可以创建一个默认消息传递层。aggr="add"表示聚合操作是求和。flow表示消息的朝向是近邻点指向目标点。做点云处理，建立的图是无向的，所以flow使用默认值即可。基类的代码比较晦涩。在这一节不做分析。MessagePassing类是任何图卷积的父类。而MessagePassing类的父类是torch.nn.Module类。

net = torch_geometric.nn.MessagePassing(aggr="add", flow="source_to_target")

3.3 搭建边缘卷积（1）

边缘卷积（Edge Convolution）是处理点云一种高效卷积方式之一。边缘卷积的搭建参考官方教程。边缘卷积是图卷积的一种特例，$x_i^{L+1}$的计算方式如下所示：

$x_i^{L+1}=\max_{j\in N(i)} h_{\theta}(x_i^L, x_j^L-x_i^L)$

其中$h_{\theta}$表示MLP函数。聚合函数是最大化。搭建的代码如下所示：

import torch
from torch.nn import Sequential as Seq, Linear, ReLU
from torch_geometric.nn import MessagePassing

class EdgeConv(MessagePassing):
    def __init__(self, in_channels, out_channels):
        # EdgeConv类继承自基类MessagePassing
        # 初始化EdgeConv需要先初始化MessagePassing，使得基类的聚合操作是Max
        super(EdgeConv, self).__init__(aggr='max') #  "Max" aggregation.
        # 定义h_theta函数，两层感知机和激活函数
        self.mlp = Seq(Linear(2 * in_channels, out_channels),
                       ReLU(),
                       Linear(out_channels, out_channels))

	# 前向计算调用基类MessagePassing的传递函数propagate
    def forward(self, x, edge_index):
        # x has shape [N, in_channels]
        # edge_index has shape [2, E]
        
        return self.propagate(edge_index, size=(x.size(0), x.size(0)), x=x)

	# 获取消息，即h_theta函数的输入
    def message(self, x_i, x_j):
        # x_i has shape [E, in_channels]
        # x_j has shape [E, in_channels]

        tmp = torch.cat([x_i, x_j - x_i], dim=1)  # tmp has shape [E, 2 * in_channels]
        return self.mlp(tmp)
	
	# 得到下一层目标节点的特征
    def update(self, aggr_out):
        # aggr_out has shape [N, out_channels]

        return aggr_out

从上述代码中不难发现，EdgeConv类继承自父类MessagePassing，EdgeConv类主要做三件事：（1）定义聚合运算（在初始化中体现）；（2）定义$h_{\theta}$函数（在初始化中体现）；（3）获取$h_{\theta}$的输入（在函数message中体现）。然后函数update和message是对父类MessagePassing中函数的重载（Override）。函数propagate将调用这些重载函数，完成对$x_i^{L+1}$计算。请注意函数propagate的三个变量，将在3.5节做分析。

3.4 搭建边缘卷积（2）

在EdgeConv的论文中，KNN近邻图会根据每一层的特征向量重新计算。即$N^L(i)$每一层都不一样。考虑这种情况，继续搭建一个新类DynamicEdgeConv。$x_i^{L+1}$的计算方式如下所示。

$x_i^{L+1}=\max_{j\in N^L(i)} h_{\theta}(x_i^L, x_j^L-x_i^L)$

类DynamicEdgeConv的搭建代码如下所示：

from torch_geometric.nn import knn_graph

# DynamicEdgeConv继承自EdgeConv，
# 初始化的时候定义父类EdgeConv的in_channels和 out_channels
class DynamicEdgeConv(EdgeConv):
    def __init__(self, in_channels, out_channels, k=6):
        super(DynamicEdgeConv, self).__init__(in_channels, out_channels)
        self.k = k

    def forward(self, x, batch=None):
    	# 每次都会重新计算KNN近邻图，调用torch_cluster库函数
        edge_index = knn_graph(x, self.k, batch, loop=False, flow=self.flow)
        # 调用父类的forward函数
        return super(DynamicEdgeConv, self).forward(x, edge_index)

3.5 消息传递层代码简析

最后看一下基类MessagePassing的底层代码。从之前节的讨论中，基类函数大概有update和message和propagate三种。其中update和message比较简单，貌似专门用来做重载的，有点像C++中的虚函数（Virtual Function）。在3.3节中，我讨论了它们的具体的重载实现。因为EdgeConv中没有消息更新函数$\gamma(\cdot)$，所以update的重载为空。

# 读读注释吧
    def message(self, x_j):  # pragma: no cover
        r"""Constructs messages in analogy to :math:`\phi_{\mathbf{\Theta}}`
        for each edge in :math:`(i,j) \in \mathcal{E}`.
        Can take any argument which was initially passed to :meth:`propagate`.
        In addition, features can be lifted to the source node :math:`i` and
        target node :math:`j` by appending :obj:`_i` or :obj:`_j` to the
        variable name, *.e.g.* :obj:`x_i` and :obj:`x_j`."""

        return x_j

    def update(self, aggr_out):  # pragma: no cover
        r"""Updates node embeddings in analogy to
        :math:`\gamma_{\mathbf{\Theta}}` for each node
        :math:`i \in \mathcal{V}`.
        Takes in the output of aggregation as first argument and any argument
        which was initially passed to :meth:`propagate`."""

        return aggr_out

函数propagate是实现核心。这个有点复杂，我还读不懂。把计算的核心代码放在下面吧：

# 对应3.1节图卷积公式
		# 计算phi()
        out = self.message(*message_args)
        # 做聚合运算
        out = scatter_(self.aggr, out, edge_index[i], dim_size=size[i])
        # 计算gamma()
        out = self.update(out, *update_args)
        return out

4. 结语

图卷积有统一的形式，有着广泛的应用。PyTorch Geometric给出它在工程上的实现。在讨论PyTorch Geometric理论细节和实现细节后，这是一个不错的有趣的project。propagate设计的很有意思，有空我会分析一下。

参考文献

1. Matthias Fey & Jan E. Lenssen, FAST GRAPH REPRESENTATION LEARNING WITH PYTORCH GEOMETRIC, Published as a workshop paper at ICLR 2019