在之前的 博文 中,咱们探讨了图机器学习的一些理论知识。这一篇咱们将探索怎么运用 Transformers 库进行图分类。(你也能够从 此处 下载演示 notebook,跟着一起做!)

目前,Transformers 中仅有可用的图 transformer 模型是微软的 Graphormer,因此本文的比如将会依据该模型。咱们等待看到大家会运用并集成哪些其他模型进 。

软件

要学习本教程,需求装置 datasetstransformers (版本号 >= 4.27.2),你能够运用 pip install -U datasets transformers 来装置。

数据

你能够运用自己的图数据集,也能够运用 Hub 上已有的数据集。本文咱们主要运用已有的数据集,你也能够随时 添加你的数据集 到 Hugging Face!

数据加载

从 Hub 加载图数据集非常简略。这儿,咱们加载 OGB 库中的 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)

这个数据集含三个拆分,trainvalidationtest,一切这些拆分每一行都表示一个图,每个图包括 5 个数据列 (edge_indexedge_attrynum_nodesnode_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 包括图上每条边对应的节点 ID,存储为包括两个节点列表的列表 (即由一个源节点列表和一个意图节点列表组成的列表)。

    • 类型: 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 每个图的猜测标签 (能够是类、特点值或是不同使命的多个二分类标签)。

    • Type: 整数列表 (用于多分类) 、浮点数 (用于回归) 或 0/1 列表 (用于二元多使命分类)
    • 示例: 咱们能够猜测图规划 (小 = 0,中 = 1,大 = 2)。本例中,y = [0]
  • node_feat 包括图中每个节点的可用特征 (假如存在),按节点 ID 排序。

    • 类型: 整数列表的列表 (可选)
    • 比如: 如上例中的节点能够有一些类型特征 (就像分子图中的节点是不同的原子,不同的原子有不同的类型相同)。打比方,本例中 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 那样小,对于大图,动态预处理成本太高,因此需求预先进行预处理,并存储预处理后的数据供后续练习实验运用。

模型

模型加载

这儿,咱们加载一个已有的预练习模型及其 checkpoint 并在咱们的下流使命上对其进行微调,该使命是一个二分类使命 (因此 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_classesignore_mismatched_sizesmodel_checkpoint 一起传给该函数。这会触发函数创建一个自界说的、特定于该下流使命的分类头,这个头与原模型中的解码器头很可能是不同的。

咱们也能够创建一个新的随机初始化的模型来从头开始练习,此时,咱们既能够复用给定检查点的超参装备,也能够自己手动挑选超参装备。

练习或微调

为了简化模型练习,咱们运用 Trainer。咱们需求界说练习相关的装备以及评估目标来实例化 Trainer。咱们主要运用 TrainingArguments类,这是一个包括一切装备项的类,用于定制练习装备。咱们要给它一个文件夹称号,用于保存模型的 checkpoint。

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,
)

对于图数据集,调整 batch size 和梯度累积步数来保证有效 batch size 够大一起又要防止内存不足,这件事尤为重要。

最后一个参数 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) 很重要,这个数据收拾器会将图转换为用于练习的 batch 数据。

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

练习完后,能够运用 push_to_hub 将模型与一切其他练习相关信息一起保存到 hub。

由于此模型比较大,因此在 CPU (Intel Core i7) 上练习/微调 20 个 epoch 大约需求一天时刻。想要更快点的话,你能够运用强大的 GPU 和并行化办法,你只需在 Colab notebook 中或直接在你挑选的其他集群上启动代码即可。

结束语

现在你现已知道怎么运用 transformers 来练习图分类模型,咱们期望你测验在 Hub 上分享你最喜欢的图 transformer 模型的 checkpoints、模型以及数据集,以供社区的其他人运用!


英文原文: hf.co/blog/graphm…

作者: Clm
译者: Matrix Yao (姚伟峰),英特尔深度学习工程师,作业方向为 transformer-family 模型在各模态数据上的应用及大规划模型的练习推理。

排版/审校: zhongdongy (阿东)