【深度学习】PyTorch框架(6)：GNN图神经网络理论和实践

【深度学习】PyTorch框架(6)：GNN图神经网络理论和实践

1.引言

在本文中，我们将探讨图神经网络（GNNs）在图上的应用。近年来，图神经网络在社交网络、知识图谱、推荐系统和生物信息学等多个领域中越来越受到关注。尽管GNNs背后的理论和数学可能初看之下颇为复杂，但其模型的实现却相对简单，有助于我们深入理解其方法论。因此，本文将重点介绍GNN的基本网络层的实现，包括图卷积和注意力层。最终，我们将演示如何在节点级别、边级别和图级别任务中应用GNN。
首先，我们从导入常用的库开始。我们将使用PyTorch Lightning，这在之前的文章中已经有所涉及。

## 导入标准库
import os
import json
import math
import numpy as np 
import time

# 导入绘图相关库
import matplotlib.pyplot as plt
%matplotlib inline 
from IPython.display import set_matplotlib_formats
set_matplotlib_formats('svg', 'pdf')  # 用于导出
from matplotlib.colors import to_rgb
import matplotlib
matplotlib.rcParams['lines.linewidth'] = 2.0
import seaborn as sns
sns.reset_orig()
sns.set()

# 导入进度条库
from tqdm.notebook import tqdm

# 导入PyTorch库
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data as data
import torch.optim as optim

# 导入Torchvision库
import torchvision
from torchvision.datasets import CIFAR10
from torchvision import transforms

# 导入PyTorch Lightning库
try:
    import pytorch_lightning as pl
except ModuleNotFoundError:  # 如果Google Colab未预装PyTorch Lightning，则在此安装
    !pip install --quiet pytorch-lightning>=1.4
    import pytorch_lightning as pl
from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint

# 设置数据集下载路径（例如CIFAR10）
DATASET_PATH = "../data"
# 设置预训练模型保存路径
CHECKPOINT_PATH = "../saved_models/tutorial7"

# 设置随机种子
pl.seed_everything(42)

# 确保在GPU上的所有操作都是确定性的（如果使用）
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

# 判断并设置设备
device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
print(device)

接下来，我们将下载一些预训练的模型。

## 下载预训练模型
import urllib.request
from urllib.error import HTTPError

# 预训练模型存储的GitHub URL
base_url = "https://raw.githubusercontent.com/phlippe/saved_models/main/tutorial7/"

# 需要下载的文件列表
pretrained_files = ["NodeLevelMLP.ckpt", "NodeLevelGNN.ckpt", "GraphLevelGraphConv.ckpt"]

# 如果检查点路径不存在，则创建它
os.makedirs(CHECKPOINT_PATH, exist_ok=True)

# 对于每个文件，检查它是否已经存在。如果没有，尝试下载它。
for file_name in pretrained_files:
    file_path = os.path.join(CHECKPOINT_PATH, file_name)
    if "/" in file_name:
        os.makedirs(file_path.rsplit("/",1)[0], exist_ok=True)
    if not os.path.isfile(file_path):
        file_url = base_url + file_name
        print(f"正在下载 {file_url}...")
        try:
            urllib.request.urlretrieve(file_url, file_path)
        except HTTPError as e:
            print("下载过程中出现问题。请尝试从其他渠道获取文件，或联系作者并提供完整的错误信息：\n", e)

2.图神经网络

2.1.图的表示方法

在深入探讨图上的特定神经网络操作之前，我们首先需要了解如何表示一个图。在数学上，图 $\mathcal{G}$ 可以定义为一个包含节点/顶点集合 $V$ 和边/链接集合 $E$ 的元组：$\mathcal{G}=(V,E)$。每条边由两个顶点组成，代表它们之间的连接。例如，考虑以下图：

图中的顶点为 V = { 1 , 2 , 3 , 4 } V=\{1,2,3,4\} V={1,2,3,4}，边为 $E=\{(1,2),(2,3),(2,4),(3,4)\}$。为了简化，我们假设图是无向的，因此不包括反向边如 $(2,1)$。在实际应用中，顶点和边通常具有特定的属性，边甚至可以是有向的。关键在于我们如何高效地以矩阵操作的方式表示这种多样性。通常，对于边的表示，我们可以选择邻接矩阵或顶点对索引列表。

邻接矩阵 $A$ 是一个方阵，其元素表示顶点对是否相邻，即是否相连。在最简单情况下，如果节点 $i$ 到 $j$ 有连接，则 $A_{ij}$ 为 1，否则为 0。如果图中有边的属性或不同类别的边，这些信息也可以添加到矩阵中。对于无向图，$A$ 是一个对称矩阵（$A_{ij}=A_{ji}$）。以上述图为例，其邻接矩阵如下：

$$A=\left[\begin{array}{cccc}0& 1& 0& 0\\ 1& 0& 1& 1\\ 0& 1& 0& 1\\ 0& 1& 1& 0\end{array}\right]$$
虽然以边的列表形式表达图在内存和计算上更为高效，但使用邻接矩阵更直观且易于实现。在下面的实现中，我们将使用邻接矩阵以简化代码。然而，常用的库可能会使用边列表，我们将在后续讨论中更详细地探讨这一点。

另外，我们也可以通过边的列表定义一个稀疏邻接矩阵，这样我们可以像处理密集矩阵一样进行操作，但更节省内存。PyTorch 通过其 torch.sparse 子包支持这一点（文档），但请注意，这仍处于测试阶段（API 可能在未来发生变化）。

2.2.图卷积

