【PyG】与networkx的图转换

原创

已于 2022-11-23 11:29:06 修改 · 5k 阅读

43 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #神经网络 #人工智能 #python

于 2022-11-23 11:20:40 首次发布

本文介绍如何在PyTorch Geometric (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']