【PYG】简单分析Planetoid()中存储Cora数据集边的数量

最新推荐文章于 2024-07-02 10:07:24 发布

原创

最新推荐文章于 2024-07-02 10:07:24 发布 · 1.3k 阅读

23 ·

CC 4.0 BY-SA版权

文章标签：

#gnn

分析一下代码生成的数据，edge_index=[2, 10556]，10556/2=5278，实际上数据集节点间有5429条边，打印前十条边可以看到(1,2)和(2,1)实际上都存进来了

Data(x=[2708, 1433], edge_index=[2, 10556], y=[2708], train_mask=[2708], val_mask=[2708], test_mask=[2708])
Number of nodes: 2708
Number of edges (directed): 10556
Average node degree: 3.90
Is undirected: True
Edge index shape: torch.Size([2, 10556])
First 10 edges:
tensor([[   0,    0,    0,    1,    1,    1,    2,    2,    2,    2],
        [ 633, 1862, 2582,    2,  652,  654,    1,  332, 1454, 1666]])

分析以下代码输出结果

from torch_geometric.datasets import Planetoid
from torch_geometric.transforms import NormalizeFeatures

# 加载并预处理Cora数据集
dataset = Planetoid(root='/tmp/Cora', name='Cora', transform=NormalizeFeatures())
data = dataset[0]

# 打印边的数量
print(f'Number of nodes: {
     
     data.num_nodes}')
print(f'Number of edges (directed): {
     
     data.num_edges}')

# 检查无向边数量
edge_index = data.edge_index.numpy()
edge_set = set()

print(f"edge_index shape {
     
     edge_index.shape}")

for i in range(edge_index.shape[1]):
    u, v = edge_index[:, i]
    print(f"u {
     
     u}, v {
     
     v}")
    if (v, u) not in edge_set:  # 如果反向边不存在，则添加当前边
        edge_set.add((u, v))
        print(f"add edge_set: {
     
     edge_set}")
    else:
        print(f"not add v {
     
     v}, u {
     
     u}")
    if (i == 9):