图卷积网络（GCNs）由 [Kipf 等人]在 2016 年在阿姆斯特丹大学提出。他还撰写了一篇关于该主题的优秀 [博客文章]，如果你希望从不同角度了解 GCNs，这篇文章是推荐的阅读。GCNs 与图像中的卷积类似，在于“滤波器”参数通常在图中的所有位置共享。同时，GCNs 依赖于消息传递方法，即顶点与邻居交换信息，并向彼此发送“消息”。在深入数学表达之前，我们先尝试直观理解 GCNs 的工作原理。首先，每个节点创建一个特征向量，代表它想要发送给所有邻居的消息。其次，这些消息被发送到邻居，使得每个节点从每个相邻节点接收到一条消息。以下是我们示例图的两个步骤的可视化：

如果我们想用更数学的语言来描述，首先需要决定如何合并一个节点接收到的所有消息。由于不同节点接收到的消息数量不同，我们需要一种适用于任何数量的操作。通常，我们选择求和或取平均。给定节点的先前特征 $H^{(l)}$，GCN 层定义如下：

$$H^{(l+1)}=\sigma \left({\hat{D}}^{-1/2}\hat{A}{\hat{D}}^{-1/2}{H}^{(l)}{W}^{(l)}\right)$$
$W^{(l)}$ 是我们将输入特征转换为消息的权重参数。我们在邻接矩阵 $A$ 中加入单位矩阵，以便每个节点也将自己的消息发送给自己：$\hat{A}=A+I$。最后，为了计算平均值而不是求和，我们计算对角矩阵 $\hat{D}$，其中 $D_{ii}$ 表示节点 $i$ 拥有的邻居数量。$\sigma$ 代表任意激活函数，不一定是 sigmoid（通常在 GNNs 中使用基于 ReLU 的激活函数）。

在 PyTorch 中实现 GCN 层时，我们可以利用张量上的灵活操作。我们不需要定义矩阵 $\hat{D}$，我们可以在之后简单地将总和消息除以邻居的数量。此外，我们用线性层替换权重矩阵，这也允许我们添加偏置。以下是 PyTorch 模块中 GCN 层的定义：

class GCNLayer(nn.Module):
def __init__(self, c_in, c_out):
super().__init__()
self.projection = nn.Linear(c_in, c_out)  # 定义线性层

def forward(self, node_feats, adj_matrix):
"""
前向传播函数
:param node_feats: 节点特征张量，形状为 [batch_size, num_nodes, c_in]
:param adj_matrix: 邻接矩阵的批次，形状为 [batch_size, num_nodes, num_nodes]
:return: 节点的新特征
"""
# 计算每个节点的邻居数量
num_neighbours = adj_matrix.sum(dim=-1, keepdims=True)
node_feats = self.projection(node_feats)  # 应用线性变换
node_feats = torch.bmm(adj_matrix, node_feats)  # 消息传递
node_feats = node_feats / num_neighbours  # 归一化
return node_feats

为了更好地理解 GCN 层，我们可以将其应用于我们之前的例子图。首先，我们定义一些节点特征和添加了自连接的邻接矩阵：

node_feats = torch.arange(8, dtype=torch.float32).view(1, 4, 2)
adj_matrix = torch.Tensor([[[1, 1, 0, 0],
[1, 1, 1, 1],
[0, 1, 1, 1],
[0, 1, 1, 1]]])

print("节点特征：\n", node_feats)
print("邻接矩阵：\n", adj_matrix)

接下来，我们对其应用一个 GCN 层。为了简化，我们将线性权重矩阵初始化为单位矩阵，这样输入特征等于消息。这使我们更容易验证消息传递操作：

layer = GCNLayer(c_in=2, c_out=2)
layer.projection.weight.data = torch.Tensor([[1., 0.], [0., 1.]])  # 初始化权重为单位矩阵
layer.projection.bias.data = torch.Tensor([0., 0.])  # 初始化偏置为零

with torch.no_grad():
out_feats = layer(node_feats, adj_matrix)

print("邻接矩阵：", adj_matrix)
print("输入特征：", node_feats)
print("输出特征：", out_feats)

如我们所见，第一个节点的输出值是它自己和第二个节点的平均值。类似地，我们可以验证所有其他节点。然而，在图神经网络中，我们也希望允许节点之间的特征交换超出其直接邻居。这可以通过应用多个 GCN 层来实现，从而构建出图神经网络的最终结构。图神经网络可以通过一系列 GCN 层和非线性激活函数（如 ReLU）来构建。下面的可视化图展示了这一点（图示来源 - [Thomas Kipf, 2016]）。

然而，从上述示例中我们可以看到一个问题，即节点 3 和 4 的输出特征是相同的，因为它们有相同的相邻节点（包括它们自己）。因此，如果我们只是简单地对所有消息求平均，GCN 层可能会使网络丢失节点特定的信息。已经提出了多种可能的改进方法。虽然最简单的选择可能是使用残差连接，但更常见的方法是要么增加自连接的权重，要么为自连接定义一个单独的权重矩阵。另外，我们可以重新考虑上一个教程中的概念：注意力。

2.3.图注意力网络

在上一个博文中，我们了解到注意力机制是一种基于输入查询和元素键动态计算权重的多个元素的加权平均。这种概念也可以应用于图结构数据，其中一种就是图注意力网络（GAT），由[Velickovic等人，2017]提出。与图卷积网络（GCN）类似，图注意力层使用线性层或权重矩阵为每个节点生成一个消息。在注意力部分，它将节点自身的消息作为查询，并将消息作为键和值进行平均（注意这也包括它自身的消息）。分数函数 实现为一个单层的多层感知机（MLP），它将查询和键映射到一个单一的值。MLP的结构如下所示（图示来源 - [Velickovic等人]：

$h_i$ 和 $h_j$ 分别是节点 $i $ 和 $j$ 的原始特征，它们代表了带有权重矩阵 $\mathbf{W} $ 的层的消息。$\mathbf{a}$ 是MLP的权重矩阵，其形状为 $[1,2\times d_{\text{message}}]$，而 ${\alpha }_{ij}$ 是从节点 $i$ 到 $j$ 的最终注意力权重。计算过程可以描述如下：

