使用 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 檢查點、模型和資料集，供社群其他人使用！