许多数据集本质上是关系型的,最适合用图来表示。例子有社交网络、分子结构、引用网络、知识图谱和推荐系统。传统深度学习结构,如 CNN 和 RNN,假定数据结构是网格状或序列状的,这使它们不适合处理图中存在的任意连接。图神经网络 (GNN) 专门设计用于直接处理图结构数据,它们学习到的表示会同时包含节点特征和图的拓扑结构。PyTorch Geometric (PyG) 是一个功能强大且被广泛采用的库,它建立在 PyTorch 之上,用于开发和应用 GNN。它提供了多种 GNN 层的优化实现、高效的图数据处理以及常见的图基准数据集。本节将指导您如何使用 PyG 来实现和理解不同的 GNN 结构。在 PyTorch Geometric 中表示图在构建 GNN 模型之前,我们需要一种标准化的方式来表示图数据。PyG 使用 torch_geometric.data.Data 对象。一个 Data 对象包含描述单个图的各种属性:x: 节点特征矩阵,形状为 [num_nodes, num_node_features]。每行代表一个节点,列代表其特征。edge_index: 图的连接信息,采用 COO (坐标) 格式,形状为 [2, num_edges]。它存储每条边的源节点和目标节点索引。对于从节点 j 到节点 i 的边,其列为 [j, i]。这种表示对于稀疏图是高效的。edge_attr: 边特征矩阵,形状为 [num_edges, num_edge_features]。表示与每条边相关的可选特征。y: 目标标签或值,取决于具体任务。对于节点级任务,形状为 [num_nodes, ...];对于图级任务,形状为 [1, ...]。pos: 节点位置特征,形状为 [num_nodes, num_dimensions]。常用于几何深度学习。下面是创建简单 Data 对象的方法:import torch from torch_geometric.data import Data # 节点特征:3 个节点,每个节点 2 个特征 x = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.float) # 边:(0 -> 1), (1 -> 0), (1 -> 2), (2 -> 1) # 表示为源节点和目标节点 edge_index = torch.tensor([[0, 1, 1, 2], # 源节点 [1, 0, 2, 1]], # 目标节点 dtype=torch.long) # 可选的边特征:4 条边,每条边 1 个特征 edge_attr = torch.tensor([[0.5], [0.5], [0.8], [0.8]], dtype=torch.float) # 可选的节点标签(例如,用于节点分类) y = torch.tensor([0, 1, 0], dtype=torch.long) # 创建 Data 对象 graph_data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr, y=y) print(graph_data) # 输出:Data(x=[3, 2], edge_index=[2, 4], edge_attr=[4, 1], y=[3])PyG 还提供了 torch_geometric.data.Dataset 和 torch_geometric.loader.DataLoader,用于高效处理图集合并创建小批量数据。DataLoader 会自动将不同大小的图整理成更大的批处理对象。消息传递方法大多数 GNN 层都基于消息传递原理运行。核心思想是每个节点通过聚合来自其局部邻域的信息,迭代地更新其特征表示(嵌入)。这个过程通常包含对层 $l$ 中每个节点 $i$ 的三个步骤:消息计算: 每个邻居节点 $j \in \mathcal{N}(i)$ 根据其自身特征 $\mathbf{h}j^{(l-1)}$,以及目标节点特征 $\mathbf{h}i^{(l-1)}$ 和边特征 $\mathbf{e}{j,i}$,计算一个消息 $\mathbf{m}{j \to i}^{(l)}$。 $$ \mathbf{m}_{j \to i}^{(l)} = \phi^{(l)}(\mathbf{h}_i^{(l-1)}, \mathbf{h}j^{(l-1)}, \mathbf{e}{j,i}) $$ 其中 $\phi^{(l)}$ 是一个可微分的消息函数(例如,一个神经网络)。聚合: 节点 $i$ 使用一个置换不变函数 $\bigoplus$(如求和、平均或最大值)来聚合来自其邻居的所有传入消息。 $$ \mathbf{a}i^{(l)} = \bigoplus{j \in \mathcal{N}(i)} \mathbf{m}_{j \to i}^{(l)} $$更新: 节点 $i$ 根据其先前的表示 $\mathbf{h}_i^{(l-1)}$ 和聚合后的消息 $\mathbf{a}_i^{(l)}$ 来更新其特征向量 $\mathbf{h}_i^{(l)}$。 $$ \mathbf{h}_i^{(l)} = \gamma^{(l)}(\mathbf{h}_i^{(l-1)}, \mathbf{a}_i^{(l)}) $$ 其中 $\gamma^{(l)}$ 是一个可微分的更新函数(例如,另一个神经网络或简单地添加聚合消息)。初始特征 $\mathbf{h}_i^{(0)}$ 通常是输入节点特征 data.x。堆叠多个消息传递层可以使信息在图中传播更远的距离。digraph G {rankdir=LR;node [shape=circle, style=filled, fillcolor="#a5d8ff", fontcolor="#1c7ed6", fontsize=10, width=0.4, height=0.4, fixedsize=true];edge [arrowsize=0.7, color="#868e96"];subgraph cluster_l_minus_1 {label="层 l-1";bgcolor="#e9ecef";style=rounded;node [fillcolor="#74c0fc"];h_i_prev [label="hᵢ(l-1)"];h_j1_prev [label="hⱼ₁(l-1)"];h_j2_prev [label="hⱼ₂(l-1)"];h_j3_prev [label="hⱼ₃(l-1)"];}subgraph cluster_l {label="层 l";bgcolor="#e9ecef";style=rounded;node [fillcolor="#4dabf7"];h_i [label="hᵢ(l)"];}h_j1_prev -> h_i [label="m₁", fontsize=9, fontcolor="#495057"];h_j2_prev -> h_i [label="m₂", fontsize=9, fontcolor="#495057"];h_j3_prev -> h_i [label="m₃", fontsize=9, fontcolor="#495057"];h_i_prev -> h_i [label="更新(hᵢ(l-1), 聚合)", style=dashed, fontsize=9, fontcolor="#495057"];edge [style=invis];h_j1_prev -> h_i_prev;h_j2_prev -> h_i_prev;h_j3_prev -> h_i_prev;} 此图表呈现了更新节点 $i$ 的消息传递理念。来自邻居 $j_1, j_2, j_3$ 的信息(消息 $m_1, m_2, m_3$)被聚合,并与节点的先前状态 $h_i^{(l-1)}$ 结合,以计算出新状态 $h_i^{(l)}$。PyG 在其层类中提供了这些步骤的优化实现。PyTorch Geometric 中常用的 GNN 层PyG 提供了多种预实现的 GNN 层。让我们看看三个流行的例子:GCN、GraphSAGE 和 GAT。图卷积网络 (GCN)GCN 层由 Kipf & Welling (2017) 提出,执行基于谱的图卷积。GCN 层的消息传递更新规则可以简化为:$$ \mathbf{H}^{(l+1)} = \sigma\left(\hat{\mathbf{D}}^{-1/2} \hat{\mathbf{A}} \hat{\mathbf{D}}^{-1/2} \mathbf{H}^{(l)} \mathbf{W}^{(l)}\right) $$其中,$\mathbf{H}^{(l)}$ 是层 $l$ 的节点嵌入矩阵,$\mathbf{W}^{(l)}$ 是一个可训练的权重矩阵,$\sigma$ 是一个激活函数(如 ReLU),$\hat{\mathbf{A}} = \mathbf{A} + \mathbf{I}$ 是添加了自循环的邻接矩阵,$\hat{\mathbf{D}}$ 是 $\hat{\mathbf{A}}$ 的对角度矩阵。项 $\hat{\mathbf{D}}^{-1/2} \hat{\mathbf{A}} \hat{\mathbf{D}}^{-1/2}$ 表示邻接矩阵的对称归一化。该层平均邻居节点(包括节点自身)的特征,然后应用线性变换,再进行非线性处理。在 PyG 中,您使用 torch_geometric.nn.GCNConv:import torch.nn.functional as F from torch_geometric.nn import GCNConv class SimpleGCN(torch.nn.Module): def __init__(self, num_node_features, num_classes, hidden_channels): super().__init__() self.conv1 = GCNConv(num_node_features, hidden_channels) self.conv2 = GCNConv(hidden_channels, 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, p=0.5, training=self.training) # 经常使用 Dropout x = self.conv2(x, edge_index) # 对于节点分类,通常使用 LogSoftmax return F.log_softmax(x, dim=1)GraphSAGEGraphSAGE (Hamilton et al., 2017) 专注于学习聚合函数,而不是固定的卷积。它被设计为归纳式的,这意味着它可以在推断时推广到未见过的节点。GraphSAGE 为每个节点采样固定大小的邻域,然后使用平均、最大值或 LSTM 池化等函数聚合邻居特征。主要步骤包括:为节点 $i$ 采样一个邻域 $\mathcal{N}(i)$。聚合邻居特征:$\mathbf{a}_{\mathcal{N}(i)}^{(l)} = \text{AGGREGATE}^{(l)}({\mathbf{h}_j^{(l-1)} \mid j \in \mathcal{N}(i)})$更新节点 $i$ 的嵌入:$\mathbf{h}_i^{(l)} = \sigma\left( \mathbf{W}^{(l)} \cdot \text{CONCAT}(\mathbf{h}i^{(l-1)}, \mathbf{a}{\mathcal{N}(i)}^{(l)}) \right)$PyG 使用 torch_geometric.nn.SAGEConv 实现此功能:from torch_geometric.nn import SAGEConv class SimpleGraphSAGE(torch.nn.Module): def __init__(self, num_node_features, num_classes, hidden_channels): super().__init__() # 默认聚合器是 'mean' self.conv1 = SAGEConv(num_node_features, hidden_channels) self.conv2 = SAGEConv(hidden_channels, 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, p=0.5, training=self.training) x = self.conv2(x, edge_index) return F.log_softmax(x, dim=1)在创建 SAGEConv 层时,您可以指定聚合器类型(例如,aggr='max'、aggr='mean')。图注意力网络 (GAT)GAT 层 (Veličković et al., 2018) 引入了注意力机制,使得节点在聚合过程中可以为其邻居分配不同的重要性权重。这使得聚合过程更加灵活,并通常带来更好的性能。节点 $i$ 和邻居 $j$ 之间的注意力系数 $e_{ij}$ 基于它们的特征计算,通常使用共享的线性变换和一个注意力机制(例如,一个单层前馈网络):$$ e_{ij} = \text{attention}(\mathbf{W}^{(l)}\mathbf{h}_i^{(l-1)}, \mathbf{W}^{(l)}\mathbf{h}_j^{(l-1)}) $$然后,这些系数使用 softmax 函数对节点 $i$ 的所有邻居进行归一化:$$ \alpha_{ij} = \text{softmax}j(e{ij}) = \frac{\exp(e_{ij})}{\sum_{k \in \mathcal{N}(i)} \exp(e_{ik})} $$聚合后的消息是转换后的邻居特征的加权和:$$ \mathbf{a}i^{(l)} = \sum{j \in \mathcal{N}(i)} \alpha_{ij} \mathbf{W}^{(l)} \mathbf{h}_j^{(l-1)} $$更新步骤将此聚合消息与节点自身的特征结合,通常使用拼接后跟激活函数:$$ \mathbf{h}_i^{(l)} = \sigma\left( \mathbf{a}_i^{(l)} \right) \quad \text{or} \quad \mathbf{h}_i^{(l)} = \sigma\left( \text{CONCAT}(\mathbf{h}_i^{(l-1)}, \mathbf{a}_i^{(l)}) \right) $$GAT 经常使用多头注意力,其中计算多个独立的注意力机制,并将其结果进行拼接或平均。PyG 使用 torch_geometric.nn.GATConv 实现此功能:from torch_geometric.nn import GATConv class SimpleGAT(torch.nn.Module): def __init__(self, num_node_features, num_classes, hidden_channels, heads=8): super().__init__() # 在第一层中使用多头注意力 self.conv1 = GATConv(num_node_features, hidden_channels, heads=heads, dropout=0.6) # 多头注意力的输出特征为 heads * hidden_channels # 对于最后一层,通常平均各头或使用单头 self.conv2 = GATConv(hidden_channels * heads, num_classes, heads=1, concat=False, dropout=0.6) def forward(self, data): x, edge_index = data.x, data.edge_index x = F.dropout(x, p=0.6, training=self.training) # 对输入特征应用 Dropout x = self.conv1(x, edge_index) x = F.elu(x) # ELU 激活在 GAT 中很常见 x = F.dropout(x, p=0.6, training=self.training) x = self.conv2(x, edge_index) return F.log_softmax(x, dim=1) 构建和训练 GNN 模型使用 PyG 层在 PyTorch 中构建 GNN 遵循标准的 PyTorch 实践。您定义一个继承自 torch.nn.Module 的类,在 __init__ 中初始化 PyG 层,并在 forward 中定义前向传播逻辑。forward 方法通常接受 Data 或 Batch 对象作为输入,并提取 x、edge_index,以及可能的 edge_attr 和 batch 索引。训练循环也类似于标准的 PyTorch 循环:遍历 DataLoader,执行前向传播,计算损失(例如,用于节点分类的 F.nll_loss,配合 log_softmax),使用 loss.backward() 计算梯度,并使用优化器更新参数。常见的 GNN 应用GNN 功能多样,可应用于各种图相关任务:节点分类: 预测图中单个节点的标签或属性(例如,对社交网络中的用户进行分类,预测蛋白质功能)。上述示例 (SimpleGCN、SimpleGraphSAGE、SimpleGAT) 均适用于节点分类。图分类: 预测整个图的标签或属性(例如,将分子分类为有毒或无毒,对社交群体进行分类)。这需要在 GNN 层之后使用图池化层(例如,torch_geometric.nn.global_mean_pool、global_max_pool)将节点嵌入聚合成单个图嵌入。链接预测: 预测两个节点之间是否存在或将存在边(例如,在社交网络中推荐朋友,预测蛋白质-蛋白质相互作用)。这通常包括学习节点嵌入,然后使用评分函数(例如,点积)对节点嵌入对进行处理。图生成: 生成具有期望属性的新图。社区检测: 在大型图中识别连接紧密的节点群组。PyTorch Geometric 为应对这些任务提供了全面的工具。通过结合其优化层、数据处理工具和标准 PyTorch 功能,您可以有效地构建和训练处理复杂图问题的精巧 GNN 模型。请记住,选择合适的 GNN 结构(GCN、GAT、SAGE 或其他)通常取决于您的图数据的具体特征和当前的任务。进行实验和理解每个层的基本原理对于成功应用非常重要。