$${\alpha }_{ij}=\frac{\exp \left(\text{LeakyReLU}\left(\mathbf{a}\left[\mathbf{W}{h}_i||\mathbf{W}{h}_j\right]\right)\right)}{{\sum }_{k\in {\mathcal{N}}_i}\exp \left(\text{LeakyReLU}\left(\mathbf{a}\left[\mathbf{W}{h}_i||\mathbf{W}{h}_k\right]\right)\right)}$$
操作符 $||$ 表示连接，${\mathcal{N}}_i$ 表示节点 $i$ 邻居的索引。注意，与通常的做法不同，我们在对元素进行softmax之前应用了非线性（这里使用LeakyReLU）。虽然一开始看起来像是一个小变化，但这对注意力依赖于原始输入至关重要。具体来说，让我们暂时去掉非线性，并尝试简化表达式：

$$\begin{array}{rl}{\alpha }_{ij}& =\frac{\exp \left(\mathbf{a}\left[\mathbf{W}{h}_i||\mathbf{W}{h}_j\right]\right)}{{\sum }_{k\in {\mathcal{N}}_i}\exp \left(\mathbf{a}\left[\mathbf{W}{h}_i||\mathbf{W}{h}_k\right]\right)}\\ & =\frac{\exp \left({\mathbf{a}}_{:,:d/2}\mathbf{W}{h}_i+{\mathbf{a}}_{:,d/2:}\mathbf{W}{h}_j\right)}{{\sum }_{k\in {\mathcal{N}}_i}\exp \left({\mathbf{a}}_{:,:d/2}\mathbf{W}{h}_i+{\mathbf{a}}_{:,d/2:}\mathbf{W}{h}_k\right)}\\ & =\frac{\exp \left({\mathbf{a}}_{:,:d/2}\mathbf{W}{h}_i\right)\cdot \exp \left({\mathbf{a}}_{:,d/2:}\mathbf{W}{h}_j\right)}{{\sum }_{k\in {\mathcal{N}}_i}\exp \left({\mathbf{a}}_{:,:d/2}\mathbf{W}{h}_i\right)\cdot \exp \left({\mathbf{a}}_{:,d/2:}\mathbf{W}{h}_k\right)}\\ & =\frac{\exp \left({\mathbf{a}}_{:,d/2:}\mathbf{W}{h}_j\right)}{{\sum }_{k\in {\mathcal{N}}_i}\exp \left({\mathbf{a}}_{:,d/2:}\mathbf{W}{h}_k\right)}\\ & \end{array}$$
我们可以看到，如果没有非线性，与 $h_i$ 相关的注意力项实际上会自我抵消，导致注意力与节点本身无关。因此，我们将面临与GCN相同的问题，即为具有相同邻居的节点创建相同的输出特征。这就是为什么LeakyReLU至关重要，并在注意力中增加了对 $h_i$ 的依赖。

一旦我们获得了所有注意力因子，我们可以通过执行加权平均来计算每个节点的输出特征：

$$h_i^{\prime }=\sigma \left(\sum _{j\in {\mathcal{N}}_i}{\alpha }_{ij}\mathbf{W}{h}_j\right)$$
$\sigma $ 是另一个非线性激活函数，类似于GCN层中的激活函数。直观上，我们可以将注意力层中的完整消息传递表示如下（图示来源 - Velickovic等人）：

为了增加图注意力网络的表达能力，Velickovic等人 提出了将其扩展到多个头部，类似于Transformer中的多头注意力块。这导致 $N$ 个注意力层并行应用。在上面的图中，它被可视化为三种不同颜色的箭头（绿色、蓝色和紫色），然后将它们连接起来。平均值仅在网络中的最终预测层应用。

在详细讨论了图注意力层之后，我们可以在下面实现它：

class GATLayer(nn.Module):

def __init__(self, c_in, c_out, num_heads=1, concat_heads=True, alpha=0.2):
"""
构造函数：
参数：
c_in - 输入特征的维度
c_out - 输出特征的维度
num_heads - 头部的数量，即并行应用的注意力机制的数量。如果 concat_heads=True，则输出特征在头部之间平均分配。
concat_heads - 如果为True，则不同头部的输出将被连接而不是平均。
alpha - LeakyReLU激活的负斜率。
"""
super().__init__()
self.num_heads = num_heads
self.concat_heads = concat_heads
if self.concat_heads:
assert c_out % num_heads == 0, "输出特征的数量必须是头部数量的倍数。"
c_out = c_out // num_heads

# 层中需要的子模块和参数
self.projection = nn.Linear(c_in, c_out * num_heads)
self.a = nn.Parameter(torch.Tensor(num_heads, 2 * c_out)) # 每个头部一个
self.leakyrelu = nn.LeakyReLU(alpha)

# 原始实现的初始化
nn.init.xavier_uniform_(self.projection.weight.data, gain=1.414)
nn.init.xavier_uniform_(self.a.data, gain=1.414)

def forward(self, node_feats, adj_matrix, print_attn_probs=False):
"""
前向传播函数：
参数：
node_feats - 节点的输入特征。形状：[batch_size, c_in]
adj_matrix - 包括自连接的邻接矩阵。形状：[batch_size, num_nodes, num_nodes]
print_attn_probs - 如果为True，在前向传播过程中打印注意力权重（用于调试目的）
"""
batch_size, num_nodes = node_feats.size(0), node_feats.size(1)

# 应用线性层并按头部排序节点
node_feats = self.projection(node_feats)
node_feats = node_feats.view(batch_size, num_nodes, self.num_heads, -1)

