使用 PyG 进行图神经网络训练

前言

最近一直在想创新点，搭模型，想尝试一下图神经网络，想着自己实现一个，但是之前也没有尝试过写 GNN 模型，对其中的实现细节也没有实际尝试过，最后找到了 PyG ，尝试一下之后发现还是挺简单的，也比较好拿到现有模型里面，于是开始挖坑。

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

目前网上对 PyG 的相关文档并不多，大本部也都是比较重复的内容，因此我主要参考的还是官方文档。具体安装方式参考 PyG Installation。

图结构

建图

首先，我们需要根据数据集进行建图，在 PyG 中，一个 Graph 的通过torch_geometric.data.Data进行实例化，它包括下面两个最主要的属性：

• data.x: 节点的特征矩阵，形状为[num_nodes, num_node_features];
• data.edge_index: 图的边索引，用 COO 稀疏矩阵格式保存。形状为[2, num_edges];

稀疏矩阵是数值计算中普遍存在的一类矩阵，主要特点是绝大部分的矩阵元为零。COO (Coordinate) 格式是将矩阵中的非零元素存储，每一个元素用一个三元组来表示，分别是（行号，列号，数值）

上面两个属性也就是整个图中最重要的部分。同时你也可以定义下面的一些其他信息：

• data.edge_attr: 边的特征矩阵，形状为[num_edges, num_edge_features];
• data.y: 计算损失所需的目标数据，比如节点级别的形状为[num_nodes, *]，或者图级别的形状为[1, *];

现在，我们来简单创建一张图：

import torch
from torch_geometric.data import Data

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)
>>> Data(edge_index=[2, 4], x=[3, 1])

需要注意的是：

1. 第一行 edge_index[0] 表示起点，第二行 edge_index[1] 表示终点；
2. 虽然只有两条边，但在 PyG 中处理无向图时实际上是互为头尾节点；
3. 矩阵中的值表示索引，指代 x 中的节点。

实际上，你不需要找对专门的属性去存取某些信息，你完全可以「自定义属性名」，并放入你想要的内容，事实上就像一个字典。例如，你想要为图再添加一个 freq 属性以表示访问频率：

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)
data.freq = torch.tensor([2, 5, 1], dtype=torch.long)
>>> Data(x=[3, 1], edge_index=[2, 4], freq=[3])

此外，这里再列举一些个人觉得比较常用的操作，感觉官方文档有些地方不清不楚，完整的类说明可以看这里：

• coalesce(): 对 edge_index 中的边排序并去重
• clone(): 创建副本
• to(cuda:0): 把数据放到 GPU
• num_edges: 返回边数量
• num_nodes: 返回节点数量

关于 train_mask

在官方的数据集中，划分「训练集」和「测试集」的方式是创建一张大图，然后指定训练节点以及测试节点，通过 train_mask 和 test_mask 来实现。

from torch_geometric.datasets import Planetoid

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

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

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

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

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

上面是官方提供的数据集中的一个例子，但个人觉得并不是所有的数据集都适合这种方法，我仍习惯于训练集和测试集分开，也就是创建两张图，因为大多数情况下都需要自定义一个 MyDataset，你可以把建图的操作放在里面，所以其实感觉这样还更方便。

data_train, data_test = split_train_test(data)

# Create dataset
dataset_train = MyDataset(data_train)
dataset_test = MyDataset(data_test)

当然，这还是要看具体做什么任务的，对某些任务来说可能只需要一张大图，对我来说可能就需要分成两张。

关于 Embedding

最开始的时候我有想过，为什么要在一开始创建 x 的时候就让我把节点的维度给定下来，这不应该是我后面模型里面 Embedding 的时候再做的事情吗，难不成建图的时候就要 Embedding？当然，对于有些情况，建图的时候完全就可以直接 Embedding。

事实上，其实就是我想多了，其实把它当做一个字典就好。

import torch
from torch import nn
from torch_geometric.data import Data

# create graph
edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]], dtype=torch.long)
x = torch.tensor([1, 0, 2], dtype=torch.long)

data = Data(x=x, edge_index=edge_index)
data.freq = torch.tensor([2, 5, 1], dtype=torch.long)

...

# model.py
# init embedding layer
emb_layer = nn.Embedding(3, 8)

data.x = emb_layer(data.x)
# data.x_emb = emb_layer(data.x)

print(data)
>>> Data(x=[3, 8], edge_index=[2, 4], freq=[3])

就像上面展示的这样，我的做法是先把节点的 ID 填进去，接着在模型里面进行 Embedding，当然你可以直接使用 data.x = emb_layer(data.x) 把原来的 ID 给替换掉；也可能你需要保留 ID，那么就可以把它放到一个新的属性中，比如 data.x_emb。总的来说其实就是不需要把 data 的属性看得那么死，完全可以按照你的需要来。

关于 Batch 和 DataLoader

Batch

在 PyG 中，你可以通过手动的方式进行批量化操作：

import torch
from torch import nn
from torch_geometric.data import Data, Batch

edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]], dtype=torch.long)
x = torch.tensor([1, 0, 2], dtype=torch.long)
g1 = Data(x=x, edge_index=edge_index)
g2 = Data(x=x, edge_index=edge_index)
g3 = Data(x=x, edge_index=edge_index)
data = [g1, g2, g3]
>>> Data(x=[3], edge_index=[2, 4])

batch = Batch().from_data_list(data)
>>> DataDataBatch(x=[9], edge_index=[2, 12], batch=[9], ptr=[4])

在上面的例子中，为了方便我创建了 3 张相同的图。现在我们来看看 batch 具体做了些什么：

print(batch.x)
>>> tensor([1, 0, 2, 1, 0, 2, 1, 0, 2])

print(batch.edge_index)
>>> tensor([[0, 1, 1, 2, 3, 4, 4, 5, 6, 7, 7, 8],
            [1, 0, 2, 1, 4, 3, 5, 4, 7, 6, 8, 7]])

print(batch.batch)
>>> tensor([0, 0, 0, 1, 1, 1, 2, 2, 2])

print(batch.ptr)
>>> tensor([0, 3, 6, 9])

实际上，Batch 的本质就是将多个图组合起来，创建一张大图。可以看出，首先在 x 中将所有节点拼接在一起，batch 指明了每个节点所属的批次，ptr 指明了每个 batch 的节点的起始索引号，然后在 edge_index 中对边进行拼接，同时为了避免冲突对索引加上了 ptr。

用公式来表示也就是：

$$\mathbf{A} = \begin{bmatrix} \mathbf{A}_1 & & \\ & \ddots & \\ & & \mathbf{A}_n \end{bmatrix}, \mathbf{X} = \begin{bmatrix} \mathbf{X}_1 \\ \vdots \\ \mathbf{X}_n \end{bmatrix}, \mathbf{Y} = \begin{bmatrix} \mathbf{Y}_1 \\ \vdots \\ \mathbf{Y}_n \end{bmatrix}$$

DataLoader

PyTorch 原生的 DataLoader 实际上对 Data 并不支持，虽然可以创建成功，但在遍历取数据的时候，你会发现如下错误：

default_collate: batch must contain tensors, numpy arrays, numbers, dicts or lists; found <class ‘torch_geometric.data.data.Data’>

PyG 有一个自己的 DataLoader，实际上只需要用它替换 PyTorch 原生的 DataLoader 就可以，个人觉得使用体验上和 PyTorch 差别不大。

from torch.utils.data import Dataset
# from torch.utils.data import DataLoader
from torch_geometric.loader import DataLoader

class MyDataset(Dataset):
    def __init__(self, data):
        super(MyDataset, self).__init__()
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        sub = self.data[idx]
        return sub

dataset = MyDataset(data)
dataloader = DataLoader(dataset)

图神经网络

讲完了图结构，以及数据集之后，现在正式进入到了模型训练阶段

Convolutional Layers

PyG 其实定义了非常多可供直接使用的 Convolutional Layers，具体你可以看这里。总的来说常用的结构基本都帮你实现好了：

• GCNConv
• GATConv
• SAGEConv
• LGConv
• …

这里以 GCNConv 为例，关于 GCN 的具体原理这里就不过多介绍了，有兴趣可以看我的另一篇文章：简单理解图神经网络 GNN。

还是以官方的例子开始：

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 并制定输入输出维度，比如 conv1。之后在 forward() 中只需要把我们一开始定义的图中的 x 和 edge_index 传进去就行了，也就是节点信息和边信息，当然，具体的输入你可以根据自己的实际情况，假如我的节点信息在 x_emb 这里，那我就把 x_emb 丢进去。其余的就和 torch 一样，我们做的只是说引入了一层。

自定义网络

当然，假如你需要的网络结构 PyG 中并没有预先给出，或者你想实现自己的聚合结构，那就需要重新定义自己的。

class GCNConv(MessagePassing):
    def __init__(self, in_channels: int, out_channels: int):
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.lin = Linear(in_channels, out_channels)

    def forward(self, x: Tensor, edge_index: Adj,
                edge_weight: OptTensor = None) -> Tensor:
        edge_index = gcn_norm(edge_index, edge_weight, x.size(self.node_dim))
        x = self.lin(x)

        # propagate_type: (x: Tensor, edge_weight: OptTensor)
        out = self.propagate(edge_index, x=x, edge_weight=edge_weight,
                             size=None)
        return out

    def message(self, x_j: Tensor, edge_weight: OptTensor) -> Tensor:
        return x_j if edge_weight is None else edge_weight.view(-1, 1) * x_j

    def message_and_aggregate(self, adj_t: SparseTensor, x: Tensor) -> Tensor:
        return matmul(adj_t, x)

上面我根据 PyG 的 GCNConv 拎出来的核心代码，主要在于传播 propagate() 以及消息传递 message()。PyG 提供了 MessagePassing 基类，它封装了「消息传递」的运行流程。关于「消息传递」的实现我自己目前并不算熟悉，这里就不展开了，有兴趣的可以去看其他的资料。

参考资料

• PyG Documentation