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

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

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

转载：公众号 阿泽的学习笔记

参考公众号 机器之心 https://zhuanlan.zhihu.com/p/98901680

目录

 

简介

1.OGB

1.1 Overview

1.2 Dataset

节点预测

连接预测

图预测

1.3 Leaderboard

数据加载与评估

 PYG

DGL

节点分类：​

链接预测：​

图分类：​

2.OGB+DGL

2.1 环境准备

2.2 数据准备

2.3 GCN

3.为什么说分割图数据是个问题？

4.Conclusion

---

简介

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

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

在 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

节点预测

• odbn-proteins：蛋白质数据集，有着蛋白质之间的关联网络，而且包括了多种生物；
• odbn-wiki：维基百科数据形成的网络；
• ogbn-products：亚马逊客户同时购买的商品的网络。

目前该基准测试所包含的数据集。

从数据集的类型来看，涵盖了现有的几大需要图表示学习的领域：生物学/分子化学、自然语言处理，以及商品推荐系统网络等。此外，这些图数据的量也非常大。例如，ogbn-wiki 的数据量已达到百万级别（节点），而最小的 ogbn-proteins 也有 100K 了。这和之前的很多图数据相比都显得更加庞大，因此也能更好地评价模型的性能表现。

连接预测

连接预测中的数据集则更多一些，包括：

• ogbi-ddi：药物相互作用网络；
• ogbi-biomed：人类生物医药知识图谱；
• ogbi-ppa：蛋白质之间的关系网络；
• ogbi-reviews：亚马逊的用户-商品评论数据集；
• ogbi-citations：微软学术的引用关系网络图谱。

相比节点数据集来说，连接预测的数据集更多一些，类型也更为多样。

图预测

OGB 同时也提供了对图进行预测的任务数据集，分别有：

• ogbg-mol：对分子进行预测，来自 MoleculeNet；
• ogbg-code：代码段的语法树结构网络；
• ogbg-ppi：蛋白质之间的交互网络；

从总体来看，数据集中偏向医药和生物的数据集很多。张峥教授认为，这可能有两个原因，首先是项目主导者 Jure 等在这一领域做了比较多的工作，因此推动这些数据集开源顺理成章。另一个原因是药分子的图数据相对干净，噪声少。而药品的结构是 3D 的，可能需要比较复杂、层数更深的模型解决相关的问题。

对于未来会有哪些数据集加入，张教授认为现在关于异构图的数据还不够多，而现实中的很多数据都是异构图表示的。但是，OGB 的作用依然明显，它能够很好地提升开源图神经网络框架的能力，推动开源社区集中力量解决实际问题。

另外，OGB 数据集中缺少金融、征信等领域的数据集，特别是反欺诈类的。这可能是因为反欺诈数据集脱敏后特征丢失过多的问题所致，但瑕不掩瑜，OGB 无疑帮助图神经网络脱离了所谓的「玩具模型」阶段，开始逐渐进入工业应用。

1.3 Leaderboard

数据加载与评估

OGB 如此庞大的数据量需要专门的代码进行提取。据悉，所有开源的数据集都可以用特定的代码进行提取和加载，使用过程和深度学习框架中的 data_loader 相似。不过在使用前，我们还需要简单地使用「pip install ogb」完成安装。目前 OGB 库主要依赖于 PyTorch、NumPy 和 Scikit-Learn 等常用建模库，当然图神经网络库也可以自由选择 DGL 或 PyTorch Geometric。

DGL：https://github.com/dmlc/dgl

PyG：https://github.com/rusty1s/pytorch_geometric

下面详细看一下 OGB 现在包含的数据集：

和数据集的统计明细：

现在以节点预测为例，OGB 同时支持 PYG 和 DGL 两个图表示学习框架中的数据加载方法，加载代码如下：

 PYG