# 我们需要为邻接矩阵中的每条边计算注意力逻辑值
# 在所有可能的节点组合上进行此操作非常昂贵
# => 创建一个 [W*h_i||W*h_j] 的张量，其中 i 和 j 是所有边的索引
edges = adj_matrix.nonzero(as_tuple=False) # 返回邻接矩阵非零位置的索引 => 边
node_feats_flat = node_feats.view(batch_size * num_nodes, self.num_heads, -1)
edge_indices_row = edges[:,0] * num_nodes + edges[:,1]
edge_indices_col = edges[:,0] * num
_nodes + edges[:,2]
a_input = torch.cat([
torch.index_select(input=node_feats_flat, index=edge_indices_row, dim=0),
torch.index_select(input=node_feats_flat, index=edge_indices_col, dim=0)
], dim=-1) # 索引选择返回一个张量，其中 node_feats_flat 在 dim=0 处被索引

# 计算注意力MLP输出（每个头部独立）
attn_logits = torch.einsum('bhc,hc->bh', a_input, self.a) 
attn_logits = self.leakyrelu(attn_logits)

# 将注意力值列表映射回矩阵
attn_matrix = attn_logits.new_zeros(adj_matrix.shape+(self.num_heads,)).fill_(-9e15)
attn_matrix[adj_matrix[...,None].repeat(1,1,1,self.num_heads) == 1] = attn_logits.reshape(-1)

# 加权平均的注意力
attn_probs = F.softmax(attn_matrix, dim=2)
if print_attn_probs:
print("Attention probs\n", attn_probs.permute(0, 3, 1, 2))
node_feats = torch.einsum('bijh,bjhc->bihc', attn_probs, node_feats)

# 如果头部应该被连接，我们可以通过重塑来做到这一点。否则，取平均值
if self.concat_heads:
node_feats = node_feats.reshape(batch_size, num_nodes, -1)
else:
node_feats = node_feats.mean(dim=2)

return node_feats 

再次，我们可以将图注意力层应用到我们上面的例子图上，以更好地理解动态。与之前一样，输入层被初始化为单位矩阵，但我们设置$\mathbf{a}$为任意数字的向量，以获得不同的注意力值。我们使用两个头部来展示在层中并行、独立工作的注意力机制。

layer = GATLayer(2, 2, num_heads=2)
layer.projection.weight.data = torch.Tensor([[1., 0.], [0., 1.]])
layer.projection.bias.data = torch.Tensor([0., 0.])
layer.a.data = torch.Tensor([[-0.2, 0.3], [0.1, -0.1]])

with torch.no_grad():
out_feats = layer(node_feats, adj_matrix, print_attn_probs=True)

print("邻接矩阵", adj_matrix)
print("输入特征", node_feats)
print("输出特征", out_feats)

我们建议你自己尝试至少为一个头部和一个节点计算注意力矩阵。条目在 $i$ 和 $j$之间不存在边的地方是0。对于其他条目，我们看到了一组多样化的注意力概率。此外，尽管节点3和4具有相同的邻居，但现在它们的输出特征是不同的。

3.PyTorch Geometric

在之前的讨论中，我们提到使用邻接矩阵实现图网络虽然简单直观，但对于大型图来说可能会带来巨大的计算开销。许多现实世界的图，如社交网络或生物信息网络，节点数可能超过20万，这种情况下使用邻接矩阵的实现方式就不再可行。幸运的是，在实现图神经网络（GNNs）时，我们可以利用一些经过优化的库来提高效率。目前，最流行的PyTorch库有[PyTorch Geometric](PyG Documentation)和Deep Graph Library（后者实际上不局限于任何特定的框架）。选择哪一个取决于你计划进行的项目和个人偏好。在本文中，我们将重点介绍作为PyTorch家族一部分的PyTorch Geometric。与PyTorch Lightning类似，pycharm上并不是默认安装的。因此，让我们在下面导入和/或安装它：

# torch geometric
try: 
    import torch_geometric
except ModuleNotFoundError:
    # 安装特定CUDA和PyTorch版本的torch geometric包。
    # 更多详情请参阅 Installation - pytorch_geometric documentation
    TORCH = torch.__version__.split('+')[0]
    CUDA = 'cu' + torch.version.cuda.replace('.','')

    !pip install torch-scatter     -f 403 Forbidden 
    !pip install torch-sparse      -f 403 Forbidden 
    !pip install torch-cluster     -f 403 Forbidden 
    !pip install torch-spline-conv -f 403 Forbidden 
    !pip install torch-geometric 
    import torch_geometric
import torch_geometric.nn as geom_nn
import torch_geometric.data as geom_data

3.1. PyTorch Geometric简介

PyTorch Geometric为我们提供了一套常见的图网络层，包括我们在上文中实现的GCN和GAT层。与PyTorch的torchvision类似，它还提供了常见的图数据集和相应的转换操作，以简化训练过程。与我们之前的实现相比，PyTorch Geometric使用一对索引列表来表示图的边。我们将在后续的实验中进一步探索这个库的细节。
在我们的任务中，我们希望能够从多种图网络层中进行选择。因此，我们在下面再次定义了一个字典，以便通过字符串来访问这些层：

gnn_layer_by_name = {
    "GCN": geom_nn.GCNConv,
    "GAT": geom_nn.GATConv,
    "GraphConv": geom_nn.GraphConv
}

3.2.图卷积层GraphConv

除了GCN和GAT，我们还增加了geom_nn.GraphConv层（[文档](torch_geometric.nn - pytorch_geometric documentation)）。GraphConv是一个具有自连接独立权重矩阵的GCN。从数学上讲，这可以表示为：
$${\mathbf{x}}_i^{(l+1)}={\mathbf{W}}_1^{(l+1)}{\mathbf{x}}_i^{(l)}+{\mathbf{W}}_2^{(\ell +1)}\sum _{j\in {\mathcal{N}}_i}{\mathbf{x}}_j^{(l)}$$
在这个公式中，邻居节点的消息是累加的，而不是平均的。然而，PyTorch Geometric提供了aggr参数，允许我们在求和、平均和最大池化之间进行切换。

