torch_geometric 入门，加载benchmark数据集

最新推荐文章于 2024-06-30 21:04:13 发布

原创

最新推荐文章于 2024-06-30 21:04:13 发布

· 862 阅读

2 ·

版权

文章标签：

#pytorch #深度学习 #python

来自官网的例子：

from torch.nn import Linear, ReLU
# from torch_geometric.nn import Sequential, GCNConv
import torch

from torch_geometric.datasets import Planetoid, TUDataset
from torch_scatter import scatter_mean
from torch_geometric.loader import DataLoader

dataset1 = Planetoid(root='/home/yingpan/disk1/lzy_data/dataset/',name='Cora')
# 图的数量，图上结点的种类，结点的特征数量
print(len(dataset1),dataset1.num_classes,dataset1.num_node_features)
data = dataset1[0]
# 图的信息，是否是无向图，训练结点的数量
print(data,data.is_undirected(),data.train_mask.sum().item())

dataset2 = TUDataset(root='/home/yingpan/disk1/lzy_data/dataset/',name='ENZYMES',use_node_attr=True)
loader = DataLoader(dataset2,batch_size=32,shuffle=True)

for data in loader:
    # 每个batch的信息，每个batch里面图的数量
    # DataBatch(edge_index=[2, 3912], x=[1063, 21], y=[32], batch=[1063], ptr=[3