from ogb.nodeproppred.dataset_pyg import PygNodePropPredDataset
dataset = PygNodePropPredDataset(name = d_name) num_tasks = dataset.num_tasks # obtaining number of prediction tasks in a dataset
splitted_idx = dataset.get_idx_split()train_idx, valid_idx, test_idx = splitted_idx["train"], splitted_idx["valid"], splitted_idx["test"]graph = dataset[0] # pyg graph object

DGL

from ogb.nodeproppred.dataset_dgl import DglNodePropPredDataset
dataset = DglNodePropPredDataset(name = d_name)num_tasks = dataset.num_tasks # obtaining number of prediction tasks in a dataset
splitted_idx = dataset.get_idx_split()train_idx, valid_idx, test_idx = splitted_idx["train"], splitted_idx["valid"], splitted_idx["test"]graph, label = dataset[0] # graph: dgl graph object, label: torch tensor of shape (num_nodes, num_tasks)

可以看出，代码非常简单，使用简便。其中「d_name」可以被替换成任何一个数据集的名字。

同时，项目提供了一些示例代码，以对每个数据集进行评估。如下所示：

from ogb.nodeproppred import Evaluator
evaluator = Evaluator(name = d_name) print(evaluator.expected_input_format) print(evaluator.expected_output_format)

这里，用户可以了解到针对这一数据集的输入和输出的特定格式。

然后，用户可以将输入字典（input dictionary）传递进评估器中，了解实际的性能：

# In most cases, input_dict is# input_dict = {"y_true": y_true, "y_pred": y_pred} result_dict = evaluator.eval(input_dict)

据悉，OGB 官方已指明上海 AWS AI 研究院主打的开源框架 DGL 作为数据导入的平台之一。目前 DGL 兼容 PyTorch、MxNet 作为后端引擎，TensorFlow 也在开发中。实际上 DGL 在异构图和可扩展性已经做了很久，因此下一步可能会和 OGB 在相关领域进行新的技术结合，推动开源框架的发展。

张峥说：「DGL 目前在制药领域已有了一个效果不错的模型库，有了 OGB 数据集对模型库进行迭代，之后应当可以进一步提升。」

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 < len(self.results)
        self.results[run].append(result)

    def print_statistics(self, run=None):
        if run is not None:
            result = 100 * torch.tensor(self.results[run])
            argmax = result[:, 1].argmax().item()
            print(f'Run {run + 1:02d}:')
            print(f'Highest Train: {result[:, 0].max():.2f}')
            print(f'Highest Valid: {result[:, 1].max():.2f}')
            print(f'  Final Train: {result[argmax, 0]:.2f}')
            print(f'   Final Test: {result[argmax, 2]:.2f}')
        else:
            result = 100 * torch.tensor(self.results)

            best_results = []
            for r in result:
                train1 = r[:, 0].max().item()
                valid = r[:, 1].max().item()
                train2 = r[r[:, 1].argmax(), 0].item()
                test = r[r[:, 1].argmax(), 2].item()
                best_results.append((train1, valid, train2, test))

            best_result = torch.tensor(best_results)

            print(f'All runs:')
            r = best_result[:, 0]
            print(f'Highest Train: {r.mean():.2f} ± {r.std():.2f}')
            r = best_result[:, 1]
            print(f'Highest Valid: {r.mean():.2f} ± {r.std():.2f}')
            r = best_result[:, 2]
            print(f'  Final Train: {r.mean():.2f} ± {r.std():.2f}')
            r = best_result[:, 3]
            print(f'   Final Test: {r.mean():.2f} ± {r.std():.2f}')

加载数据

device = f'cuda:{device_id}' if torch.cuda.is_available() else 'cpu'
device = torch.device(device)

# 加载数据，name 为 'ogbn-' + 数据集名
# 自己可以打印出 dataset 看一下
dataset = DglNodePropPredDataset(name='ogbn-arxiv')
split_idx = dataset.get_idx_split()
g, labels = dataset[0]
feats = g.ndata['feat']
g = dgl.to_bidirected(g)
feats, labels = feats.to(device), labels.to(device)
train_idx = split_idx['train'].to(device)