4.图结构数据上的实践

在图结构化数据上的任务可以被归类为三个层级：节点层、边层和图层。这些不同的层级描述了我们希望在哪个层级上执行分类或回归任务。下面我们将更详细地讨论这三种类型。

4.1.节点层任务：半监督节点分类

节点层任务的目标是在一个图中对节点进行分类。通常，我们会得到一个大型图，包含超过1000个节点，其中一部分节点已经被标记。在训练过程中，我们学习对这些已标记的样本进行分类，并尝试将这种学习推广到未标记的节点上。

本文中，我们将使用一个流行的示例——Cora数据集，这是一个论文之间的引用网络。Cora包含2708篇科学出版物，它们之间通过链接相互引用。我们的任务是将每篇出版物分类到七个类别中的一个。每篇出版物由一个词袋模型向量表示，这意味着我们为每篇出版物拥有一个1433维的向量，其中第 个特征上的1表示预定义字典中的第 个词出现在文章中。当需要非常简单的编码，并且已经对网络中预期出现的词有直观理解时，通常会使用二进制词袋模型表示。虽然存在更好的方法，但我们将留给自然语言处理课程来讨论。
我们将在下面加载数据集：

cora_dataset = torch_geometric.datasets.Planetoid(root=DATASET_PATH, name="Cora")

让我们看看PyTorch Geometric是如何表示图数据的。需要注意的是，尽管我们只有一个图，但PyTorch Geometric返回一个数据集以兼容其他数据集。

cora_dataset[0]

图由一个Data对象表示（[文档](torch_geometric.data - pytorch_geometric documentation)），我们可以像访问标准Python命名空间一样访问它。边索引张量是图中边的列表，对于无向图，它包含每条边的镜像版本。train_mask、val_mask和test_mask是布尔掩码，指示我们应使用哪些节点进行训练、验证和测试。x张量是我们2708篇出版物的特征张量，而y是所有节点的标签。
在了解了数据之后，我们可以实现一个简单的图神经网络。GNN应用一系列图层（如GCN、GAT或GraphConv）、ReLU作为激活函数，以及dropout用于正则化。以下是具体的实现代码。

class GNNModel(nn.Module):
    
    def __init__(self, c_in, c_hidden, c_out, num_layers=2, layer_name="GCN", dp_rate=0.1, **kwargs):
        """
        参数：
            c_in - 输入特征的维度
            c_hidden - 隐藏特征的维度
            c_out - 输出特征的维度，通常是分类任务中的类别数
            num_layers - "隐藏"图层的数量
            layer_name - 要使用的图层类型
            dp_rate - 整个网络中应用的dropout率
            kwargs - 图层的额外参数（例如GAT的头数）
        """
        super().__init__()
        gnn_layer = gnn_layer_by_name[layer_name]
        
        layers = []
        in_channels, out_channels = c_in, c_hidden
        for l_idx in range(num_layers-1):
            layers += [
                gnn_layer(in_channels=in_channels, 
                          out_channels=out_channels,
                          **kwargs),
                nn.ReLU(inplace=True),
                nn.Dropout(dp_rate)
            ]
            in_channels = c_hidden
        layers += [gnn_layer(in_channels=in_channels, 
                             out_channels=c_out,
                             **kwargs)]
        self.layers = nn.ModuleList(layers)
    
    def forward(self, x, edge_index):
        """
        参数：
            x - 每个节点的输入特征
            edge_index - 表示图中边的顶点索引对列表（PyTorch Geometric的表示法）
        """
        for l in self.layers:
            # 对于图层，我们需要将"edge_index"张量作为额外输入
            # 所有PyTorch Geometric图层都继承自"MessagePassing"类，因此
            # 我们可以通过检查类类型来确定
            if isinstance(l, geom_nn.MessagePassing):
                x = l(x, edge_index)
            else:
                x = l(x)
        return x

在节点层任务中，一个良好的实践是创建一个独立应用于每个节点的MLP基线。这样我们可以验证是否向模型中添加图信息确实可以改进预测结果。也有可能每个节点的特征本身就足够表达，能够清晰地指向特定类别。为了验证这一点，我们在下面实现了一个简单的MLP。

class MLPModel(nn.Module):
    
    def __init__(self, c_in, c_hidden, c_out, num_layers=2, dp_rate=0.1):
        """
        参数：
            c_in - 输入特征的维度
            c_hidden - 隐藏特征的维度
            c_out - 输出特征的维度，通常是分类任务中的类别数
            num_layers - 隐藏层的数量
            dp_rate - 整个网络中应用的dropout率
        """
        super().__init__()
        layers = []
        in_channels, out_channels = c_in, c_hidden
        for l_idx in range(num_layers-1):
            layers += [
                nn.Linear(in_channels, out_channels),
                nn.ReLU(inplace=True),
                nn.Dropout(dp_rate)
            ]
            in_channels = c_hidden
        layers += [nn.Linear(in_channels, c_out)]
        self.layers = nn.Sequential(*layers)
    
    def forward(self, x, *args, **kwargs):
        """
        参数：
            x - 每个节点的输入特征
        """
        return self.layers(x)

最后，我们可以将模型合并到一个PyTorch Lightning模块中，它将为我们处理训练、验证和测试。

