OGB数据加载器深度教程：与PyTorch Geometric和DGL的无缝集成

【免费下载链接】ogb Benchmark datasets, data loaders, and evaluators for graph machine learning 项目地址: https://gitcode.com/gh_mirrors/og/ogb

OGB（Open Graph Benchmark）是一个用于图机器学习的基准数据集、数据加载器和评估器项目，提供了与PyTorch Geometric和DGL等主流图深度学习框架的无缝集成，帮助开发者轻松处理图数据。

什么是OGB数据加载器？

OGB数据加载器是OGB项目的核心组件之一，它能够自动下载、解析和预处理各种图数据集，并将其转换为适合PyTorch Geometric或DGL框架使用的格式。这极大地简化了图机器学习研究和应用的流程，让开发者可以专注于模型设计和实验。

OGB数据加载器架构示意图，展示了数据处理的完整流程

安装OGB

要使用OGB数据加载器，首先需要安装OGB库。可以通过以下命令从Git仓库克隆并安装：

git clone https://gitcode.com/gh_mirrors/og/ogb
cd ogb
pip install .

与PyTorch Geometric集成

OGB提供了专门针对PyTorch Geometric的数据集类PygGraphPropPredDataset，位于ogb/graphproppred/dataset_pyg.py文件中。使用这个类可以轻松加载图属性预测数据集。

基本用法

from ogb.graphproppred import PygGraphPropPredDataset

# 加载数据集
dataset = PygGraphPropPredDataset(name='ogbg-molhiv')

# 获取数据集分割
split_index = dataset.get_idx_split()
train_idx, valid_idx, test_idx = split_index['train'], split_index['valid'], split_index['test']

# 获取单个图数据
graph = dataset[0]
print(graph)  # 输出图的基本信息

数据集类的核心功能

PygGraphPropPredDataset类提供了以下关键方法和属性：

• __init__(): 初始化数据集，自动下载和处理数据
• get_idx_split(): 获取训练、验证和测试集的索引
• num_classes: 获取分类任务的类别数
• __getitem__(): 获取指定索引的图数据

与DGL集成

同样，OGB也提供了针对DGL的数据集类DglGraphPropPredDataset，位于ogb/graphproppred/dataset_dgl.py文件中。

基本用法

from ogb.graphproppred import DglGraphPropPredDataset

# 加载数据集
dataset = DglGraphPropPredDataset(name='ogbg-molhiv')

# 获取数据集分割
split_index = dataset.get_idx_split()
train_idx, valid_idx, test_idx = split_index['train'], split_index['valid'], split_index['test']

# 获取单个图数据
graph, label = dataset[0]
print(graph)  # 输出DGL图对象

批处理图数据

OGB还提供了collate_dgl函数，用于将多个图数据批处理成一个批次：

from ogb.graphproppred import collate_dgl
from dgl.dataloading import DataLoader

# 创建数据加载器
dataloader = DataLoader(dataset, batch_size=32, shuffle=True, collate_fn=collate_dgl)

# 迭代处理批次数据
for batch_graph, batch_labels in dataloader:
    print(batch_graph)  # 输出批处理后的图
    print(batch_labels)  # 输出对应的标签

OGB支持多GPU数据处理，加速大规模图数据集的训练过程

支持的数据集类型

OGB数据加载器支持多种类型的图数据集，包括：

1. 图属性预测：如分子性质预测数据集ogbg-molhiv
2. 节点属性预测：如学术网络数据集ogbn-arxiv
3. 链接属性预测：如知识图谱数据集ogbl-wikikg2

这些数据集可以通过修改name参数来加载，例如：

# 加载节点属性预测数据集
dataset = PygNodePropPredDataset(name='ogbn-arxiv')

# 加载链接属性预测数据集
dataset = PygLinkPropPredDataset(name='ogbl-wikikg2')

高级功能

数据集分割

OGB数据加载器提供了灵活的数据集分割方式，可以通过get_idx_split()方法获取不同的分割策略：

# 获取默认分割
split_index = dataset.get_idx_split()

# 获取特定分割类型
split_index = dataset.get_idx_split(split_type='time')  # 按时间分割

数据转换

可以在加载数据时应用转换函数，对图数据进行预处理：

from torch_geometric.transforms import AddSelfLoops

# 添加自环
dataset = PygGraphPropPredDataset(name='ogbg-molhiv', transform=AddSelfLoops())

OGB数据加载器处理知识图谱数据的示例

实际应用示例

以下是一个使用OGB数据加载器和PyTorch Geometric进行图分类的简单示例：

from ogb.graphproppred import PygGraphPropPredDataset, Evaluator
import torch_geometric.nn as pyg_nn
import torch.nn.functional as F
from torch_geometric.loader import DataLoader

# 加载数据集和评估器
dataset = PygGraphPropPredDataset(name='ogbg-molhiv')
evaluator = Evaluator(name='ogbg-molhiv')

# 分割数据集
split_index = dataset.get_idx_split()
train_loader = DataLoader(dataset[split_index['train']], batch_size=32, shuffle=True)
valid_loader = DataLoader(dataset[split_index['valid']], batch_size=32, shuffle=False)

# 定义模型
class GNN(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = pyg_nn.GCNConv(dataset.num_features, hidden_channels)
        self.conv2 = pyg_nn.GCNConv(hidden_channels, out_channels)

    def forward(self, x, edge_index, batch):
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index)
        x = pyg_nn.global_mean_pool(x, batch)
        return x

model = GNN(hidden_channels=64, out_channels=dataset.num_classes)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# 训练模型
model.train()
for batch in train_loader:
    out = model(batch.x, batch.edge_index, batch.batch)
    loss = F.binary_cross_entropy_with_logits(out, batch.y)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

# 评估模型
model.eval()
valid_preds = []
valid_labels = []
for batch in valid_loader:
    out = model(batch.x, batch.edge_index, batch.batch)
    valid_preds.append(out)
    valid_labels.append(batch.y)
valid_preds = torch.cat(valid_preds, dim=0)
valid_labels = torch.cat(valid_labels, dim=0)
roc_auc = evaluator.eval({'y_pred': valid_preds, 'y_true': valid_labels})['rocauc']
print(f'Validation ROC-AUC: {roc_auc:.4f}')

总结

OGB数据加载器为图机器学习提供了强大而便捷的数据处理工具，通过与PyTorch Geometric和DGL的无缝集成，大大降低了图数据处理的门槛。无论是学术研究还是工业应用，OGB数据加载器都能帮助开发者快速构建和训练图神经网络模型。

通过本文介绍的方法，你可以轻松开始使用OGB数据加载器处理各种图数据集，加速你的图机器学习项目开发。更多详细信息和高级用法，请参考OGB项目的官方文档和示例代码。

OGB数据加载器与其他方法在知识图谱任务上的性能对比

【免费下载链接】ogb Benchmark datasets, data loaders, and evaluators for graph machine learning 项目地址: https://gitcode.com/gh_mirrors/og/ogb