Pytorch图神经网络库——PyG基础操作

图数据处理
常用的基准数据库
- Mini-Batches
Data Transforms
图学习方法
参考资料

PyG（PyTorch Geometric）是一个基于PyTorch构建的库，可轻松编写和训练图神经网络（GNN），用于与结构化数据相关的广泛应用。

它包括从各种已发表的论文中的图和其他不规则结构（也称为几何深度学习geometric deep learning）的各种方法。此外，它由易于使用的mini-batch加载器组成，用于在许多小型和单个巨型图上操作，多GPU支持，DataPipe支持，通过Quiver分布式图形学习，大量常见的基准数据集（基于简单的交互创建自己的界面），Graphgym实验管理器和有用的转换，可以在任意图以及3D网格或点云上学习。
PyG

图数据处理

单张图被PyG表示为torch_geometric.data.Data类型，有如下属性：
data.x: 节点的特征矩阵，形状为 [num_nodes, num_node_features]
data.edge_index: COO格式的图的边 shape [2, num_edges] and type torch.long
data.edge_attr：边的特征矩阵 shape [num_edges, num_edge_features]
data.y：训练数据的标签，节点级的目标 shape [num_nodes, *] or 图级的目标 shape [1, *]
data.pos: 节点的位置矩阵 shape [num_nodes, num_dimensions]

这些属性都不是必需的。事实上，Data 对象甚至不限于这些属性。例如，我们可以通过 data.face 扩展它，以保存形状为[3，num _ faces ]和torch.long 类型的张量中3D 网格中三角形的连通性。

PyTorch和torchvision将示例定义为图像和目标的元组。我们在PyG中省略了这个符号，以便以清晰易懂的方式实现各种数据结构。

我们展示了一个具有三个节点和四条边的未加权无向图的简单示例。每个节点仅包含一个特征：

import torch
from torch_geometric.data import Data

# 一个无权无向图的示例
# 边的索引 COO格式 就是说第一行代表行索引 第二行代表列索引
# 节点0和1之间有一条边；节点1和2之间有一条边 
edge_index = torch.tensor([[0, 1, 1, 2], 
                           [1, 0, 2, 1]], dtype=torch.long)
# 节点0 1 2的特征分别是
x = torch.tensor([[-1], [0], [1]], dtype=torch.float)

data = Data(x=x, edge_index=edge_index)
#>>> Data(x=[3, 1], edge_index=[2, 4])

图示意图
注意 edge _ index，即定义所有边的源节点和目标节点的张量，不是索引元组的列表。如果想这样写你的索引，应该在将它们传递给数据构造函数之前对它进行转置和调用:

# 不是COO格式的边,而是使用边的两个节点的元组形式
edge_index = torch.tensor([[0, 1],
                           [1, 0],
                           [1, 2],
                           [2, 1]], dtype=torch.long)
x = torch.tensor([[-1], [0], [1]], dtype=torch.float)

data = Data(x=x, edge_index=edge_index.t().contiguous())

虽然图只有两条边，但是我们需要定义四个索引元组来表示边的两个方向。
注意，edge _ index 中的元素必须只保存范围{0，... ，num _ node-1}内的索引。这是必需的，因为我们希望我们的最终数据表示尽可能紧凑，例如，我们希望通过 x [0]和 x [1]分别索引第一个边(0,1)的源和目标节点特征。

对于data提供了如下方法进行访问：

print(data.keys)
>>> ['x', 'edge_index']

print(data['x'])
>>> tensor([[-1.0],
            [0.0],
            [1.0]])

for key, item in data:
    print(f'{key} found in data')
>>> x found in data
>>> edge_index found in data

'edge_attr' in data
>>> False

data.num_nodes
>>> 3

data.num_edges
>>> 4

data.num_node_features
>>> 1

data.has_isolated_nodes()
>>> False

data.has_self_loops()
>>> False

data.is_directed()
>>> False

# Transfer data object to GPU.
device = torch.device('cuda')
data = data.to(device)

可以在torch_geometric.data.data中找到所有方法的完整列表。

常用的基准数据库

PyG包含大量常见的基准数据集，例如所有Planetoid数据集（Cora、Citseeer、Pubmed）、来自http://graphkernels.cs.tu-dortmund.de的所有图分类数据集以及它们的简化版本，QM7和QM9数据集，以及少量3D网格/点云数据集，如FAUST、ModelNet10/40和ShapeNet。

初始化数据集很简单。数据集的初始化将自动下载其原始文件并将其处理为前面描述的 Data 格式。例如，要加载 ENZYMES 数据集(包括6个类中的600个图形) ，输入:

from torch_geometric.datasets import TUDataset
dataset = TUDataset(root='../DemoData/ENZYMES', name='ENZYMES')

len(dataset) # 600
dataset.num_classes # 6
dataset.num_node_features # 3
# 查看第一张图
data = dataset[0]
data # Data(edge_index=[2, 168], x=[37, 3], y=[1])
data.is_undirected() # True

我们可以看到数据集中的第一个图包含37个节点，每个节点有3个特性。有168/2 = 84条无向边，图被分配到一个类中。此外，数据对象只持有一个图级目标。
我们甚至可以使用切片、long或bool张量来分割数据集。例如，要创建90/10的训练/测试分割，输入:

train_dataset = dataset[:540]
>>> ENZYMES(540)

test_dataset = dataset[540:]
>>> ENZYMES(60)