class NodeLevelGNN(pl.LightningModule):
    
    def __init__(self, model_name, **model_kwargs):
        super().__init__()
        # 保存超参数
        self.save_hyperparameters()
        
        if model_name == "MLP":
            self.model = MLPModel(**model_kwargs)
        else:
            self.model = GNNModel(**model_kwargs)
        self.loss_module = nn.CrossEntropyLoss()

    def forward(self, data, mode="train"):
        x, edge_index = data.x, data.edge_index
        x = self.model(x, edge_index)
        
        # 只有在对应的掩码节点上计算损失
        if mode == "train":
            mask = data.train_mask
        elif mode == "val":
            mask = data.val_mask
        elif mode == "test":
            mask = data.test_mask
        else:
            assert False, f"未知的前向模式: {mode}"
        
        loss = self.loss_module(x[mask], data.y[mask])
        acc = (x[mask].argmax(dim=-1) == data.y[mask]).sum().float() / mask.sum()
        return loss, acc

    def configure_optimizers(self):
        # 我们这里使用SGD，但Adam同样适用
        optimizer = optim.SGD(self.parameters(), lr=0.1, momentum=0.9, weight_decay=2e-3)
        return optimizer

    def training_step(self, batch, batch_idx):
        loss, acc = self.forward(batch, mode="train")
        self.log('train_loss', loss)
        self.log('train_acc', acc)
        return loss

    def validation_step(self, batch, batch_idx):
        _, acc = self.forward(batch, mode="val")
        self.log('val_acc', acc)

    def test_step(self, batch, batch_idx):
        _, acc = self.forward(batch, mode="test")
        self.log('test_acc', acc)

除了Lightning模块，我们下面定义了一个训练函数。由于我们只有一个图，我们为数据加载器使用批量大小为1，并为训练、验证和测试集共享相同的数据加载器（掩码在Lightning模块中选择）。此外，我们将参数enable_progress_bar设置为False，因为它通常显示每个epoch的进度，但一个epoch只包含一个步骤。其余的代码与我们在教程5和6中看到的非常相似。

def train_node_classifier(model_name, dataset, **model_kwargs):
    pl.seed_everything(42)
    node_data_loader = geom_data.DataLoader(dataset, batch_size=1)
    
    # 创建一个带有生成回调的PyTorch Lightning训练器
    root_dir = os.path.join(CHECKPOINT_PATH, "NodeLevel" + model_name)
    os.makedirs(root_dir, exist_ok=True)
    trainer = pl.Trainer(default_root_dir=root_dir,
                         callbacks=[ModelCheckpoint(save_weights_only=True, mode="max", monitor="val_acc")],
                         accelerator="gpu" if str(device).startswith("cuda") else "cpu",
                         devices=1,
                         max_epochs=200,
                         enable_progress_bar=False) # 因为epoch大小为1，所以设置为False
    trainer.logger._default_hp_metric = None # 可选的日志记录参数，我们不需要

    # 检查预训练模型是否存在。
如果存在，加载它并跳过训练
    pretrained_filename = os.path.join(CHECKPOINT_PATH, f"NodeLevel{model_name}.ckpt")
    if os.path.isfile(pretrained_filename):
        print("找到预训练模型，正在加载...")
        model = NodeLevelGNN.load_from_checkpoint(pretrained_filename)
    else:
        pl.seed_everything()
        model = NodeLevelGNN(model_name=model_name, c_in=dataset.num_node_features, c_out=dataset.num_classes, **model_kwargs)
        trainer.fit(model, node_data_loader, node_data_loader)
        model = NodeLevelGNN.load_from_checkpoint(trainer.checkpoint_callback.best_model_path)
    
    # 在测试集上测试最佳模型
    test_result = trainer.test(model, node_data_loader, verbose=False)
    batch = next(iter(node_data_loader))
    batch = batch.to(model.device)
    _, train_acc = model.forward(batch, mode="train")
    _, val_acc = model.forward(batch, mode="val")
    result = {"train": train_acc,
              "val": val_acc,
              "test": test_result[0]['test_acc']}
    return model, result

# 小函数，用于打印测试分数
def print_results(result_dict):
    if "train" in result_dict:
        print(f"训练准确率: {(100.0*result_dict['train']):4.2f}%")
    if "val" in result_dict:
        print(f"验证准确率:   {(100.0*result_dict['val']):4.2f}%")
    print(f"测试准确率:  {(100.0*result_dict['test']):4.2f}%")

# 训练简单的MLP
node_mlp_model, node_mlp_result = train_node_classifier(model_name="MLP",
                                                        dataset=cora_dataset,
                                                        c_hidden=16,
                                                        num_layers=2,
                                                        dp_rate=0.1)

print_results(node_mlp_result)

尽管MLP可能会因为高维输入特征而在训练数据集上过拟合，但它在测试集上的表现并不理想。让我们看看我们的图网络是否可以超越这个分数：

node_gnn_model, node_gnn_result = train_node_classifier(model_name="GNN",
                                                        layer_name="GCN",
                                                        dataset=cora_dataset, 
                                                        c_hidden=16, 
                                                        num_layers=2,
                                                        dp_rate=0.1)
print_results(node_gnn_result)

正如我们所希望的，GNN模型以相当大的优势超越了MLP。这表明使用图信息确实可以改进我们的预测，并让我们更好地泛化。
模型中的超参数被选择为创建一个相对较小的网络。这是因为输入维度为1433的第一层对于大型图来说可能相对昂贵。一般来说，对于非常大的图，GNN可能会变得相对昂贵。这就是为什么这样的GNN要么有一个小的隐藏尺寸，要么使用特殊的批处理策略，其中我们采样大的原始图的连通子图。

4.2.边缘级任务：链接预测

在某些应用场景中，我们可能需要对图的边缘进行预测，而非单个节点。图神经网络（GNN）中最常见的边缘级任务便是链接预测。所谓链接预测，即在给定一个图的情况下，预测两个节点之间是否存在或应该存在连接。例如，在社交网络中，Facebook等平台会利用这一技术向你推荐可能认识的朋友。显然，图的整体信息在执行这一任务时至关重要。通常，预测结果的生成会基于节点特征对的相似性度量，如果两个节点间应该存在连接，则结果为1；否则接近0。为了保持本教程的简洁性，我们不会亲自实现这一任务。然而，如果你对这一任务感兴趣，可以参考以下资源深入了解。
相关教程和论文包括：

