PyG-使用 networkx 对 Graph 进行可视化

参考来源：
CSDN：PyG-使用 networkx 对 Graph 进行可视化
知乎：在 PyTorch 框架下使用 PyG 和 networkx 对 Graph 进行可视化
NetworkX 文档：Drawing

方法一

根据 networkx 的文档：networkx.drawing.nx_pylab.draw_networkx
我们可以写出来一个非常简单的例子，如下：

import networkx as nx
import matplotlib.pyplot as plt
G = nx.Graph()
edge_index = [(1, 2), (1, 3), (2, 3), (3, 4)]
G.add_edges_from(edge_index)
nx.draw(G)
plt.show()

运行程序之后，可以得到下面的图，（偷了一个懒，没有加 label 之类的信息）

这个例子给我们的启发就是，我们可以将 PyG 得到的 edge_index 转成 numpy 的格式,然后传给 nx，下面是根据这个写的一个函数：

在 PyG 中,边的表示放在了 edge_index 中，由一个二维的矩阵构成，edge_index[0] 表示节点，edge_index[1] 表示另一个节点。

def draw(edge_index, name=None):
    G = nx.Graph(node_size=15, font_size=8)
    src = edge_index[0].cpu().numpy()
    dst = edge_index[1].cpu().numpy()
    edgelist = zip(src, dst)
    for i, j in edgelist:
        G.add_edge(i, j)
    plt.figure(figsize=(20, 14)) # 设置画布的大小
    nx.draw_networkx(G)
    plt.savefig('{}.png'.format(name if name else 'path'))

注：该方法可以用于模型中的 forward 函数，用于分析 cov、pool 等操作。

下面是与上面思想一致可以直接运行的一个例子

from torch_geometric.datasets import KarateClub
import networkx as nx
import matplotlib.pyplot as plt
dataset = KarateClub()
edge, x, y = dataset[0]
# edge, x, y 每个维度都为2，其中第一维度是name，第二个维度是data
# x表示的是结点，y表示的标签，edge表示的连边, 由两个维度的tensor构成
x_np = x[1].numpy()
y_np = y[1].numpy()
g = nx.Graph()
name, edgeinfo = edge
src = edgeinfo[0].numpy()
dst = edgeinfo[1].numpy()
edgelist = zip(src, dst)
for i, j in edgelist:
    g.add_edge(i, j)
nx.draw(g)
plt.savefig('test.png')
plt.show()

方法二

其实，torch_geometric.utils 中已经带有 to_networkx 的函数可以直接将格式为 torch_geometric.data.Data 的数据转换为 networkx.DiGraph 的格式，该格式可以直接 networkx 处理，但是我们提前要得到 torch_geometric.data.Data 的数据格式。

import networkx as nx
from torch_geometric.utils.convert import to_networkx
def draw(Data):
    G = to_networkx(Data)
    nx.draw(G)
    plt.savefig("path.png")
    plt.show()

这个一般可以用于在 model 加载数据之前数据的分析，比如下面的例子：

for i, data in enumerate(train_loader):
    draw(data)
    data = data.to(args.device)
    out = model(data)
    loss = F.nll_loss(out, data.y)
    print("Training loss:{}".format(loss.item()))
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

上面的函数是在 graph classification 进行分析的一段代码，可以把 batch size 的设置为 1，那么 for 循环中得到就是一个 graph 的数据，在把数据 feed 给模型之前，我们可以通过该方法分析一下原始的数据是什么样子的。