使用 Transformers 進行圖形分類

Hugging Face 提供了一份關於使用 Transformers 套件實作圖形分類的技術指南。目前套件中可用於此任務的主要模型是微軟的 Graphormer，使用者可以將基於 transformer 的架構套用於圖結構資料，執行二元分類、多類別分類或回歸等任務。

圖形資料格式化與載入

在 Hugging Face Hub 上的圖形資料集主要以 jsonl 格式儲存，作為圖形列表。每個圖形以字典形式表示，包含以下必須與可選欄位：

  • edge_index：一個列表，內含兩個平行的整數列表，代表邊中節點的索引（例如 [[1, 1, 3], [2, 3, 1]] 表示邊 1–2、1–3、以及 3–1）。
  • num_nodes：一個整數，表示圖形中節點的總數，以確保孤立節點也被計算在內。
  • y：預測的目標值，可為整數（用於多類別分類）、浮點數（用於回歸），或二元標籤列表（用於多任務分類）。
  • node_feat（可選）：一個整數列表的列表，包含每個節點的特徵，依節點索引排序。
  • edge_attr（可選）：一個整數列表的列表，包含每條邊的屬性，依 edge_index 的順序排列。

使用者可以直接透過 datasets 套件從 Hub 載入這些資料集，例如來自 Open Graph Benchmark（OGB）的 ogbg-molhiv 資料集。

Graphormer 的前處理

圖形 transformer 框架需要特定的前處理，以產生有助於學習的特徵。對於 Graphormer，這包括產生節點間的最短路徑矩陣以及進出度資訊。

前處理可以透過兩種方式處理：

  1. 靜態前處理：使用 dataset.map(preprocess_item, batched=False) 在訓練前處理資料。
  2. 即時前處理：在 GraphormerDataCollator 中設定 on_the_fly_processing=True，於訓練時即時處理資料，對大型圖形更節省記憶體。

模型實作與微調

圖形分類使用 GraphormerForGraphClassification 類別實作。工作流程包括載入預訓練檢查點，並將分類頭調整至特定的下游任務。

載入與設定

透過 from_pretrained 載入模型時，使用者可以指定 num_classes 以符合其特定任務（例如二元分類時 num_classes=2）。將 ignore_mismatched_sizes=True 設定為真時，套件會建立自訂的分類頭，取代預訓練檢查點中的原始解碼頭。

訓練工作流程

訓練過程透過 Trainer API 管理。主要的技術考量包括：

  • 資料整理器：必須將 GraphormerDataCollator 傳入 Trainer，以將單一圖形轉換為批次。
  • 記憶體管理：由於模型尺寸，調整 per_device_train_batch_size 與 gradient_accumulation_steps 對防止記憶體不足（OOM）錯誤至關重要。
  • 硬體：雖然可以在 CPU 上訓練，但指南指出在 Intel Core i7 上訓練 20 個 epoch 大約需要一天，實務上強烈建議使用 GPU 加速。

Sources