• [PyTorch Geometric 示例]
• [图神经网络：方法和应用的综述]，周等人，2019年
• [基于图神经网络的链接预测]，张和陈，2018年。

4.3.图级别任务：图分类

在本教程的这一部分，我们将深入探讨如何将图神经网络（GNN）应用于图分类任务。与之前不同，我们的目标是分类整个图，而非单个节点或边缘。因此，我们手头也有一个包含多个图的数据集，需要根据图的结构属性对这些图进行分类。图分类中最常见的任务之一是分子属性预测，其中分子被表示为图的形式。每个原子对应一个节点，而原子之间的键则作为图的边。例如，参考下图。

在图的左侧，我们展示了一个包含不同原子的小分子，而右侧则展示了其图表示。原子类型被抽象为节点特征（例如，一个独热向量），而不同的键类型则作为边特征。为了简化，本教程将忽略边属性，但您可以通过使用像[关系图卷积]这样的方法来包含它们，该方法为每种边类型使用不同的权重矩阵。
我们将使用的数据集称为MUTAG数据集。它是图分类算法的常用小型基准，包含188个图，每个图平均有18个节点和20条边。图节点有7种不同的标签/原子类型，二元图标签代表“它们对特定革兰氏阴性细菌的突变效应”（标签的具体含义在这里不是很重要）。该数据集是一系列不同的图分类数据集的一部分，称为TUDatasets，可以通过PyTorch Geometric中的torch_geometric.datasets.TUDataset直接访问（[文档](torch_geometric.datasets - pytorch_geometric documentation))。下面我们将加载这个数据集。

tu_dataset = torch_geometric.datasets.TUDataset(root=DATASET_PATH, name="MUTAG")

让我们来查看一下数据集的一些统计信息：

print("数据对象：", tu_dataset.data)
print("长度：", len(tu_dataset))
print(f"平均标签：{tu_dataset.data.y.float().mean().item():4.2f}")

第一行显示了数据集如何存储不同的图。每个图的节点、边和标签被连接成一个张量，数据集存储了相应地分割张量的索引。数据集的长度是我们拥有的图的数量，而“平均标签”表示标签为1的图的百分比。只要百分比在0.5的范围内，我们就有一个相对平衡的数据集。图数据集往往非常不平衡，因此检查类别平衡始终是一个好的实践。
接下来，我们将数据集分为训练和测试部分。这次我们不使用验证集，因为数据集的大小很小。因此，我们的模型可能会略微过拟合验证集，但由于评估的噪声，我们仍然可以得到对未训练数据性能的估计。

torch.manual_seed(42)
tu_dataset.shuffle()
train_dataset = tu_dataset[:150]
test_dataset = tu_dataset[150:]

在使用数据加载器时，我们遇到了一个与批处理$N$个图有关的问题。批次中的每个图可以具有不同数量的节点和边，因此我们需要大量的填充来获得一个单一的张量。Torch geometric 使用了一种不同的、更有效的方法：我们可以将批次中的$N$个图视为一个带有连接节点和边列表的大型图。由于来自两个不同图的任何节点之间没有边，因此在大图上运行GNN层会给我们与分别在每个图上运行GNN相同的输出。从视觉上看，这种批处理策略在下面可视化（图示来源 - PyTorch Geometric团队，[教程在这里])。

对于来自两个不同图的任何节点，邻接矩阵为零，否则根据各个图的邻接矩阵。幸运的是，这种策略已经在torch geometric中实现，因此我们可以使用相应的数据加载器：

graph_train_loader = geom_data.DataLoader(train_dataset, batch_size=64, shuffle=True)
graph_val_loader = geom_data.DataLoader(test_dataset, batch_size=64) # 如果你想更改为更大的数据集，可以添加额外的加载器
graph_test_loader = geom_data.DataLoader(test_dataset, batch_size=64)

让我们在下面加载一个批次，看看批处理的实际操作：

batch = next(iter(graph_test_loader))
print("批次：", batch)
print("标签：", batch.y[:10])
print("批次索引：", batch.batch[:40])

我们为测试数据集堆叠了38个图。存储在batch中的批次索引显示前12个节点属于第一个图，接下来的22个属于第二个图，以此类推。这些索引对于进行最终预测很重要。为了对整个图进行预测，我们通常在运行GNN模型后对所有节点执行池化操作。在这种情况下，我们将使用平均池化。因此，我们需要知道哪些节点应该包含在哪个平均池中。使用这种池化，我们已经可以创建下面的图网络。具体来说，我们重用之前的GNNModel类，并简单地添加了一个平均池和用于图预测任务的单个线性层。

class GraphGNNModel(nn.Module):
    
    def __init__(self, c_in, c_hidden, c_out, dp_rate_linear=0.5, **kwargs):
        """
        参数：
            c_in - 输入特征的维度
            c_hidden - 隐藏特征的维度
            c_out - 输出特征的维度（通常是类别数量）
            dp_rate_linear - 线性层前的丢弃率（通常比GNN内部的高得多）
            kwargs - GNNModel对象的额外参数
        """
        super().__init__()
        self.GNN = GNNModel(c_in=c_in, 
                            c_hidden=c_hidden, 
                            c_out=c_hidden, # 还不是我们的预测输出！
                            **kwargs)
        self.head = nn.Sequential(
            nn.Dropout(dp_rate_linear),
            nn.Linear(c_hidden, c_out)
        )

    def forward(self, x, edge_index, batch_idx):
        """
        参数：
            x - 每个节点的输入特征
            edge_index - 表示图中边的顶点索引对列表（PyTorch几何符号）
            batch_idx - 每个节点的批次元素索引
        """
        x = self.GNN(x, edge_index)
        x = geom_nn.global_mean_pool(x, batch_idx) # 平均池化
        x = self.head(x)
        return x

