基于pyg下载obg数据集

1. 安装pyg

pip install torch
pip install torch-geometric

2. 下载obg数据集

import os
import urllib.request

url = 'https://github.com/graphdeeplearning/benchmarking-gnns/raw/master/data/ogbn-proteins.zip'
filename = 'ogbn-proteins.zip'
urllib.request.urlretrieve(url, filename)

os.system('unzip ogbn-proteins.zip -d data')

3. 加载obg数据集

from ogb.graphproppred import PygGraphPropPredDataset

dataset = PygGraphPropPredDataset(name='ogbn-proteins')
split_idx = dataset.get_idx_split()

train_idx = split_idx['train']
val_idx = split_idx['valid']
test_idx = split_idx['test']

data = dataset[0]

4. 查看数据集信息

print(data)
print(data.num_nodes)
print(data.num_edges)

5. 可以使用DataLoader来批量加载数据

from torch_geometric.data import DataLoader

train_loader = DataLoader(dataset[train_idx], batch_size=32, shuffle=True)
val_loader = DataLoader(dataset[val_idx], batch_size=32, shuffle=False)
test_loader = DataLoader(dataset[test_idx], batch_size=32, shuffle=False)