【PyG】与networkx的图转换

在使用图神经网络的过程中，往往需要使用到相关的 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 提供的接口进行绘制。