最后，我们可以创建一个PyTorch Lightning模块来处理训练。它与我们之前见过的模块类似，在训练方面没有令人惊讶的地方。由于我们有一个二元分类任务，我们使用二元交叉熵损失。

class GraphLevelGNN(pl.LightningModule):
    
    def __init__(self, **model_kwargs):
        super().__init__()
        # 保存超参数
        self.save_hyperparameters()
        
        self.model = GraphGNNModel(**model_kwargs)
        self.loss_module = nn.BCEWithLogitsLoss() if self.hparams.c_out == 1 else nn.CrossEntropyLoss()

    def forward(self, data, mode="train"):
        x, edge_index, batch_idx = data.x, data.edge_index, data.batch
        x = self.model(x, edge_index, batch_idx)
        x = x.squeeze(dim=-1)
        
        ifself.hparams.c_out == 1:
            preds = (x > 0).float() # 预测为正类的概率
            data.y = data.y.float() # 将标签转换为浮点数
        else:
            preds = x.argmax(dim=-1) # 预测最可能的类别
        loss = self.loss_module(x, data.y) # 计算损失
        acc = (preds == data.y).sum().float() / preds.shape[0] # 计算准确率
        return loss, acc

    def configure_optimizers(self):
        optimizer = optim.AdamW(self.parameters(), lr=1e-2, weight_decay=0.0) # 高学习率，因为数据集和模型很小
        return optimizer

    def training_step(self, batch, batch_idx):
        loss, acc = self.forward(batch, mode="train")
        self.log('train_loss', loss)
        self.log('train_acc', acc)
        return loss

    def validation_step(self, batch, batch_idx):
        _, acc = self.forward(batch, mode="val")
        self.log('val_acc', acc)

    def test_step(self, batch, batch_idx):
        _, acc = self.forward(batch, mode="test")
        self.log('test_acc', acc)

下面我们在数据集上训练模型。它类似于我们迄今为止见过的典型训练函数。

def train_graph_classifier(model_name, **model_kwargs):
    pl.seed_everything(42)
    
    # 创建一个带有生成回调的PyTorch Lightning训练器
    root_dir = os.path.join(CHECKPOINT_PATH, "GraphLevel" + model_name)
    os.makedirs(root_dir, exist_ok=True)
    trainer = pl.Trainer(default_root_dir=root_dir,
                         callbacks=[ModelCheckpoint(save_weights_only=True, mode="max", monitor="val_acc")],
                         accelerator="gpu" if str(device).startswith("cuda") else "cpu",
                         devices=1,
                         max_epochs=500,
                         enable_progress_bar=False)
    trainer.logger._default_hp_metric = None # 可选的日志记录参数，我们不需要

    # 检查是否存在预训练模型。如果存在，加载它并跳过训练
    pretrained_filename = os.path.join(CHECKPOINT_PATH, f"GraphLevel{model_name}.ckpt")
    if os.path.isfile(pretrained_filename):
        print("找到预训练模型，正在加载...")
        model = GraphLevelGNN.load_from_checkpoint(pretrained_filename)
    else:
        pl.seed_everything(42)
        model = GraphLevelGNN(c_in=tu_dataset.num_node_features, 
                              c_out=1 if tu_dataset.num_classes==2 else tu_dataset.num_classes, 
                              **model_kwargs)
        trainer.fit(model, graph_train_loader, graph_val_loader)
        model = GraphLevelGNN.load_from_checkpoint(trainer.checkpoint_callback.best_model_path)
    # 在验证集和测试集上测试最佳模型
    train_result = trainer.test(model, graph_train_loader, verbose=False)
    test_result = trainer.test(model, graph_test_loader, verbose=False)
    result = {"test": test_result[0]['test_acc'], "train": train_result[0]['test_acc']} 
    return model, result

最后，让我们进行训练和测试。请随意尝试不同的GNN层、超参数等。

model, result = train_graph_classifier(model_name="GraphConv", 
                                       c_hidden=256, 
                                       layer_name="GraphConv", 
                                       num_layers=3, 
                                       dp_rate_linear=0.5,
                                       dp_rate=0.0)
print(f"训练性能：{100.0*result['train']:4.2f}%")
print(f"测试性能：  {100.0*result['test']:4.2f}%")

测试性能显示我们在数据集的未见过的部分上获得了相当好的分数。需要注意的是，由于我们一直在使用测试集进行验证，我们可能对这个集合有过拟合。尽管如此，实验表明GNN确实可以强大地预测图和/或分子的属性。

5.结论

通过本文的学习，我们深入理解了神经网络在处理图结构数据方面的应用。我们探讨了图的两种常见表示方法——邻接矩阵和边列表，并详细分析了图卷积网络（GCN）和图注意力网络（GAT）这两种典型的图神经网络层的实现。这些实现让我们认识到，虽然图神经网络的理论基础可能相当复杂，但其在实际应用中的操作往往更为直观和简单。
在本文的实践部分，我们尝试了多种不同的任务，涵盖了从单个节点到整个图的不同层级。这些实践进一步证实了图信息在预测任务中的重要性，特别是在需要捕捉节点间复杂关系的场景中。我们发现，将图结构信息融入预测模型，对于提升模型性能至关重要。
目前，图神经网络已经在众多领域展现出其强大的应用潜力，并且随着研究的深入，它们在未来的重要性和影响力预计会持续增长。无论是在社交网络分析、生物信息学、交通网络优化，还是在推荐系统等领域，GNNs都正在发挥着越来越关键的作用。随着技术的不断进步，我们可以预见，图神经网络将在未来的人工智能领域扮演更加重要的角色。