使用 Transformer 进行图分类

在之前的博客中，我们探讨了图机器学习的一些理论知识。本篇博客将探讨如何使用 Transformers 库进行图分类。(你也可以通过下载演示 notebook 在此处跟进！)

目前，Transformers 中唯一可用的图 Transformer 模型是微软的 Graphormer，所以我们将在这里使用它。我们期待看到大家会使用和集成哪些其他模型 🤗

环境要求

要学习本教程，你需要安装 datasets 和 transformers (版本 >= 4.27.2)，你可以通过 pip install -U datasets transformers 来安装。

数据

要使用图数据，你可以使用自己的数据集，也可以使用 Hub 上可用的数据集。我们将重点介绍如何使用已有的数据集，但也欢迎你添加自己的数据集！

加载数据

从 Hub 加载图数据集非常简单。让我们加载 ogbg-mohiv 数据集 (这是斯坦福大学 Open Graph Benchmark 的一个基准数据集)，它存储在 OGB 仓库中。

from datasets import load_dataset

# There is only one split on the hub
dataset = load_dataset("OGB/ogbg-molhiv")

dataset = dataset.shuffle(seed=0)

这个数据集已经有三个划分：train、validation 和 test，所有这些划分都包含我们感兴趣的 5 个列 (edge_index、edge_attr、y、num_nodes、node_feat)，你可以通过执行 print(dataset) 来查看。

如果你有其他图处理库，你可以用它们来绘制图并进一步检查数据集。例如，使用 PyGeometric 和 matplotlib。

import networkx as nx
import matplotlib.pyplot as plt

# We want to plot the first train graph
graph = dataset["train"][0]

edges = graph["edge_index"]
num_edges = len(edges[0])
num_nodes = graph["num_nodes"]

# Conversion to networkx format
G = nx.Graph()
G.add_nodes_from(range(num_nodes))
G.add_edges_from([(edges[0][i], edges[1][i]) for i in range(num_edges)])

# Plot
nx.draw(G)

格式

在 Hub 上，图数据集主要以图列表的形式存储 (使用 jsonl 格式)。

单个图是一个字典，以下是我们图分类数据集的预期格式。

• edge_index 包含边中节点的索引，存储为包含两个平行边索引列表的列表。
  • 类型: 包含 2 个整数列表的列表。
  • 示例: 一个包含四个节点 (0、1、2 和 3) 且连接为 1->2、1->3 和 3->1 的图，其 edge_index = [[1, 1, 3], [2, 3, 1]]。你可能会注意到节点 0 并未出现，因为它本身不属于任何边。这就是下一个属性很重要的原因。
• num_nodes 表示图中可用的节点总数 (默认情况下，假定节点是按顺序编号的)。
  • 类型: 整数
  • 示例: 在上面的例子中，num_nodes = 4。
• y 将每个图映射到我们想要预测的内容 (可以是一个类别、一个属性值，或者针对不同任务的多个二进制标签)。
  • 类型: 整数列表 (用于多类别分类)、浮点数列表 (用于回归) 或由 1 和 0 组成的列表的列表 (用于二元多任务分类)。
  • 示例: 我们可以预测图的大小 (小 = 0, 中 = 1, 大 = 2)。在这里，y = [0]。
• node_feat 包含图中每个节点可用的特征 (如果存在)，按节点索引排序。
  • 类型: 整数列表的列表 (可选)
  • 示例: 上面的节点可以有类型 (比如分子中的不同原子)。这可能得到 node_feat = [[1], [0], [1], [1]]。
• edge_attr 包含图中每条边可用的属性 (如果存在)，遵循 edge_index 的顺序。
  • 类型: 整数列表的列表 (可选)
  • 示例: 上面的边可以有类型 (比如分子键)。这可能得到 edge_attr = [[0], [1], [1]]。

预处理

图 Transformer 框架通常对其数据集应用特定的预处理，以生成有助于底层学习任务 (在我们的案例中是分类) 的附加特征和属性。在这里，我们使用 Graphormer 的默认预处理，它会生成入/出度信息、节点间最短路径矩阵以及模型感兴趣的其他属性。

from transformers.models.graphormer.collating_graphormer import preprocess_item, GraphormerDataCollator

dataset_processed = dataset.map(preprocess_item, batched=False)

也可以在 DataCollator 的参数中即时应用此预处理 (通过将 on_the_fly_processing 设置为 True)：并非所有数据集都像 ogbg-molhiv 那么小，对于大图，预先存储所有预处理数据可能成本太高。

模型

加载模型

在这里，我们加载一个已有的预训练模型/检查点，并在我们的下游任务上进行微调，这是一个二元分类任务 (因此 num_classes = 2)。我们也可以在回归任务 (num_classes = 1) 或多任务分类上微调我们的模型。

from transformers import GraphormerForGraphClassification

model = GraphormerForGraphClassification.from_pretrained(
    "clefourrier/pcqm4mv2_graphormer_base",
    num_classes=2, # num_classes for the downstream task 
    ignore_mismatched_sizes=True,
)

让我们更详细地看一下。

在我们的模型上调用 from_pretrained 方法会为我们下载并缓存权重。由于 (用于预测的) 类别数量取决于数据集，我们将新的 num_classes 以及 ignore_mismatched_sizes 和 model_checkpoint 一起传递。这确保了会创建一个自定义的分类头，特定于我们的任务，因此可能与原始的解码器头不同。

也可以创建一个新的随机初始化的模型从头开始训练，可以遵循给定检查点的已知参数，也可以手动选择它们。

训练或微调

为了简单地训练我们的模型，我们将使用 Trainer。要实例化它，我们需要定义训练配置和评估指标。最重要的是 TrainingArguments，这是一个包含所有用于自定义训练的属性的类。它需要一个文件夹名称，该名称将用于保存模型的检查点。

from transformers import TrainingArguments, Trainer

training_args = TrainingArguments(
    "graph-classification",
    logging_dir="graph-classification",
    per_device_train_batch_size=64,
    per_device_eval_batch_size=64,
    auto_find_batch_size=True, # batch size can be changed automatically to prevent OOMs
    gradient_accumulation_steps=10,
    dataloader_num_workers=4, #1, 
    num_train_epochs=20,
    evaluation_strategy="epoch",
    logging_strategy="epoch",
    push_to_hub=False,
)

对于图数据集，调整批大小和梯度累积步数尤其重要，以便在避免内存不足错误的同时训练足够多的样本。

最后一个参数 push_to_hub 允许 Trainer 在训练期间定期将模型推送到 Hub，即在每个保存步骤中。

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset_processed["train"],
    eval_dataset=dataset_processed["validation"],
    data_collator=GraphormerDataCollator(),
)

在用于图分类的 Trainer 中，传递针对给定图数据集的特定数据整理器 (data collator) 很重要，它将把单个图转换为用于训练的批次。

train_results = trainer.train()
trainer.push_to_hub()

当模型训练完成后，可以使用 push_to_hub 将其与所有相关的训练产出物一起保存到 Hub。

由于这个模型相当大，在 CPU (IntelCore i7) 上训练/微调 20 个 epoch 大约需要一天的时间。为了加快速度，你可以使用强大的 GPU 和并行化，可以在 Colab notebook 中或直接在你选择的集群上运行代码。

结语

既然你已经知道如何使用 transformers 来训练图分类模型，我们希望你会尝试在 Hub 上分享你最喜欢的图 transformer 检查点、模型和数据集，供社区其他人使用！