PyG 图神经网络设计指南从 MessagePassing 基类到异构图学习实战【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric本指南以 PyTorch GeometricPyG官方教程《Design of Graph Neural Networks》为主线系统讲解如何利用MessagePassing基类从零实现 GCN、EdgeConv 等消息传递网络并进一步掌握异构图Heterogeneous Graph的建模、变换与三类异构 GNN 构建方案。读完本文你将能够自定义任意消息传递算子并把同构图模型一键迁移到 ogbn-mag 等真实异构数据集上完成训练。消息传递范式GNN 的统一数学框架把卷积算子推广到不规则域如图、点云上通常被表达为**邻域聚合neighborhood aggregation或消息传递message passing**范式。设 $\mathbf{x}^{(k-1)}i \in \mathbb{R}^F$ 表示节点 $i$ 在第 $(k-1)$ 层的特征$\mathbf{e}{j,i} \in \mathbb{R}^D$ 表示从节点 $j$ 指向节点 $i$ 的可选边特征则消息传递图神经网络可以写成$$ \mathbf{x}_i^{(k)} \gamma^{(k)} \left( \mathbf{x}i^{(k-1)}, \bigoplus{j \in \mathcal{N}(i)} , \phi^{(k)}\left(\mathbf{x}_i^{(k-1)}, \mathbf{x}j^{(k-1)},\mathbf{e}{j,i}\right) \right) $$其中 $\bigoplus$ 表示一个可微的、置换不变的聚合函数例如 sum、mean 或 max$\gamma$ 与 $\phi$ 表示可微函数例如 MLP多层感知机。直观地说每个节点先从邻居收集消息由 $\phi$ 生成再按置换不变的方式聚合由 $\bigoplus$ 完成最后结合自身信息更新表示由 $\gamma$ 完成。几乎所有主流 GNN 算子——GCN、GraphSAGE、GAT、GIN、EdgeConv——都是该公式的特例这正是 PyG 设计MessagePassing基类的理论根基。MessagePassing 基类PyG 提供的消息传递脚手架PyG 在 message_passing.py 中提供了MessagePassing基类它自动处理消息的传播propagation流程用户只需定义三个要素$\phi$即message()函数$\gamma$即update()函数聚合方案即aggradd、aggrmean或aggrmax。构造函数参数MessagePassing的构造签名见 源码 L110-L118如下参数说明默认值aggr聚合方案add/sum、mean、min、max、mul也可以是Aggregation模块或字符串列表列表时各聚合结果在最后一维拼接sum与add等价aggr_kwargs传递给自动解析出的聚合函数的额外参数Noneflow消息传递方向source_to_target或target_to_sourcesource_to_targetnode_dim沿哪个轴传播消息批量维度处理-2decomposed_layers特征分解层数用于降低峰值内存、加速 CPU 推理1注意源码中flow只接受source_to_target与target_to_source两个取值传入其他值会抛出ValueError见 L121-L123。decomposed_layers通过把特征维切片成多个子层分批做聚合可在 CPU 上显著降低峰值内存但不适用于注意力类 GNN消息无法简单分解。四个核心方法propagate(edge_index, sizeNone, **kwargs)启动消息传播的入口源码 L421。它接收边索引以及构造消息、更新节点嵌入所需的全部附加数据。关键特性propagate不只限于形状为 $[N, N]$ 的方阵邻接矩阵也可以处理形状为 $[N, M]$ 的一般稀疏赋值矩阵即二部图只需传入size(N, M)若为None则默认是方阵。对于具有两类独立节点的二部图若两类节点各自持有信息可用元组传入例如x(x_N, x_M)。message(...)对应于 $\phi$为每条边 $(j,i) \in \mathcal{E}$当flowsource_to_target时构造发往节点 $i$ 的消息。它可以接收任何传给propagate的参数更妙的是把变量名加上_i或_j后缀PyG 会自动把张量映射到对应的目标/源节点例如x_i、x_j。这里约定 $i$ 是聚合信息的中心节点$j$ 是邻居节点。aggregate(...)执行 $\bigoplus$ 聚合源码 L577。默认委托给__init__中由aggr解析出的Aggregation模块。update(aggr_out, ...)对应于 $\gamma$对每个节点 $i \in \mathcal{V}$ 更新节点嵌入源码 L609。第一个参数是聚合输出其余参数来自propagate的输入。propagate内部会依次调用message→aggregate→update见 源码 L499-L550。其底层实现中有两个值得关注的机制Inspector 参数收集构造时用Inspector反射message/aggregate/update的签名自动确定需要从propagate的**kwargs中收集哪些参数L138-L151。张量抬升lifting_lift通过index_select(self.node_dim, edge_index[dim])把节点特征按边索引展开成边级张量这就是x_j的来源_collect则根据flow决定_i/_j分别取edge_index的第 1/0 行。只要张量持有源或目标节点特征任何张量都可以用_i/_j后缀自动抬升。下面通过重实现两个经典算子——GCN 层Kipf Welling与 EdgeConv 层Wang et al.——来验证这套机制。实战一从零实现 GCN 层GCN 层的数学定义为$$ \mathbf{x}i^{(k)} \sum{j \in \mathcal{N}(i) \cup { i }} \frac{1}{\sqrt{\deg(i)} \cdot \sqrt{\deg(j)}} \cdot \left( \mathbf{W}^{\top} \cdot \mathbf{x}_j^{(k-1)} \right) \mathbf{b} $$即邻居节点特征先经权重矩阵 $\mathbf{W}$ 线性变换再按度数归一化最后求和对聚合输出施加偏置向量 $\mathbf{b}$。这个公式可拆解为六个步骤给邻接矩阵添加自环self-loops线性变换节点特征矩阵计算归一化系数在 $\phi$message中归一化节点特征求和聚合邻居节点特征add聚合施加最终偏置向量。步骤 1–3 通常在消息传递前完成步骤 4–5 可直接交给MessagePassing基类。完整实现如下与教程 create_gnn.rst 一致import torch from torch.nn import Linear, Parameter from torch_geometric.nn import MessagePassing from torch_geometric.utils import add_self_loops, degree class GCNConv(MessagePassing): def __init__(self, in_channels, out_channels): super().__init__(aggradd) # Add aggregation (Step 5). self.lin Linear(in_channels, out_channels, biasFalse) self.bias Parameter(torch.empty(out_channels)) self.reset_parameters() def reset_parameters(self): self.lin.reset_parameters() self.bias.data.zero_() def forward(self, x, edge_index): # x has shape [N, in_channels] # edge_index has shape [2, E] # Step 1: Add self-loops to the adjacency matrix. edge_index, _ add_self_loops(edge_index, num_nodesx.size(0)) # Step 2: Linearly transform node feature matrix. x self.lin(x) # Step 3: Compute normalization. row, col edge_index deg degree(col, x.size(0), dtypex.dtype) deg_inv_sqrt deg.pow(-0.5) deg_inv_sqrt[deg_inv_sqrt float(inf)] 0 norm deg_inv_sqrt[row] * deg_inv_sqrt[col] # Step 4-5: Start propagating messages. out self.propagate(edge_index, xx, normnorm) # Step 6: Apply a final bias vector. out out self.bias return out def message(self, x_j, norm): # x_j has shape [E, out_channels] # Step 4: Normalize node features. return norm.view(-1, 1) * x_j逐段解读super().__init__(aggradd)选择了求和聚合步骤 5。PyG 内部把add与sum视为同一聚合见FUSE_AGGRS {add, sum, mean, min, max}message_passing.py L35。add_self_loops来自 utils/loop.py为edge_index补充 $(i,i)$ 自环步骤 1。PyG 官方实现的GCNConvgcn_conv.py更进一步使用add_remaining_self_loops避免重复添加并通过gcn_normgcn_conv.py L45-L113在稀疏矩阵与稠密张量两种表示下统一完成自环 对称归一化。degree(col, ...)来自 utils/_degree.py计算每个节点的度数deg_inv_sqrt中无穷大被置 0 以处理孤立节点步骤 3。propagate(edge_index, xx, normnorm)内部依次调用message、aggregate、update。在message(self, x_j, norm)中x_j是抬升lifted张量包含每条边源节点的特征即各节点的邻居特征norm同样按边索引展开后与x_j逐元素相乘完成归一化步骤 4。初始化与调用极其简单该层可直接作为深度架构的积木conv GCNConv(16, 32) x conv(x, edge_index)实战二实现 Edge 卷积层EdgeConv 层用于处理图或点云数学定义为$$ \mathbf{x}i^{(k)} \max{j \in \mathcal{N}(i)} h_{\mathbf{\Theta}} \left( \mathbf{x}_i^{(k-1)}, \mathbf{x}_j^{(k-1)} - \mathbf{x}_i^{(k-1)} \right) $$其中 $h_{\mathbf{\Theta}}$ 是一个 MLP。与 GCN 类似这次改用max聚合import torch from torch.nn import Sequential as Seq, Linear, ReLU from torch_geometric.nn import MessagePassing class EdgeConv(MessagePassing): def __init__(self, in_channels, out_channels): super().__init__(aggrmax) # Max aggregation. self.mlp Seq(Linear(2 * in_channels, out_channels), ReLU(), Linear(out_channels, out_channels)) def forward(self, x, edge_index): # x has shape [N, in_channels] # edge_index has shape [2, E] return self.propagate(edge_index, xx) def message(self, x_i, x_j): # x_i has shape [E, in_channels] # x_j has shape [E, in_channels] tmp torch.cat([x_i, x_j - x_i], dim1) # tmp has shape [E, 2 * in_channels] return self.mlp(tmp)在message中我们同时用到了x_i与x_jx_i是目标节点特征x_j - x_i是相对源节点特征刻画了每条边两端点的差分信息二者在特征维拼接后送入 MLP。官方 edge_conv.py 中的EdgeConv实现与此一致L60-L61差别仅在于官方把 MLP 作为nn参数从外部注入、message用dim-1拼接且forward支持二部图输入x (x, x)。EdgeConv 本质上是动态卷积每一层都在特征空间中用最近邻重新构造图。PyG 提供了 GPU 加速的批量 k-NN 图生成方法torch_geometric.nn.pool.knn_graphfrom torch_geometric.nn import knn_graph class DynamicEdgeConv(EdgeConv): def __init__(self, in_channels, out_channels, k6): super().__init__(in_channels, out_channels) self.k k def forward(self, x, batchNone): edge_index knn_graph(x, self.k, batch, loopFalse, flowself.flow) return super().forward(x, edge_index)knn_graph计算最近邻图再调用EdgeConv.forward完成消息传递。官方 DynamicEdgeConv 还额外要求pyg-lib0.6.0支持并支持num_workers并行计算 k-NN。调用接口干净利落conv DynamicEdgeConv(3, 128, k6) x conv(x, batch)动手练习验证你的理解教程 create_gnn.rst 提供了一个练习数据集import torch from torch_geometric.data import Data edge_index torch.tensor([[0, 1], [1, 0], [1, 2], [2, 1]], dtypetorch.long) x torch.tensor([[-1], [0], [1]], dtypetorch.float) data Data(xx, edge_indexedge_index.t().contiguous())围绕GCNConv思考以下问题答案都可从源码与上述解读推出row和col分别保存什么信息row edge_index[0]为源节点col edge_index[1]为目标节点degree函数做了什么统计每个节点的度数即入边数量为什么用degree(col, ...)而不是degree(row, ...)flowsource_to_target下目标节点按col聚合消息度数应统计目标端deg_inv_sqrt[col]和deg_inv_sqrt[row]分别做什么对称归一化中目标端与源端的归一化因子在message中x_j保存什么若self.lin为恒等函数x_j的具体内容是什么每条边源节点的原始特征给GCNConv添加一个update函数把经变换的中心节点特征加到聚合输出上相当于实现残差连接。围绕EdgeConv的问题x_i和x_j - x_i分别是什么目标节点特征边两端点特征之差torch.cat([x_i, x_j - x_i], dim1)做了什么为什么是dim1按特征维拼接因为每行是一条边、每列是一个特征通道异构图学习为什么需要专门的数据结构大量真实数据集以异构图heterogeneous graph形式存储——推荐领域的社交图就是典型例子图中同时存在多种实体类型用户、商品、商家与多种关系类型。这类图上不同节点/边类型携带不同维度、不同类型的特征单一的特征张量无法容纳整张图的信息因此需要按类型分别维护数据张量相应地消息传递公式也需允许消息函数与更新函数按节点/边类型条件化。引导示例ogbn-mag 网络教程以 OGB 数据集套件中的ogbn-mag网络为引导示例示意图见 hg_example.svg该图共有1,939,743 个节点分为四种节点类型author作者、paper论文、institution机构与field of study研究领域共有21,111,007 条边也分四种类型writes作者撰写某篇论文affiliated with作者隶属于某机构cites论文引用另一篇论文has topic论文属于某个研究领域。任务目标是根据图中存储的信息推断每篇论文的发表 venue会议或期刊。创建异构图HeteroData 详解首先创建torch_geometric.data.HeteroData对象按类型分别定义节点特征张量、边索引张量与边特征张量from torch_geometric.data import HeteroData data HeteroData() data[paper].x ... # [num_papers, num_features_paper] data[author].x ... # [num_authors, num_features_author] data[institution].x ... # [num_institutions, num_features_institution] data[field_of_study].x ... # [num_field, num_features_field] data[paper, cites, paper].edge_index ... # [2, num_edges_cites] data[author, writes, paper].edge_index ... # [2, num_edges_writes] data[author, affiliated_with, institution].edge_index ... # [2, num_edges_affiliated] data[paper, has_topic, field_of_study].edge_index ... # [2, num_edges_topic] data[paper, cites, paper].edge_attr ... # [num_edges_cites, num_features_cites] data[author, writes, paper].edge_attr ... # [num_edges_writes, num_features_writes] data[author, affiliated_with, institution].edge_attr ... # [num_edges_affiliated, num_features_affiliated] data[paper, has_topic, field_of_study].edge_attr ... # [num_edges_topic, num_features_topic]节点/边张量在首次访问时自动创建并以字符串键索引。节点类型用单个字符串标识边类型用三元组(source_node_type, edge_type, destination_node_type)标识。因此该数据对象天然允许每种类型拥有不同的特征维度。按属性名而非按节点/边类型分组后的异构字典可直接作为 GNN 模型的输入model HeteroGNN(...) output model(data.x_dict, data.edge_index_dict, data.edge_attr_dict)直接加载 OGB_MAG 数据集若数据集在 PyG 的数据集列表中可直接导入使用——数据集会被自动下载到root并完成预处理from torch_geometric.datasets import OGB_MAG dataset OGB_MAG(root./data, preprocessmetapath2vec) data dataset[0]打印该data对象验证结构HeteroData( paper{ x[736389, 128], y[736389], train_mask[736389], val_mask[736389], test_mask[736389] }, author{ x[1134649, 128] }, institution{ x[8740, 128] }, field_of_study{ x[59965, 128] }, (author, affiliated_with, institution){ edge_index[2, 1043998] }, (author, writes, paper){ edge_index[2, 7145660] }, (paper, cites, paper){ edge_index[2, 5416271] }, (paper, has_topic, field_of_study){ edge_index[2, 7505078] } )注意原始ogbn-mag网络只给 paper 节点提供特征。PyG 的OGB_MAG提供了下载已处理版本的选项——用metapath2vec或TransE得到的结构特征填充到无特征节点上这正是 OGB 排行榜顶部方案常用的做法见 datasets/ogb_mag.py。实用工具函数HeteroData提供多种修改与分析图的工具# 单独索引某个节点或边存储 paper_node_data data[paper] cites_edge_data data[paper, cites, paper] # 若边类型可由节点对或边类型唯一确定可简写 cites_edge_data data[paper, paper] cites_edge_data data[cites] # 新增/删除节点类型或张量 data[paper].year ... # Setting a new paper attribute del data[field_of_study] # Deleting field_of_study node type del data[has_topic] # Deleting has_topic edge type # 查看元数据所有节点/边类型 node_types, edge_types data.metadata() print(node_types) [paper, author, institution] print(edge_types) [(paper, cites, paper), (author, writes, paper), (author, affiliated_with, institution)] # 设备迁移 data data.to(cuda:0) data data.cpu() # 图性质分析 data.has_isolated_nodes() data.has_self_loops() data.is_undirected()还可通过to_homogeneous()转换为同构带类型图当各类型特征维度一致时可保留特征homogeneous_data data.to_homogeneous() print(homogeneous_data) Data(x[1879778, 128], edge_index[2, 13605929], edge_type[13605929])其中homogeneous_data.edge_type是一个边级向量以整数记录每条边的边类型。异构图变换大多数预处理变换同样适用于异构data对象import torch_geometric.transforms as T data T.ToUndirected()(data) data T.AddSelfLoops()(data) data T.NormalizeFeatures()(data)ToUndirected把有向图转换为PyG 表示下的无向图——为所有边添加反向边必要时会为异构图添加反向边类型使后续消息传递沿两个方向进行AddSelfLoops对类型为node_type的所有节点、以及所有形如(node_type, edge_type, node_type)的现有边类型添加自环结果是每个节点可能收到一个或多个每个合适的边类型一个来自自身的信息NormalizeFeatures与同构图一致把所有类型的指定特征归一化到和为 1。创建异构 GNN 的三种方式标准消息传递 GNNMP-GNN无法直接应用于异构数据不同类型的节点/边特征不能用同一函数处理维度、语义不同。自然的思路是为每种边类型单独实现消息函数、每种节点类型单独实现更新函数——运行时按边类型字典迭代计算消息、按节点类型字典更新节点。为避免不必要的运行时开销并简化建模PyG 提供了三种构建异构 GNN 模型的方式自动转换用torch_geometric.nn.to_hetero或to_hetero_with_bases把同构 GNN 自动转换为异构 GNNHeteroConv 包装器用torch_geometric.nn.conv.HeteroConv为不同边类型定义各自的卷积直接部署现成或自研异构算子如torch_geometric.nn.conv.HGTConv。方式一自动转换同构 GNNto_hetero / to_hetero_with_basesPyG 内置to_hetero与to_hetero_with_bases函数实现见 to_hetero_transformer.py可把任意 PyG GNN 模型自动转换为异构输入模型。以 to_hetero_mag.py 为例import torch_geometric.transforms as T from torch_geometric.datasets import OGB_MAG from torch_geometric.nn import SAGEConv, to_hetero dataset OGB_MAG(root./data, preprocessmetapath2vec, transformT.ToUndirected()) data dataset[0] class GNN(torch.nn.Module): def __init__(self, hidden_channels, out_channels): super().__init__() self.conv1 SAGEConv((-1, -1), hidden_channels) self.conv2 SAGEConv((-1, -1), out_channels) def forward(self, x, edge_index): x self.conv1(x, edge_index).relu() x self.conv2(x, edge_index) return x model GNN(hidden_channels64, out_channelsdataset.num_classes) model to_hetero(model, data.metadata(), aggrsum)转换过程会复制消息函数使其按每种边类型独立工作如下图所示来源 to_hetero.svg。转换后模型期望的输入从同构图中的单一张量变为以节点/边类型为键的字典。注意SAGEConv传入的是(-1, -1)形状的in_channels元组——这是为了支持二部图异构边类型的消息传递。惰性初始化lazy initialization由于不同类型间的输入特征数与张量尺寸各不相同PyG 用-1作为in_channels表示惰性初始化避免手动计算计算图中所有张量尺寸惰性初始化对所有 PyG 算子生效。只需调用一次模型即可完成参数初始化with torch.no_grad(): # Initialize lazy modules. out model(data.x_dict, data.edge_index_dict)to_hetero/to_hetero_with_bases对可自动转换的同构架构非常灵活跳连skip-connection、Jumping Knowledge 等技术开箱即用。例如实现带可学习跳连的异构图注意力网络只需from torch_geometric.nn import GATConv, Linear, to_hetero class GAT(torch.nn.Module): def __init__(self, hidden_channels, out_channels): super().__init__() self.conv1 GATConv((-1, -1), hidden_channels, add_self_loopsFalse) self.lin1 Linear(-1, hidden_channels) self.conv2 GATConv((-1, -1), out_channels, add_self_loopsFalse) self.lin2 Linear(-1, out_channels) def forward(self, x, edge_index): x self.conv1(x, edge_index) self.lin1(x) x x.relu() x self.conv2(x, edge_index) self.lin2(x) return x model GAT(hidden_channels64, out_channelsdataset.num_classes) model to_hetero(model, data.metadata(), aggrsum)这里特意用add_self_loopsFalse关闭自环在二部图中自环概念不成立边类型两端节点类型不同若开启会错误地在二部图上添加[(0, 0), (1, 1), ...]这类边。为保留中心节点信息改用可学习跳连conv(x, edge_index) lin(x)注意力消息从源节点传到目标节点输出再与既有目标节点特征相加。转换后的模型按标准流程训练def train(): model.train() optimizer.zero_grad() out model(data.x_dict, data.edge_index_dict) mask data[paper].train_mask loss F.cross_entropy(out[paper][mask], data[paper].y[mask]) loss.backward() optimizer.step() return float(loss)方式二使用异构卷积包装器 HeteroConvHeteroConvhetero_conv.py允许从零为每种边类型定义自定义消息与更新函数构建任意异构 MP-GNN。与to_hetero所有边类型共用同一算子不同包装器为不同边类型指定不同算子它接收一个以边类型为键的子模块字典。以 hetero_conv_dblp.py 为例import torch_geometric.transforms as T from torch_geometric.datasets import OGB_MAG from torch_geometric.nn import HeteroConv, GCNConv, SAGEConv, GATConv, Linear dataset OGB_MAG(root./data, preprocessmetapath2vec, transformT.ToUndirected()) data dataset[0] class HeteroGNN(torch.nn.Module): def __init__(self, hidden_channels, out_channels, num_layers): super().__init__() self.convs torch.nn.ModuleList() for _ in range(num_layers): conv HeteroConv({ (paper, cites, paper): GCNConv(-1, hidden_channels), (author, writes, paper): SAGEConv((-1, -1), hidden_channels), (paper, rev_writes, author): GATConv((-1, -1), hidden_channels, add_self_loopsFalse), }, aggrsum) self.convs.append(conv) self.lin Linear(hidden_channels, out_channels) def forward(self, x_dict, edge_index_dict): for conv in self.convs: x_dict conv(x_dict, edge_index_dict) x_dict {key: x.relu() for key, x in x_dict.items()} return self.lin(x_dict[author]) model HeteroGNN(hidden_channels64, out_channelsdataset.num_classes, num_layers2)初始化与训练方式同前with torch.no_grad(): # Initialize lazy modules. out model(data.x_dict, data.edge_index_dict)从源码看HeteroConv内部会为字典中每个边类型检查add_self_loops设置并警告那些只作为源类型出现、从不作为目标类型更新表示的节点类型hetero_conv.py L70-L80指向同一目标节点的多种关系的结果会按aggr聚合sum/mean/min/max/cat/None见 L13-L26。方式三部署现成的异构算子如 HGTConvPyG 提供专为异构图设计的算子例如torch_geometric.nn.conv.HGTConv可直接用于搭建异构 GNN 模型参考 hgt_dblp.pyimport torch_geometric.transforms as T from torch_geometric.datasets import OGB_MAG from torch_geometric.nn import HGTConv, Linear dataset OGB_MAG(root./data, preprocessmetapath2vec, transformT.ToUndirected()) data dataset[0] class HGT(torch.nn.Module): def __init__(self, hidden_channels, out_channels, num_heads, num_layers): super().__init__() self.lin_dict torch.nn.ModuleDict() for node_type in data.node_types: self.lin_dict[node_type] Linear(-1, hidden_channels) self.convs torch.nn.ModuleList() for _ in range(num_layers): conv HGTConv(hidden_channels, hidden_channels, data.metadata(), num_heads, groupsum) self.convs.append(conv) self.lin Linear(hidden_channels, out_channels) def forward(self, x_dict, edge_index_dict): for node_type, x in x_dict.items(): x_dict[node_type] self.lin_dictnode_type.relu_() for conv in self.convs: x_dict conv(x_dict, edge_index_dict) return self.lin(x_dict[author]) model HGT(hidden_channels64, out_channelsdataset.num_classes, num_heads2, num_layers2)同样地初始化与训练流程与前面完全一致先with torch.no_grad()调用一次触发惰性初始化再按标准训练函数迭代。大规模异构图采样器与 mini-batch 训练PyG 为异构图采样提供了多种功能标准torch_geometric.loader.NeighborLoader同时支持同构与异构图也有专用异构采样器如torch_geometric.loader.HGTLoader。这对大规模异构图的高效表示学习尤为重要——全量处理邻居在计算上不可承受。所有异构图加载器输出一个HeteroData对象原数据的子集差异主要在采样过程因此从全批量训练切换到 mini-batch 训练只需极少的代码改动。用NeighborLoader做邻居采样的示例同样参考 to_hetero_mag.pyimport torch_geometric.transforms as T from torch_geometric.datasets import OGB_MAG from torch_geometric.loader import NeighborLoader transform T.ToUndirected() # Add reverse edge types. data OGB_MAG(root./data, preprocessmetapath2vec, transformtransform)[0] train_loader NeighborLoader( data, # Sample 15 neighbors for each node and each edge type for 2 iterations: num_neighbors[15] * 2, # Use a batch size of 128 for sampling training nodes of type paper: batch_size128, input_nodes(paper, data[paper].train_mask), ) batch next(iter(train_loader))NeighborLoader在异构图上还可按边类型做更细粒度的采样控制非必需num_neighbors {key: [15] * 2 for key in data.edge_types}input_nodes参数指定采样的局部邻域起点类型与索引——这里即train_mask标记的全部 paper 训练节点。打印batch得到HeteroData( paper{ x[20799, 256], y[20799], train_mask[20799], val_mask[20799], test_mask[20799], batch_size128 }, author{ x[4419, 128] }, institution{ x[302, 128] }, field_of_study{ x[2605, 128] }, (author, affiliated_with, institution){ edge_index[2, 0] }, (author, writes, paper){ edge_index[2, 5927] }, (paper, cites, paper){ edge_index[2, 11829] }, (paper, has_topic, field_of_study){ edge_index[2, 10573] }, (institution, rev_affiliated_with, author){ edge_index[2, 829] }, (paper, rev_writes, author){ edge_index[2, 5512] }, (field_of_study, rev_has_topic, paper){ edge_index[2, 10499] } )batch共包含 28,187 个节点用于计算 128 个 paper 节点的嵌入。采样节点总是按采样顺序排序因此batch[paper]中前batch[paper].batch_size个节点即原始 mini-batch 节点集合通过切片即可方便地取出最终输出嵌入。mini-batch 训练与全批量训练类似区别在于遍历train_loader产生的 mini-batch 并逐个优化def train(): model.train() total_examples total_loss 0 for batch in train_loader: optimizer.zero_grad() batch batch.to(cuda:0) batch_size batch[paper].batch_size out model(batch.x_dict, batch.edge_index_dict) loss F.cross_entropy(out[paper][:batch_size], batch[paper].y[:batch_size]) loss.backward() optimizer.step() total_examples batch_size total_loss float(loss) * batch_size return total_loss / total_examples关键点损失计算只使用前 128 个 paper 节点——通过batch[paper].batch_size对标签batch[paper].y与预测out[paper]同时切片二者分别对应原始 mini-batch 节点的标签与最终输出。进一步阅读本指南对应的原始教程gnn_design.rstgallery 入口、create_gnn.rst、heterogeneous.rst基类与算子源码message_passing.py、gcn_conv.py、edge_conv.py、hetero_conv.py、to_hetero_transformer.py完整可运行示例to_hetero_mag.py、hetero_conv_dblp.py、hgt_dblp.py、bipartite_sage.py异构数据与数据集hetero_data.py、ogb_mag.py异构图测试覆盖test/nn/conv、test/data/test_hetero_data.py掌握MessagePassing基类与三种异构建模方案后你既可以按需定制任意图算子也能把成熟同构模型快速迁移到真实的异构大规模图数据上是进阶 PyG 深度定制与工业级落地的关键一步。【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 PyG 图神经网络设计指南：从 MessagePassing 基类到异构图学习实战
PyG 图神经网络设计指南从 MessagePassing 基类到异构图学习实战【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric本指南以 PyTorch GeometricPyG官方教程《Design of Graph Neural Networks》为主线系统讲解如何利用MessagePassing基类从零实现 GCN、EdgeConv 等消息传递网络并进一步掌握异构图Heterogeneous Graph的建模、变换与三类异构 GNN 构建方案。读完本文你将能够自定义任意消息传递算子并把同构图模型一键迁移到 ogbn-mag 等真实异构数据集上完成训练。消息传递范式GNN 的统一数学框架把卷积算子推广到不规则域如图、点云上通常被表达为**邻域聚合neighborhood aggregation或消息传递message passing**范式。设 $\mathbf{x}^{(k-1)}i \in \mathbb{R}^F$ 表示节点 $i$ 在第 $(k-1)$ 层的特征$\mathbf{e}{j,i} \in \mathbb{R}^D$ 表示从节点 $j$ 指向节点 $i$ 的可选边特征则消息传递图神经网络可以写成$$ \mathbf{x}_i^{(k)} \gamma^{(k)} \left( \mathbf{x}i^{(k-1)}, \bigoplus{j \in \mathcal{N}(i)} , \phi^{(k)}\left(\mathbf{x}_i^{(k-1)}, \mathbf{x}j^{(k-1)},\mathbf{e}{j,i}\right) \right) $$其中 $\bigoplus$ 表示一个可微的、置换不变的聚合函数例如 sum、mean 或 max$\gamma$ 与 $\phi$ 表示可微函数例如 MLP多层感知机。直观地说每个节点先从邻居收集消息由 $\phi$ 生成再按置换不变的方式聚合由 $\bigoplus$ 完成最后结合自身信息更新表示由 $\gamma$ 完成。几乎所有主流 GNN 算子——GCN、GraphSAGE、GAT、GIN、EdgeConv——都是该公式的特例这正是 PyG 设计MessagePassing基类的理论根基。MessagePassing 基类PyG 提供的消息传递脚手架PyG 在 message_passing.py 中提供了MessagePassing基类它自动处理消息的传播propagation流程用户只需定义三个要素$\phi$即message()函数$\gamma$即update()函数聚合方案即aggradd、aggrmean或aggrmax。构造函数参数MessagePassing的构造签名见 源码 L110-L118如下参数说明默认值aggr聚合方案add/sum、mean、min、max、mul也可以是Aggregation模块或字符串列表列表时各聚合结果在最后一维拼接sum与add等价aggr_kwargs传递给自动解析出的聚合函数的额外参数Noneflow消息传递方向source_to_target或target_to_sourcesource_to_targetnode_dim沿哪个轴传播消息批量维度处理-2decomposed_layers特征分解层数用于降低峰值内存、加速 CPU 推理1注意源码中flow只接受source_to_target与target_to_source两个取值传入其他值会抛出ValueError见 L121-L123。decomposed_layers通过把特征维切片成多个子层分批做聚合可在 CPU 上显著降低峰值内存但不适用于注意力类 GNN消息无法简单分解。四个核心方法propagate(edge_index, sizeNone, **kwargs)启动消息传播的入口源码 L421。它接收边索引以及构造消息、更新节点嵌入所需的全部附加数据。关键特性propagate不只限于形状为 $[N, N]$ 的方阵邻接矩阵也可以处理形状为 $[N, M]$ 的一般稀疏赋值矩阵即二部图只需传入size(N, M)若为None则默认是方阵。对于具有两类独立节点的二部图若两类节点各自持有信息可用元组传入例如x(x_N, x_M)。message(...)对应于 $\phi$为每条边 $(j,i) \in \mathcal{E}$当flowsource_to_target时构造发往节点 $i$ 的消息。它可以接收任何传给propagate的参数更妙的是把变量名加上_i或_j后缀PyG 会自动把张量映射到对应的目标/源节点例如x_i、x_j。这里约定 $i$ 是聚合信息的中心节点$j$ 是邻居节点。aggregate(...)执行 $\bigoplus$ 聚合源码 L577。默认委托给__init__中由aggr解析出的Aggregation模块。update(aggr_out, ...)对应于 $\gamma$对每个节点 $i \in \mathcal{V}$ 更新节点嵌入源码 L609。第一个参数是聚合输出其余参数来自propagate的输入。propagate内部会依次调用message→aggregate→update见 源码 L499-L550。其底层实现中有两个值得关注的机制Inspector 参数收集构造时用Inspector反射message/aggregate/update的签名自动确定需要从propagate的**kwargs中收集哪些参数L138-L151。张量抬升lifting_lift通过index_select(self.node_dim, edge_index[dim])把节点特征按边索引展开成边级张量这就是x_j的来源_collect则根据flow决定_i/_j分别取edge_index的第 1/0 行。只要张量持有源或目标节点特征任何张量都可以用_i/_j后缀自动抬升。下面通过重实现两个经典算子——GCN 层Kipf Welling与 EdgeConv 层Wang et al.——来验证这套机制。实战一从零实现 GCN 层GCN 层的数学定义为$$ \mathbf{x}i^{(k)} \sum{j \in \mathcal{N}(i) \cup { i }} \frac{1}{\sqrt{\deg(i)} \cdot \sqrt{\deg(j)}} \cdot \left( \mathbf{W}^{\top} \cdot \mathbf{x}_j^{(k-1)} \right) \mathbf{b} $$即邻居节点特征先经权重矩阵 $\mathbf{W}$ 线性变换再按度数归一化最后求和对聚合输出施加偏置向量 $\mathbf{b}$。这个公式可拆解为六个步骤给邻接矩阵添加自环self-loops线性变换节点特征矩阵计算归一化系数在 $\phi$message中归一化节点特征求和聚合邻居节点特征add聚合施加最终偏置向量。步骤 1–3 通常在消息传递前完成步骤 4–5 可直接交给MessagePassing基类。完整实现如下与教程 create_gnn.rst 一致import torch from torch.nn import Linear, Parameter from torch_geometric.nn import MessagePassing from torch_geometric.utils import add_self_loops, degree class GCNConv(MessagePassing): def __init__(self, in_channels, out_channels): super().__init__(aggradd) # Add aggregation (Step 5). self.lin Linear(in_channels, out_channels, biasFalse) self.bias Parameter(torch.empty(out_channels)) self.reset_parameters() def reset_parameters(self): self.lin.reset_parameters() self.bias.data.zero_() def forward(self, x, edge_index): # x has shape [N, in_channels] # edge_index has shape [2, E] # Step 1: Add self-loops to the adjacency matrix. edge_index, _ add_self_loops(edge_index, num_nodesx.size(0)) # Step 2: Linearly transform node feature matrix. x self.lin(x) # Step 3: Compute normalization. row, col edge_index deg degree(col, x.size(0), dtypex.dtype) deg_inv_sqrt deg.pow(-0.5) deg_inv_sqrt[deg_inv_sqrt float(inf)] 0 norm deg_inv_sqrt[row] * deg_inv_sqrt[col] # Step 4-5: Start propagating messages. out self.propagate(edge_index, xx, normnorm) # Step 6: Apply a final bias vector. out out self.bias return out def message(self, x_j, norm): # x_j has shape [E, out_channels] # Step 4: Normalize node features. return norm.view(-1, 1) * x_j逐段解读super().__init__(aggradd)选择了求和聚合步骤 5。PyG 内部把add与sum视为同一聚合见FUSE_AGGRS {add, sum, mean, min, max}message_passing.py L35。add_self_loops来自 utils/loop.py为edge_index补充 $(i,i)$ 自环步骤 1。PyG 官方实现的GCNConvgcn_conv.py更进一步使用add_remaining_self_loops避免重复添加并通过gcn_normgcn_conv.py L45-L113在稀疏矩阵与稠密张量两种表示下统一完成自环 对称归一化。degree(col, ...)来自 utils/_degree.py计算每个节点的度数deg_inv_sqrt中无穷大被置 0 以处理孤立节点步骤 3。propagate(edge_index, xx, normnorm)内部依次调用message、aggregate、update。在message(self, x_j, norm)中x_j是抬升lifted张量包含每条边源节点的特征即各节点的邻居特征norm同样按边索引展开后与x_j逐元素相乘完成归一化步骤 4。初始化与调用极其简单该层可直接作为深度架构的积木conv GCNConv(16, 32) x conv(x, edge_index)实战二实现 Edge 卷积层EdgeConv 层用于处理图或点云数学定义为$$ \mathbf{x}i^{(k)} \max{j \in \mathcal{N}(i)} h_{\mathbf{\Theta}} \left( \mathbf{x}_i^{(k-1)}, \mathbf{x}_j^{(k-1)} - \mathbf{x}_i^{(k-1)} \right) $$其中 $h_{\mathbf{\Theta}}$ 是一个 MLP。与 GCN 类似这次改用max聚合import torch from torch.nn import Sequential as Seq, Linear, ReLU from torch_geometric.nn import MessagePassing class EdgeConv(MessagePassing): def __init__(self, in_channels, out_channels): super().__init__(aggrmax) # Max aggregation. self.mlp Seq(Linear(2 * in_channels, out_channels), ReLU(), Linear(out_channels, out_channels)) def forward(self, x, edge_index): # x has shape [N, in_channels] # edge_index has shape [2, E] return self.propagate(edge_index, xx) def message(self, x_i, x_j): # x_i has shape [E, in_channels] # x_j has shape [E, in_channels] tmp torch.cat([x_i, x_j - x_i], dim1) # tmp has shape [E, 2 * in_channels] return self.mlp(tmp)在message中我们同时用到了x_i与x_jx_i是目标节点特征x_j - x_i是相对源节点特征刻画了每条边两端点的差分信息二者在特征维拼接后送入 MLP。官方 edge_conv.py 中的EdgeConv实现与此一致L60-L61差别仅在于官方把 MLP 作为nn参数从外部注入、message用dim-1拼接且forward支持二部图输入x (x, x)。EdgeConv 本质上是动态卷积每一层都在特征空间中用最近邻重新构造图。PyG 提供了 GPU 加速的批量 k-NN 图生成方法torch_geometric.nn.pool.knn_graphfrom torch_geometric.nn import knn_graph class DynamicEdgeConv(EdgeConv): def __init__(self, in_channels, out_channels, k6): super().__init__(in_channels, out_channels) self.k k def forward(self, x, batchNone): edge_index knn_graph(x, self.k, batch, loopFalse, flowself.flow) return super().forward(x, edge_index)knn_graph计算最近邻图再调用EdgeConv.forward完成消息传递。官方 DynamicEdgeConv 还额外要求pyg-lib0.6.0支持并支持num_workers并行计算 k-NN。调用接口干净利落conv DynamicEdgeConv(3, 128, k6) x conv(x, batch)动手练习验证你的理解教程 create_gnn.rst 提供了一个练习数据集import torch from torch_geometric.data import Data edge_index torch.tensor([[0, 1], [1, 0], [1, 2], [2, 1]], dtypetorch.long) x torch.tensor([[-1], [0], [1]], dtypetorch.float) data Data(xx, edge_indexedge_index.t().contiguous())围绕GCNConv思考以下问题答案都可从源码与上述解读推出row和col分别保存什么信息row edge_index[0]为源节点col edge_index[1]为目标节点degree函数做了什么统计每个节点的度数即入边数量为什么用degree(col, ...)而不是degree(row, ...)flowsource_to_target下目标节点按col聚合消息度数应统计目标端deg_inv_sqrt[col]和deg_inv_sqrt[row]分别做什么对称归一化中目标端与源端的归一化因子在message中x_j保存什么若self.lin为恒等函数x_j的具体内容是什么每条边源节点的原始特征给GCNConv添加一个update函数把经变换的中心节点特征加到聚合输出上相当于实现残差连接。围绕EdgeConv的问题x_i和x_j - x_i分别是什么目标节点特征边两端点特征之差torch.cat([x_i, x_j - x_i], dim1)做了什么为什么是dim1按特征维拼接因为每行是一条边、每列是一个特征通道异构图学习为什么需要专门的数据结构大量真实数据集以异构图heterogeneous graph形式存储——推荐领域的社交图就是典型例子图中同时存在多种实体类型用户、商品、商家与多种关系类型。这类图上不同节点/边类型携带不同维度、不同类型的特征单一的特征张量无法容纳整张图的信息因此需要按类型分别维护数据张量相应地消息传递公式也需允许消息函数与更新函数按节点/边类型条件化。引导示例ogbn-mag 网络教程以 OGB 数据集套件中的ogbn-mag网络为引导示例示意图见 hg_example.svg该图共有1,939,743 个节点分为四种节点类型author作者、paper论文、institution机构与field of study研究领域共有21,111,007 条边也分四种类型writes作者撰写某篇论文affiliated with作者隶属于某机构cites论文引用另一篇论文has topic论文属于某个研究领域。任务目标是根据图中存储的信息推断每篇论文的发表 venue会议或期刊。创建异构图HeteroData 详解首先创建torch_geometric.data.HeteroData对象按类型分别定义节点特征张量、边索引张量与边特征张量from torch_geometric.data import HeteroData data HeteroData() data[paper].x ... # [num_papers, num_features_paper] data[author].x ... # [num_authors, num_features_author] data[institution].x ... # [num_institutions, num_features_institution] data[field_of_study].x ... # [num_field, num_features_field] data[paper, cites, paper].edge_index ... # [2, num_edges_cites] data[author, writes, paper].edge_index ... # [2, num_edges_writes] data[author, affiliated_with, institution].edge_index ... # [2, num_edges_affiliated] data[paper, has_topic, field_of_study].edge_index ... # [2, num_edges_topic] data[paper, cites, paper].edge_attr ... # [num_edges_cites, num_features_cites] data[author, writes, paper].edge_attr ... # [num_edges_writes, num_features_writes] data[author, affiliated_with, institution].edge_attr ... # [num_edges_affiliated, num_features_affiliated] data[paper, has_topic, field_of_study].edge_attr ... # [num_edges_topic, num_features_topic]节点/边张量在首次访问时自动创建并以字符串键索引。节点类型用单个字符串标识边类型用三元组(source_node_type, edge_type, destination_node_type)标识。因此该数据对象天然允许每种类型拥有不同的特征维度。按属性名而非按节点/边类型分组后的异构字典可直接作为 GNN 模型的输入model HeteroGNN(...) output model(data.x_dict, data.edge_index_dict, data.edge_attr_dict)直接加载 OGB_MAG 数据集若数据集在 PyG 的数据集列表中可直接导入使用——数据集会被自动下载到root并完成预处理from torch_geometric.datasets import OGB_MAG dataset OGB_MAG(root./data, preprocessmetapath2vec) data dataset[0]打印该data对象验证结构HeteroData( paper{ x[736389, 128], y[736389], train_mask[736389], val_mask[736389], test_mask[736389] }, author{ x[1134649, 128] }, institution{ x[8740, 128] }, field_of_study{ x[59965, 128] }, (author, affiliated_with, institution){ edge_index[2, 1043998] }, (author, writes, paper){ edge_index[2, 7145660] }, (paper, cites, paper){ edge_index[2, 5416271] }, (paper, has_topic, field_of_study){ edge_index[2, 7505078] } )注意原始ogbn-mag网络只给 paper 节点提供特征。PyG 的OGB_MAG提供了下载已处理版本的选项——用metapath2vec或TransE得到的结构特征填充到无特征节点上这正是 OGB 排行榜顶部方案常用的做法见 datasets/ogb_mag.py。实用工具函数HeteroData提供多种修改与分析图的工具# 单独索引某个节点或边存储 paper_node_data data[paper] cites_edge_data data[paper, cites, paper] # 若边类型可由节点对或边类型唯一确定可简写 cites_edge_data data[paper, paper] cites_edge_data data[cites] # 新增/删除节点类型或张量 data[paper].year ... # Setting a new paper attribute del data[field_of_study] # Deleting field_of_study node type del data[has_topic] # Deleting has_topic edge type # 查看元数据所有节点/边类型 node_types, edge_types data.metadata() print(node_types) [paper, author, institution] print(edge_types) [(paper, cites, paper), (author, writes, paper), (author, affiliated_with, institution)] # 设备迁移 data data.to(cuda:0) data data.cpu() # 图性质分析 data.has_isolated_nodes() data.has_self_loops() data.is_undirected()还可通过to_homogeneous()转换为同构带类型图当各类型特征维度一致时可保留特征homogeneous_data data.to_homogeneous() print(homogeneous_data) Data(x[1879778, 128], edge_index[2, 13605929], edge_type[13605929])其中homogeneous_data.edge_type是一个边级向量以整数记录每条边的边类型。异构图变换大多数预处理变换同样适用于异构data对象import torch_geometric.transforms as T data T.ToUndirected()(data) data T.AddSelfLoops()(data) data T.NormalizeFeatures()(data)ToUndirected把有向图转换为PyG 表示下的无向图——为所有边添加反向边必要时会为异构图添加反向边类型使后续消息传递沿两个方向进行AddSelfLoops对类型为node_type的所有节点、以及所有形如(node_type, edge_type, node_type)的现有边类型添加自环结果是每个节点可能收到一个或多个每个合适的边类型一个来自自身的信息NormalizeFeatures与同构图一致把所有类型的指定特征归一化到和为 1。创建异构 GNN 的三种方式标准消息传递 GNNMP-GNN无法直接应用于异构数据不同类型的节点/边特征不能用同一函数处理维度、语义不同。自然的思路是为每种边类型单独实现消息函数、每种节点类型单独实现更新函数——运行时按边类型字典迭代计算消息、按节点类型字典更新节点。为避免不必要的运行时开销并简化建模PyG 提供了三种构建异构 GNN 模型的方式自动转换用torch_geometric.nn.to_hetero或to_hetero_with_bases把同构 GNN 自动转换为异构 GNNHeteroConv 包装器用torch_geometric.nn.conv.HeteroConv为不同边类型定义各自的卷积直接部署现成或自研异构算子如torch_geometric.nn.conv.HGTConv。方式一自动转换同构 GNNto_hetero / to_hetero_with_basesPyG 内置to_hetero与to_hetero_with_bases函数实现见 to_hetero_transformer.py可把任意 PyG GNN 模型自动转换为异构输入模型。以 to_hetero_mag.py 为例import torch_geometric.transforms as T from torch_geometric.datasets import OGB_MAG from torch_geometric.nn import SAGEConv, to_hetero dataset OGB_MAG(root./data, preprocessmetapath2vec, transformT.ToUndirected()) data dataset[0] class GNN(torch.nn.Module): def __init__(self, hidden_channels, out_channels): super().__init__() self.conv1 SAGEConv((-1, -1), hidden_channels) self.conv2 SAGEConv((-1, -1), out_channels) def forward(self, x, edge_index): x self.conv1(x, edge_index).relu() x self.conv2(x, edge_index) return x model GNN(hidden_channels64, out_channelsdataset.num_classes) model to_hetero(model, data.metadata(), aggrsum)转换过程会复制消息函数使其按每种边类型独立工作如下图所示来源 to_hetero.svg。转换后模型期望的输入从同构图中的单一张量变为以节点/边类型为键的字典。注意SAGEConv传入的是(-1, -1)形状的in_channels元组——这是为了支持二部图异构边类型的消息传递。惰性初始化lazy initialization由于不同类型间的输入特征数与张量尺寸各不相同PyG 用-1作为in_channels表示惰性初始化避免手动计算计算图中所有张量尺寸惰性初始化对所有 PyG 算子生效。只需调用一次模型即可完成参数初始化with torch.no_grad(): # Initialize lazy modules. out model(data.x_dict, data.edge_index_dict)to_hetero/to_hetero_with_bases对可自动转换的同构架构非常灵活跳连skip-connection、Jumping Knowledge 等技术开箱即用。例如实现带可学习跳连的异构图注意力网络只需from torch_geometric.nn import GATConv, Linear, to_hetero class GAT(torch.nn.Module): def __init__(self, hidden_channels, out_channels): super().__init__() self.conv1 GATConv((-1, -1), hidden_channels, add_self_loopsFalse) self.lin1 Linear(-1, hidden_channels) self.conv2 GATConv((-1, -1), out_channels, add_self_loopsFalse) self.lin2 Linear(-1, out_channels) def forward(self, x, edge_index): x self.conv1(x, edge_index) self.lin1(x) x x.relu() x self.conv2(x, edge_index) self.lin2(x) return x model GAT(hidden_channels64, out_channelsdataset.num_classes) model to_hetero(model, data.metadata(), aggrsum)这里特意用add_self_loopsFalse关闭自环在二部图中自环概念不成立边类型两端节点类型不同若开启会错误地在二部图上添加[(0, 0), (1, 1), ...]这类边。为保留中心节点信息改用可学习跳连conv(x, edge_index) lin(x)注意力消息从源节点传到目标节点输出再与既有目标节点特征相加。转换后的模型按标准流程训练def train(): model.train() optimizer.zero_grad() out model(data.x_dict, data.edge_index_dict) mask data[paper].train_mask loss F.cross_entropy(out[paper][mask], data[paper].y[mask]) loss.backward() optimizer.step() return float(loss)方式二使用异构卷积包装器 HeteroConvHeteroConvhetero_conv.py允许从零为每种边类型定义自定义消息与更新函数构建任意异构 MP-GNN。与to_hetero所有边类型共用同一算子不同包装器为不同边类型指定不同算子它接收一个以边类型为键的子模块字典。以 hetero_conv_dblp.py 为例import torch_geometric.transforms as T from torch_geometric.datasets import OGB_MAG from torch_geometric.nn import HeteroConv, GCNConv, SAGEConv, GATConv, Linear dataset OGB_MAG(root./data, preprocessmetapath2vec, transformT.ToUndirected()) data dataset[0] class HeteroGNN(torch.nn.Module): def __init__(self, hidden_channels, out_channels, num_layers): super().__init__() self.convs torch.nn.ModuleList() for _ in range(num_layers): conv HeteroConv({ (paper, cites, paper): GCNConv(-1, hidden_channels), (author, writes, paper): SAGEConv((-1, -1), hidden_channels), (paper, rev_writes, author): GATConv((-1, -1), hidden_channels, add_self_loopsFalse), }, aggrsum) self.convs.append(conv) self.lin Linear(hidden_channels, out_channels) def forward(self, x_dict, edge_index_dict): for conv in self.convs: x_dict conv(x_dict, edge_index_dict) x_dict {key: x.relu() for key, x in x_dict.items()} return self.lin(x_dict[author]) model HeteroGNN(hidden_channels64, out_channelsdataset.num_classes, num_layers2)初始化与训练方式同前with torch.no_grad(): # Initialize lazy modules. out model(data.x_dict, data.edge_index_dict)从源码看HeteroConv内部会为字典中每个边类型检查add_self_loops设置并警告那些只作为源类型出现、从不作为目标类型更新表示的节点类型hetero_conv.py L70-L80指向同一目标节点的多种关系的结果会按aggr聚合sum/mean/min/max/cat/None见 L13-L26。方式三部署现成的异构算子如 HGTConvPyG 提供专为异构图设计的算子例如torch_geometric.nn.conv.HGTConv可直接用于搭建异构 GNN 模型参考 hgt_dblp.pyimport torch_geometric.transforms as T from torch_geometric.datasets import OGB_MAG from torch_geometric.nn import HGTConv, Linear dataset OGB_MAG(root./data, preprocessmetapath2vec, transformT.ToUndirected()) data dataset[0] class HGT(torch.nn.Module): def __init__(self, hidden_channels, out_channels, num_heads, num_layers): super().__init__() self.lin_dict torch.nn.ModuleDict() for node_type in data.node_types: self.lin_dict[node_type] Linear(-1, hidden_channels) self.convs torch.nn.ModuleList() for _ in range(num_layers): conv HGTConv(hidden_channels, hidden_channels, data.metadata(), num_heads, groupsum) self.convs.append(conv) self.lin Linear(hidden_channels, out_channels) def forward(self, x_dict, edge_index_dict): for node_type, x in x_dict.items(): x_dict[node_type] self.lin_dictnode_type.relu_() for conv in self.convs: x_dict conv(x_dict, edge_index_dict) return self.lin(x_dict[author]) model HGT(hidden_channels64, out_channelsdataset.num_classes, num_heads2, num_layers2)同样地初始化与训练流程与前面完全一致先with torch.no_grad()调用一次触发惰性初始化再按标准训练函数迭代。大规模异构图采样器与 mini-batch 训练PyG 为异构图采样提供了多种功能标准torch_geometric.loader.NeighborLoader同时支持同构与异构图也有专用异构采样器如torch_geometric.loader.HGTLoader。这对大规模异构图的高效表示学习尤为重要——全量处理邻居在计算上不可承受。所有异构图加载器输出一个HeteroData对象原数据的子集差异主要在采样过程因此从全批量训练切换到 mini-batch 训练只需极少的代码改动。用NeighborLoader做邻居采样的示例同样参考 to_hetero_mag.pyimport torch_geometric.transforms as T from torch_geometric.datasets import OGB_MAG from torch_geometric.loader import NeighborLoader transform T.ToUndirected() # Add reverse edge types. data OGB_MAG(root./data, preprocessmetapath2vec, transformtransform)[0] train_loader NeighborLoader( data, # Sample 15 neighbors for each node and each edge type for 2 iterations: num_neighbors[15] * 2, # Use a batch size of 128 for sampling training nodes of type paper: batch_size128, input_nodes(paper, data[paper].train_mask), ) batch next(iter(train_loader))NeighborLoader在异构图上还可按边类型做更细粒度的采样控制非必需num_neighbors {key: [15] * 2 for key in data.edge_types}input_nodes参数指定采样的局部邻域起点类型与索引——这里即train_mask标记的全部 paper 训练节点。打印batch得到HeteroData( paper{ x[20799, 256], y[20799], train_mask[20799], val_mask[20799], test_mask[20799], batch_size128 }, author{ x[4419, 128] }, institution{ x[302, 128] }, field_of_study{ x[2605, 128] }, (author, affiliated_with, institution){ edge_index[2, 0] }, (author, writes, paper){ edge_index[2, 5927] }, (paper, cites, paper){ edge_index[2, 11829] }, (paper, has_topic, field_of_study){ edge_index[2, 10573] }, (institution, rev_affiliated_with, author){ edge_index[2, 829] }, (paper, rev_writes, author){ edge_index[2, 5512] }, (field_of_study, rev_has_topic, paper){ edge_index[2, 10499] } )batch共包含 28,187 个节点用于计算 128 个 paper 节点的嵌入。采样节点总是按采样顺序排序因此batch[paper]中前batch[paper].batch_size个节点即原始 mini-batch 节点集合通过切片即可方便地取出最终输出嵌入。mini-batch 训练与全批量训练类似区别在于遍历train_loader产生的 mini-batch 并逐个优化def train(): model.train() total_examples total_loss 0 for batch in train_loader: optimizer.zero_grad() batch batch.to(cuda:0) batch_size batch[paper].batch_size out model(batch.x_dict, batch.edge_index_dict) loss F.cross_entropy(out[paper][:batch_size], batch[paper].y[:batch_size]) loss.backward() optimizer.step() total_examples batch_size total_loss float(loss) * batch_size return total_loss / total_examples关键点损失计算只使用前 128 个 paper 节点——通过batch[paper].batch_size对标签batch[paper].y与预测out[paper]同时切片二者分别对应原始 mini-batch 节点的标签与最终输出。进一步阅读本指南对应的原始教程gnn_design.rstgallery 入口、create_gnn.rst、heterogeneous.rst基类与算子源码message_passing.py、gcn_conv.py、edge_conv.py、hetero_conv.py、to_hetero_transformer.py完整可运行示例to_hetero_mag.py、hetero_conv_dblp.py、hgt_dblp.py、bipartite_sage.py异构数据与数据集hetero_data.py、ogb_mag.py异构图测试覆盖test/nn/conv、test/data/test_hetero_data.py掌握MessagePassing基类与三种异构建模方案后你既可以按需定制任意图算子也能把成熟同构模型快速迁移到真实的异构大规模图数据上是进阶 PyG 深度定制与工业级落地的关键一步。【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考