# 同样 也可以打乱数据
dataset = dataset.shuffle()
>>> ENZYMES(600)

# 等于这个操作
perm = torch.randperm(len(dataset))
dataset = dataset[perm]
>> ENZYMES(600)

让我们再试一个! 我们下载Cora，半监督图节点分类的标准基准数据集，这个数据集是图神经网络论文中经常看见的数据集，是一个论文的引用图，节点的特征为论文的词向量:

from torch_geometric.datasets import Planetoid

dataset = Planetoid(root='/tmp/Cora', name='Cora')
>>> Cora()

len(dataset)
>>> 1 # 数据集只有一张图

dataset.num_classes
>>> 7 # 分为7类

dataset.num_node_features
>>> 1433 # 每个节点有1433维特征

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

注意：可能会出现：TimeoutError: [WinError 10060] 由于连接方在一段时间后没有正确答复或连接的主机没有反应，连接尝试失败。
解决方法可以更改源码中的下载地址。

Data对象保存每个节点的标签，以及额外的节点级属性:train_mask、val_mask和test_mask，其中

Train_mask表示要训练哪个节点(140个节点)，
Val_mask表示使用哪些节点进行验证，例如，执行早期停止(500个节点)，
Test_mask表示要测试哪个节点(1000个节点)。

比如data.train_mask是一个一维的bool类型的tensor，里面的数据是True就代表是训练数据。

data.train_mask.sum().item()
>>> 140

data.val_mask.sum().item()
>>> 500

data.test_mask.sum().item()
>>> 1000

Mini-Batches

神经网络通常以批处理方式进行训练。PyG通过创建稀疏块对角邻接矩阵（由edge_index定义）并在节点维度中连接特征矩阵和目标矩阵，在小批量上实现并行化。这种组合允许在一个批处理中不同数量的节点和边：

PyG包含它自己的torch_geometric.loader.DataLoader，它已经处理了这个连接过程。其实使用的过程和torch里面的DataLoader是一样的：
mini-batches
torch_geometric.data.Batch从torch_geometric.data.Data继承，并包含称为batch的附加属性。batch是一个列向量，保存了batch中每个节点和对应图的映射关系。可以使用batch来计算batch中每个图中各个节点的平均特征。

from torch_scatter import scatter_mean

for data in loader:
    print(data) # DataBatch(batch=[1082], edge_index=[2, 4066], x=[1082, 21], y=[32])
    print(data.num_graphs) # 32
    x = scatter_mean(data.x, data.batch, dim=0)
    print(x.size()) # torch.Size([32, 21])

Data Transforms

下面以一个例子来解释，ShapeNet数据集是包含17000个3D点云的数据集。

from torch_geometric.datasets import ShapeNet

dataset = ShapeNet(root='/tmp/ShapeNet', categories=['Airplane'])

dataset[0]
>>> Data(pos=[2518, 3], y=[2518])

我们可以通过transforms将点云生成最近邻图，将点云数据集转换为图数据集:

import torch_geometric.transforms as T
from torch_geometric.datasets import ShapeNet
# 主要就是加了最后一个参数pre_transform
dataset = ShapeNet(root='/tmp/ShapeNet', categories=['Airplane'],pre_transform=T.KNNGraph(k=6))

dataset[0]
>>> Data(edge_index=[2, 15108], pos=[2518, 3], y=[2518])

pre_transform参数可以在数据存入磁盘之前进行转换，在下次使用时，数据集将自动包含边。

此外，我们可以使用transform参数来随机增强一个Data对象，例如，将每个节点的位置转换为一个小数字：

import torch_geometric.transforms as T
from torch_geometric.datasets import ShapeNet

dataset = ShapeNet(root='../DemoData/ShapeNet', categories=['Airplane'], pre_transform=T.KNNGraph(k=6), transform=T.RandomTranslate(0.01))
dataset[0] # Data(x=[2518, 3], y=[2518], pos=[2518, 3], category=[1])

图学习方法

在学习了PyG中的数据处理、datasets, loader 以及 transforms之后，是时候实现我们的第一个图形神经网络了!

我们将使用一个简单的 GCN 层，并在 Cora 引文数据集上复制实验。要了解 GCN 的高层解释，可以看看它的博客文章。

首先，加载数据集：

from torch_geometric.datasets import Planetoid

dataset = Planetoid(root='/data/Cora', name='Cora')

然后，定义一个两层的GCN，和torch里面的定义方式一样，只是在传入GCN网络时传的是两个参数，第一个是数据特征，第二个是边的index。数据特征x就是每个节点的特征向量，edge_index就是（2，节点数）的Tensor，表示哪两个节点之间有边。卷积层后面的ReLU激活函数以及dropout层也是常用的，也不做介绍。

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)

然后对网络进行训练，训练的方式也是和torch中形式一样的：

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}')
>>> Accuracy: 0.8150

最后做个总结：使用PyG创建图神经网络模型的过程，和基于Pytorch的过程差不多，其中最重要的是DataLoader的写法。
需要注意的是，如果使用 torch_geometric.datasets中没有的数据集的话，将数据转化为图的形式可能需要一些时间。

参考资料

[1] https://pytorch-geometric.readthedocs.io/en/latest/notes/introduction.html
[2] https://github.com/pyg-team/pytorch_geometric
[3] https://blog.csdn.net/qq_40344307/article/details/122160733
[4] GRAPH CONVOLUTIONAL NETWORKS