2.3 GCN

实现一个基本的 GCN，这里对每一层都进行了一个 Batch Normalization，去掉的话，精度会下降 2% 左右。

from dgl.nn import GraphConv

class GCN(nn.Module):
    def __init__(self,
                 in_feats,
                 n_hiddens,
                 n_classes,
                 n_layers,
                 dropout):
        super(GCN, self).__init__()

        self.layers = nn.ModuleList()
        self.bns = nn.ModuleList()
        self.layers.append(GraphConv(in_feats, n_hiddens, 'both'))
        self.bns.append(nn.BatchNorm1d(n_hiddens))
        for _ in range(n_layers - 2):
            self.layers.append(GraphConv(n_hiddens, n_hiddens, 'both'))
            self.bns.append(nn.BatchNorm1d(n_hiddens))
        self.layers.append(GraphConv(n_hiddens, n_classes, 'both'))
        self.dropout = dropout

    def reset_parameters(self):
        for layer in self.layers:
            layer.reset_parameters()
        for bn in self.bns:
            bn.reset_parameters()

    def forward(self, g, x):
        for i, layer in enumerate(self.layers[:-1]):
            x = layer(g, x)
            x = self.bns[i](x)
            x = F.relu(x)
            x = F.dropout(x, p=self.dropout, training=self.training)
        x = self.layers[-1](g, x)
        return x.log_softmax(dim=-1)

model = GCN(in_feats=feats.size(-1),
            n_hiddens=n_hiddens,
            n_classes=dataset.num_classes,
            n_layers=n_layers,
            dropout=dropout).to(device)

evaluator = Evaluator(name='ogbn-arxiv')
logger = Logger(runs)

for run in range(runs):
    model.reset_parameters()
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    for epoch in range(1, 1 + epochs):
        loss = train(model, g, feats, labels, train_idx, optimizer)
        result = test(model, g, feats, labels, split_idx, evaluator)
        logger.add_result(run, result)

        if epoch % log_steps == 0:
            train_acc, valid_acc, test_acc = result
            print(f'Run: {run + 1:02d}, '
                  f'Epoch: {epoch:02d}, '
                  f'Loss: {loss:.4f}, '
                  f'Train: {100 * train_acc:.2f}%, '
                  f'Valid: {100 * valid_acc:.2f}% '
                  f'Test: {100 * test_acc:.2f}%')

    logger.print_statistics(run)
logger.print_statistics()

Run: 01, Epoch: 50, Loss: 1.1489, Train: 68.71%, Valid: 68.93% Test: 68.32%
Run: 01, Epoch: 100, Loss: 1.0565, Train: 71.29%, Valid: 69.61% Test: 68.03%
Run: 01, Epoch: 150, Loss: 1.0010, Train: 72.28%, Valid: 70.57% Test: 70.00%
Run: 01, Epoch: 200, Loss: 0.9647, Train: 73.18%, Valid: 69.79% Test: 67.97%
Training time/epoch 0.2617543590068817
Run 01:
Highest Train: 73.54
Highest Valid: 71.16
  Final Train: 73.08
   Final Test: 70.43
Run: 02, Epoch: 50, Loss: 1.1462, Train: 68.83%, Valid: 68.69% Test: 68.50%
Run: 02, Epoch: 100, Loss: 1.0583, Train: 71.17%, Valid: 69.54% Test: 68.06%
Run: 02, Epoch: 150, Loss: 1.0013, Train: 71.98%, Valid: 69.71% Test: 68.06%
Run: 02, Epoch: 200, Loss: 0.9626, Train: 73.23%, Valid: 69.76% Test: 67.79%
Training time/epoch 0.26154680013656617
Run 02:
Highest Train: 73.34
Highest Valid: 70.87
  Final Train: 72.56
   Final Test: 70.42
