PyG学习笔记1-INTRODUCTION BY EXAMPLE

官方示例

torch_geometric.data.Data

节点和节点之间的边构成了图。所以在 PyG 中，如果要构建图，那么需要两个要素：节点和边。PyG 提供了torch_geometric.data.Data 用于构建图，包括 5 个属性，每一个属性都不是必须的，可以为空。

• x: 用于存储每个节点的特征，形状是[num_nodes, num_node_features]。
• edge_index: 用于存储节点之间的边，形状是 [2, num_edges]。
• pos: 存储节点的坐标，形状是[num_nodes, num_dimensions]。
• y: 存储样本标签。如果是每个节点都有标签，那么形状是[num_nodes, *]；如果是整张图只有一个标签，那么形状是[1, *]。
• edge_attr: 存储边的特征。形状是[num_edges, num_edge_features]。

实际上，Data对象不仅仅限制于这些属性，我们可以通过data.face来扩展Data，以张量保存三维网格中三角形的连接性。

需要注意的的是，在Data里包含了样本的 label，这意味和 PyTorch 稍有不同。在PyTorch中，我们重写Dataset的__getitem__()，根据 index 返回对应的样本和 label。在 PyG 中，我们使用的不是这种写法，而是在get()函数中根据 index 返回torch_geometric.data.Data类型的数据，在Data里包含了数据和 label。

下面一个例子是未加权无向图 ( 未加权指边上没有权值 )，包括 3 个节点和 4 条边。

import torch
from torch_geometric.data import Data
# 由于是无向图，因此有 4 条边：(0 -> 1), (1 -> 0), (1 -> 2), (2 -> 1)
edge_index = torch.tensor([[0, 1, 1, 2],
                           [1, 0, 2, 1]], dtype=torch.long)
# 节点的特征                           
x = torch.tensor([[-1], [0], [1]], dtype=torch.float)

data = Data(x=x, edge_index=edge_index)

Dataset 与 DataLoader

PyG 的 Dataset继承自torch.utils.data.Dataset，自带了很多图数据集，我们以TUDataset为例，通过以下代码就可以加载数据集，root参数设置数据下载的位置。通过索引可以访问每一个数据。

from torch_geometric.datasets import TUDataset
dataset = TUDataset(root='/tmp/ENZYMES', name='ENZYMES')
data = dataset[0]

在一个图中，由edge_index和edge_attr可以决定所有节点的邻接矩阵。PyG 通过创建稀疏的对角邻接矩阵，并在节点维度中连接特征矩阵和 label 矩阵，实现了在 mini-batch 的并行化。PyG 允许在一个 mini-batch 中的每个Data (图) 使用不同数量的节点和边。

让我们再试一个！让我们下载 Cora，半监督图节点分类的标准基准数据集：（需要开VPN下载数据）

from torch_geometric.datasets import Planetoid

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

len(dataset)
>>> 1

dataset.num_classes
>>> 7

dataset.num_node_features
>>> 1433

这一次，Data对象包含每个节点的标签，以及其他节点级属性：、和 ，其中train_mask``val_mask``test_mask

• train_mask表示针对哪些节点进行训练（140 个节点），
• val_mask表示要用于验证的节点，例如，执行早期停止（500 个节点），
• test_mask表示要针对哪些节点进行测试（1000 个节点）。

Mini-Batches

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

from torch_geometric.datasets import TUDataset
from torch_geometric.loader import DataLoader

dataset = TUDataset(root='/tmp/ENZYMES', name='ENZYMES', use_node_attr=True)
loader = DataLoader(dataset, batch_size=32, shuffle=True)

for batch in loader:
    batch
    >>> DataBatch(batch=[1082], edge_index=[2, 4066], x=[1082, 21], y=[32])

    batch.num_graphs
    >>> 32

torch_geometric.data.Batch 继承自 torch_geometric.data.Data 并包含一个名为 batch 的附加属性。batch是一个列向量，它将每个节点映射到批处理中的相应图中。

Data Transforms

变换是变换图像和执行增强的常用方法。PyG 附带了自己的转换，这些转换期望Data对象作为输入，并返回新的转换数据对象。转换可以使用torch_geometric.transforms.Compose链接在一起，并在将已处理的数据集保存到磁盘 （） 上之前或在访问数据集 （） 中的图形之前应用。torchvision``pre_transform``transform

我们可以通过转换从点云生成最近邻图，从而将点云数据集转换为图形数据集：

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

dataset = ShapeNet(root='/tmp/ShapeNet', categories=['Airplane'],
                    pre_transform=T.KNNGraph(k=6))

dataset[0]

Learning Methods on Graphs

我们将使用一个简单的GCN图层，并在Cora引文数据集上复制实验。

我们首先需要加载Cora数据集：

from torch_geometric.datasets import Planetoid

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

现在，让我们实现一个双层 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)

构造函数定义了两个GCNConv层.

请注意，非线性没有集成到调用中，因此需要在之后应用（这在PyG中的所有运算符中是一致的）。在这里，我们选择使用ReLU作为我们的中间非线性，并最终输出类数的softmax分布。让我们在训练节点上训练此模型 200 个 epoch：

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}')

自定义 Dataset

尽管 PyG 已经包含许多有用的数据集，我们也可以通过继承torch_geometric.data.Dataset使用自己的数据集。提供 2 种不同的Dataset：

• InMemoryDataset：使用这个Dataset会一次性把数据全部加载到内存中。
• Dataset: 使用这个Dataset每次加载一个数据到内存中，比较常用。

我们需要在自定义的Dataset的初始化方法中传入数据存放的路径，然后 PyG 会在这个路径下再划分 2 个文件夹：

• raw_dir: 存放原始数据的路径，一般是 csv、mat 等格式
• processed_dir: 存放处理后的数据，一般是 pt 格式 ( 由我们重写process()方法实现)。

Transforms

transforms在计算机视觉领域是一种很常见的数据增强。PyG 有自己的transforms，输出是Data类型，输出也是Data类型。可以使用torch_geometric.transforms.Compose封装一系列的transforms。我们以 ShapeNet 数据集 (包含 17000 个 point clouds，每个 point 分类为 16 个类别的其中一个) 为例，我们可以使用transforms从 point clouds 生成最近邻图：

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

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])

还可以通过transform在一定范围内随机平移每个点，增加坐标上的扰动，做数据增强：

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

dataset = ShapeNet(root='/tmp/ShapeNet', categories=['Airplane'],
                    pre_transform=T.KNNGraph(k=6),
                    transform=T.RandomTranslate(0.01))
# dataset[0]: Data(edge_index=[2, 15108], pos=[2518, 3], y=[2518])

模型训练

这里只是展示一个简单的 GCN 模型构造和训练过程，没有用到Dataset和DataLoader。

我们将使用一个简单的 GCN 层，并在 Cora 数据集上实验。有关 GCN 的更多内容，请查看**这篇博客**。

我们首先加载数据集：

from torch_geometric.datasets import Planetoid

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

然后定义 2 层的 GCN：

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

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)

然后训练 200 个 epochs：

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))

参考链接

PyG Documentation — pytorch_geometric 2.0.2 documentation (pytorch-geometric.readthedocs.io)