【PyG】与networkx的图转换
在使用图神经网络的过程中,往往需要使用到相关的 GNN 库,而在这些 GNN 库中,一款比较高效热门的图神经网络库是 PyTorch 中的 PyG 库。PyG 提供了很多经典的图神经网络模型和图数据集,通常在使用 PyG 框架来构建和训练图模型时,需要事先选择合适的图数据结构来构造图,PyG 提供的选择包括 Data、HeteroData、TemporalData。
在使用图神经网络的过程中,往往需要使用到相关的 GNN 库,而在这些 GNN 库中,一款比较高效热门的图神经网络库是 PyTorch 中的 PyG 库。PyG 提供了很多经典的图神经网络模型和图数据集,通常在使用 PyG 框架来构建和训练图模型时,需要事先选择合适的图数据结构来构造图,PyG 提供的选择包括 Data、HeteroData、TemporalData。而在实验的过程中,可能需要使用到 networkx 提供的一些功能来实现与图相关的操作,这时图数据需要在两个框架提供的图结构之间进行转换,基于此,本文主要针对转换操作进行了整理和总结。
一、数据准备
本文以简单图为例,同构图与异构图(无向图)如下所示:
1、构建 PyG 同构图
import torch
from torch_geometric.data import Data
data = Data()
# 初始化节点特征
data.x = torch.tensor([[-1], [0], [1]], dtype=torch.float)
# 初始化边索引
data.edge_index = torch.tensor([[0, 1, 1, 2],
[1, 0, 2, 1]], dtype=torch.long)
2、构建 PyG 异构图
import torch
from torch_geometric.data import HeteroData
data = HeteroData()
# 初始化结点特征
# [num_papers, num_features_paper]
data['paper'].x = torch.tensor([[0, 1, 2]], dtype=torch.float)
# [num_authors, num_features_author]
data['author'].x = torch.tensor([[-1], [1]], dtype=torch.float)
# 初始化边索引
# [2, num_edges_writes]
data['author', 'writes', 'paper'].edge_index = torch.tensor([[0, 1],
[0, 0]], dtype=torch.long)
data['paper', 'belongs', 'author'].edge_index = torch.tensor([[0, 0],
[0, 1]], dtype=torch.long)
3、构建 networkx 同构图
import networkx as nx
# 创建无向图
G = nx.Graph()
# 两种添加节点的方式 add_node 和 add_nodes_from
G.add_nodes_from([0, 1, 2])
# 两种添加连边的方式,add_edge 和 add_edges_from
G.add_edges_from([[0, 1], [1, 2]])
4、构建 networkx 异构图
import networkx as nx
# 创建无向图
G = nx.Graph()
# 为节点添加 type 属性(属性名可自定义)来区分节点类型
G.add_nodes_from([0, 2], type='author')
G.add_nodes_from([1], type='paper')
# 为连边添加 type 属性(属性名可自定义)来区分连边类型
G.add_edges_from([[0, 1], [1, 2]], type='writes')
# 获取节点 & 连边类型
node_labels = nx.get_node_attributes(G, 'type')
edge_labels = nx.get_edge_attributes(G, 'type')
二、同构图转换
1、PyG 转 networkx
(1)利用 to_networkx方法直接转换
from torch_geometric.utils.convert import to_networkx
G = to_networkx(data)
- 优点:简单,高效
- 缺点:无法处理规模较大的图(内存不足)
(2)以添加节点与连边的方式转换
import numpy as np
G = nx.Graph()
# 使用 add_nodes_from 批处理的效率比 add_node 高
G.add_nodes_from([i for i in range(data.x.shape[0])])
# 使用 add_edges_from 批处理的效率比 add_edge 高
edges = np.array(data.edge_index.T, dtype=int)
G.add_edges_from(edges)
- 优点:适用于规模较大的图
- 缺点:较为复杂
2、networkx 转 PyG
import torch
import numpy as np
# 创建节点特征矩阵
x = torch.ones((G.number_of_nodes(),1), dtype=torch.float)
# 获取图G邻接矩阵的稀疏表示
adj = nx.to_scipy_sparse_array(G).tocoo()
# 获取非零元素行索引
row = torch.from_numpy(adj.row.astype(np.int64)).to(torch.long)
# 获取非零元素列索引
col = torch.from_numpy(adj.col.astype(np.int64)).to(torch.long)
# 将行和列进行拼接,shape变为[2, num_edges], 包含两个列表,第一个是row, 第二个是col
edge_index = torch.stack([row, col], dim=0)
data = Data(x=x, edge_index=edge_index)
三、异构图转换
1、PyG 转 networkx
(1)利用 to_networkx方法直接转换
from torch_geometric.utils.convert import to_networkx
data = data.to_homogeneous()
G = to_networkx(data)
- 优点:简单,高效
- 缺点:无法处理规模较大的图(内存不足)
(2)以添加节点与连边的方式转换
import numpy as np
G = nx.Graph()
# 需要为节点重新排序
node_num = 0
nt_start = {}
for nt in data.node_types:
nt_start[nt] = node_num
node_num += data[nt].x.shape[0]
# 使用 add_nodes_from 批处理的效率比 add_node 高
for nt in data.node_types:
G.add_nodes_from([nt_start[nt] + i for i in range(data[nt].x.shape[0])], node_type=nt)
# 使用 add_edges_from 批处理的效率比 add_edge 高
for et in data.edge_types:
edges = np.array(data[et].edge_index.T, dtype=int)
G.add_edges_from([[nt_start[et[0]] + e[0], nt_start[et[2]] + e[0]] for e in edges],
edge_type=et[1])
- 优点:适用于规模较大的图
- 缺点:较为复杂
2、networkx 转 PyG
利用 networkx 框架将异构图转 PyG 图结构的情况一般不常见,通常是在 PyG 中创建了图,但为了绘制图结构,才需要转换为 networkx 框架下的图,再利用 networkx 提供的接口进行绘制。

GitCode 天启AI是一款由 GitCode 团队打造的智能助手,基于先进的LLM(大语言模型)与多智能体 Agent 技术构建,致力于为用户提供高效、智能、多模态的创作与开发支持。它不仅支持自然语言对话,还具备处理文件、生成 PPT、撰写分析报告、开发 Web 应用等多项能力,真正做到“一句话,让 Al帮你完成复杂任务”。
更多推荐
所有评论(0)