Run: 03, Epoch: 50, Loss: 1.1508, Train: 68.93%, Valid: 68.49% Test: 67.14%
Run: 03, Epoch: 100, Loss: 1.0527, Train: 70.90%, Valid: 69.75% Test: 68.77%
Run: 03, Epoch: 150, Loss: 1.0042, Train: 72.54%, Valid: 70.71% Test: 69.36%
Run: 03, Epoch: 200, Loss: 0.9679, Train: 73.13%, Valid: 69.92% Test: 68.05%
Training time/epoch 0.26173179904619853
Run 03:
Highest Train: 73.44
Highest Valid: 71.04
  Final Train: 73.06
   Final Test: 70.53
Run: 04, Epoch: 50, Loss: 1.1507, Train: 69.02%, Valid: 68.81% Test: 68.09%
Run: 04, Epoch: 100, Loss: 1.0518, Train: 71.30%, Valid: 70.19% Test: 68.78%
Run: 04, Epoch: 150, Loss: 0.9951, Train: 72.05%, Valid: 68.20% Test: 65.38%
Run: 04, Epoch: 200, Loss: 0.9594, Train: 72.98%, Valid: 70.47% Test: 69.26%
Training time/epoch 0.2618525844812393
Run 04:
Highest Train: 73.34
Highest Valid: 70.88
  Final Train: 72.86
   Final Test: 70.60
Run: 05, Epoch: 50, Loss: 1.1500, Train: 68.82%, Valid: 69.00% Test: 68.47%
Run: 05, Epoch: 100, Loss: 1.0566, Train: 71.13%, Valid: 70.15% Test: 69.47%
Run: 05, Epoch: 150, Loss: 0.9999, Train: 72.48%, Valid: 70.88% Test: 70.27%
Run: 05, Epoch: 200, Loss: 0.9648, Train: 73.37%, Valid: 70.51% Test: 68.96%
Training time/epoch 0.261941517829895
Run 05:
Highest Train: 73.37
Highest Valid: 70.93
  Final Train: 72.77
   Final Test: 70.24
Run: 06, Epoch: 50, Loss: 1.1495, Train: 69.00%, Valid: 68.76% Test: 67.89%
Run: 06, Epoch: 100, Loss: 1.0541, Train: 71.24%, Valid: 69.74% Test: 68.21%
Run: 06, Epoch: 150, Loss: 0.9947, Train: 71.89%, Valid: 69.81% Test: 69.77%
Run: 06, Epoch: 200, Loss: 0.9579, Train: 73.45%, Valid: 70.50% Test: 69.60%
Training time/epoch 0.2620268513758977
Run 06:
Highest Train: 73.70
Highest Valid: 70.97
  Final Train: 73.70
   Final Test: 70.12
Run: 07, Epoch: 50, Loss: 1.1544, Train: 68.93%, Valid: 68.81% Test: 67.97%
Run: 07, Epoch: 100, Loss: 1.0562, Train: 71.17%, Valid: 69.79% Test: 68.45%
Run: 07, Epoch: 150, Loss: 1.0016, Train: 72.41%, Valid: 70.65% Test: 69.87%
Run: 07, Epoch: 200, Loss: 0.9627, Train: 73.12%, Valid: 69.97% Test: 68.20%
Training time/epoch 0.2620680228301457
Run 07:
Highest Train: 73.40
Highest Valid: 71.02
  Final Train: 73.08
   Final Test: 70.49
Run: 08, Epoch: 50, Loss: 1.1508, Train: 68.89%, Valid: 68.42% Test: 67.68%
Run: 08, Epoch: 100, Loss: 1.0536, Train: 71.24%, Valid: 69.24% Test: 67.01%
Run: 08, Epoch: 150, Loss: 1.0015, Train: 72.36%, Valid: 69.57% Test: 67.76%
Run: 08, Epoch: 200, Loss: 0.9593, Train: 73.42%, Valid: 70.86% Test: 70.02%
Training time/epoch 0.2621182435750961
Run 08:
Highest Train: 73.43
Highest Valid: 70.93
  Final Train: 73.43
   Final Test: 69.92
