【GNN基准】图机器学习的基准测试数据集

图神经网络的ImageNet？斯坦福大学等开源百万量级OGB基准测试数据集

在满是「MNIST」这样的小数据里，图神经网络也需要「ImageNet」这样的大基准？近日，斯坦福大学的 Jure Leskovec 教授在 NeurlPS 2019 大会演讲中宣布开源 Open Graph Benchmark，这是迈向图神经网络建模统一基准的重要一步。

转载：

公众号机器之心 图神经网络的ImageNet？斯坦福大学等开源百万量级OGB基准测试数据集 - 知乎

公众号

图神经网络是近来发展较快的机器学习分支领域。通过将非结构数据转换为结构化的节点和边的图，然后采用图神经网络进行学习，往往能够取得更好的效果。

然而，图神经网络发展到现在，尚无一个公认的基准测试数据集。许多论文采用的方法往往是针对较小的、缺乏节点和边特征的数据集上进行的。因此，在这些数据集上取得的模型性能很难说是最好的，也不一定可靠，这对进一步发展造成阻碍。

在 NeurlPS 2019 大会的图表示学习演讲中，Jure Leskovec 宣布开源图神经网络的通用性能评价基准数据集 OGB（Open Graph Benchmark）。通过这一数据集，可以更好地评估模型性能等方面的指标。

• 项目地址：http://ogb.stanford.edu
• 图表示学习演讲合集：https://slideslive.com/38921872/graph-representation-learning-3

本次演讲的嘉宾为 Jure Leskovec，是斯坦福大学计算机科学的副教授。他主要的研究兴趣是社会信息网络的挖掘和建模等，特别是针对大规模数据、网络和媒体数据。

值得注意的是，OGB 数据集也支持了 PYG 和 DGL 这两个常用的图神经网络框架。DGL 项目的发起人之一、AWS 上海 AI 研究院院长，上海纽约大学张峥教授（学术休假中）说：「现阶段，我认为 OGB 的最大作用是促成学界走出玩具型数据集。一个统一的、更复杂、更多样的数据集会使得研究人员重新聚集力量，虽然还会有模型过拟合标准数据集带来的弊端，但对提升模型和算法效果、提高 DGL 等平台的能力有着重要作用。」

张峥教授表示，Open Graph Benchmark 这种多样与统一的基准，对于图神经网络来说，是非常有必要的一步。

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, feats)
    y_pred = out.argmax(dim=-1, keepdim=True)

    train_acc = evaluator.eval({
        'y_true': y_true[split_idx['train']],
        'y_pred': y_pred[split_idx['train']],
    })['acc']
    valid_acc = evaluator.eval({
        'y_true': y_true[split_idx['valid']],
        'y_pred': y_pred[split_idx['valid']],
    })['acc']
    test_acc = evaluator.eval({
        'y_true': y_true[split_idx['test']],
        'y_pred': y_pred[split_idx['test']],
    })['acc']

    return train_acc, valid_acc, test_acc

class Logger(object):
    """ 用于日志记录
    """
    def __init__(self, runs, info=None):
        self.info = info
        self.results = [[] for _ in range(runs)]

    def add_result(self, run, result):
        assert len(result) == 3
        assert run >= 0 and run