使用 Transformers 進行圖分類

HuggingFace發表於2023-04-18
ORM

在之前的 博文 中，我們探討了圖機器學習的一些理論知識。這一篇我們將探索如何使用 Transformers 庫進行圖分類。(你也可以從 此處 下載演示 notebook，跟著一起做！)

目前，Transformers 中唯一可用的圖 transformer 模型是微軟的 Graphormer，因此本文的例子將會基於該模型。我們期待看到大家會使用並整合哪些其他模型進 ?。

軟體

要學習本教程，需要安裝 datasets 和 transformers (版本號 >= 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)

這個資料集含三個拆分，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 包含圖上每條邊對應的節點 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_classes 和 ignore_mismatched_sizes 與 model_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、模型以及資料集，以供社群的其他人使用！

英文原文: <url>https://hf.co/blog/graphml-classification</url>

作者: Clém
譯者: Matrix Yao (姚偉峰)，英特爾深度學習工程師，工作方向為 transformer-family 模型在各模態資料上的應用及大規模模型的訓練推理。

排版/審校: zhongdongy (阿東)

相關文章