Run: 09, Epoch: 50, Loss: 1.1457, Train: 69.17%, Valid: 68.83% Test: 67.67%
Run: 09, Epoch: 100, Loss: 1.0496, Train: 71.45%, Valid: 69.86% Test: 68.53%
Run: 09, Epoch: 150, Loss: 0.9941, Train: 72.51%, Valid: 69.38% Test: 67.02%
Run: 09, Epoch: 200, Loss: 0.9587, Train: 73.49%, Valid: 70.35% Test: 68.59%
Training time/epoch 0.2621259101231893
Run 09:
Highest Train: 73.64
Highest Valid: 70.97
  Final Train: 73.22
   Final Test: 70.46
Run: 10, Epoch: 50, Loss: 1.1437, Train: 69.16%, Valid: 68.43% Test: 67.17%
Run: 10, Epoch: 100, Loss: 1.0473, Train: 71.43%, Valid: 70.33% Test: 69.29%
Run: 10, Epoch: 150, Loss: 0.9936, Train: 71.98%, Valid: 67.93% Test: 65.06%
Run: 10, Epoch: 200, Loss: 0.9583, Train: 72.93%, Valid: 68.05% Test: 65.43%
Training time/epoch 0.26213142466545103
Run 10:
Highest Train: 73.44
Highest Valid: 70.93
  Final Train: 73.44
   Final Test: 70.26
All runs:
Highest Train: 73.46 ± 0.12
Highest Valid: 70.97 ± 0.09
  Final Train: 73.12 ± 0.34
   Final Test: 70.35 ± 0.21

3.为什么说分割图数据是个问题？

在 Jure Leskovec 的演讲中，他特意强调了 OGB 所采用的数据分割方法，它能构建更合理的评估方案。他表示，似乎随机数据分割并不是一件值得关注的事，但当我们将数据随机分割为训练、验证和测试集时，很可能预测准确率看上去非常不错。但实际上，采用随机分割的模型验证，其效果是过于高估的。

Jure Leskovec 举了个例子，比如说自然科学研究者，他们每次收集的数据肯定不是重复的，他们每次都需要做一系列新实验，因此模型每次都在做分布外的预测。这就要求数据的分割方式需要非常合理，需要模型的泛化能力足够强大以处理这些分布外的数据预测。

谈及数据分割问题，张峥教授说：「我们在和制药行业的研究人员讨论时，都被提醒在训练集上做随机切分是不可取的，因为分子图样本有结构性质，独立同分布假设会对模型的泛化能力有影响，我认为其他领域也有同样的问题。」

为了处理这种情况，OGB 采用的数据分割方法也非常有意思。例如对于分子图数据集，分割方法可以是分子支架（scaffold），具体而言，我们可以通过分子的子结构做聚类，然后将常用的集群作为训练集，将其它非常见集群作为验证与测试集。这种处理方式会迫使神经网络获得更高的泛化性，不然它完全无法预测那些子结构不同的分子。

按物种分割或按代码库分割也是相同的道理，本质上这些数据分割都尝试把某一小部分整体移出来做测试。

最后，Jure Leskovec 也表明，他们预想 OGB 不仅能作为一种广泛使用的研究资源，同时也能作为各种新任务或新模型的真实测试环境。在不久的将来，OGB 将进一步支持更多的图数据集、更多的图建模任务，并同时提供一份开放的 LeadBoard。有了这样的 LeadBoard，我们就能更直观地评估各种图神经网络的特点，了解哪种情况下它们的效果是最好的。

4.Conclusion

目前，OGB 才刚刚起步，5 月 4 号刚发布第一个主要版本，未来还会扩展到千万级别节点的数据集。OGB 这样的多样且统一的基准的出现对 GNN 来说是非常重要的一步，希望也能形成与 NLP、CV 等领域类似的 Leaderboard，不至于每次论文都是在 Cora, CiteSeer 等玩具型数据集上做实验了。