【Code】OGB：图机器学习的基准测试数据集

大家好，欢迎来到IT知识分享网。

1.OGB

1.1 Overview

Open Graph Benchmark（以下简称 OGB）是斯坦福大学的同学开源的 Python 库，其包含了图机器学习（以下简称图 ML）的基准数据集、数据加载器和评估器，目的在于促进可扩展的、健壮的、可复现的图 ML 的研究。

OGB 包含了多种图机器学习的多种任务，并且涵盖从社会和信息网络到生物网络，分子图和知识图的各种领域。没有数据集都有特定的数据拆分和评估指标，从而提供统一的评估协议。

OGB 提供了一个自动的端到端图 ML 的 pipeline，该 pipeline 简化并标准化了图数据加载，实验设置和模型评估的过程。如下图所示：

下图展示了 OGB 的三个维度，包括任务类型（Tasks）、可扩展性（Scale）、领域（Rich domains）。

1.2 Dataset

来看一下 OGB 现在包含的数据集：

和数据集的统计明细：

1.3 Leaderboard

OGB 也提供了标准化的评估人员和排行榜，以跟踪最新的结果，我们来看下不同任务下的部分 Leaderboard。

2.OGB+DGL

官方给出的例子都是基于 PyG 实现的，我们这里实现一个基于 DGL 例子。

2.1 环境准备

导入数据包

import dgl import ogb import math import time import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from ogb.nodeproppred import DglNodePropPredDataset, Evaluator 

查看版本

print(dgl.__version__) print(torch.__version__) print(ogb.__version__) 

0.4.3post2 1.5.0+cu101 1.1.1 

cuda 相关信息

print(torch.version.cuda) print(torch.cuda.is_available()) print(torch.cuda.device_count()) print(torch.cuda.get_device_name(0)) print(torch.cuda.current_device()) 

10.1 True 1 Tesla P100-PCIE-16GB 0 

2.2 数据准备

设置参数

device_id=0 # GPU 的使用 id n_layers=3 # 输入层 + 隐藏层 + 输出层的数量 n_hiddens=256 # 隐藏层节点的数量 dropout=0.5 lr=0.01 epochs=300 runs=10 # 跑 10 次，取平均 log_steps=50 

定义训练函数、测试函数和日志记录

def train(model, g, feats, y_true, train_idx, optimizer): """ 训练函数 """ model.train() optimizer.zero_grad() out = model(g, feats)[train_idx] loss = F.nll_loss(out, y_true.squeeze(1)[train_idx]) loss.backward() optimizer.step() return loss.item() @torch.no_grad() def test(model, g, feats, y_true, split_idx, evaluator): """ 测试函数 """ model.eval() out = model(g