Vision Transformers 在 Hugging Face Optimum Graphcore 上

Hugging Face 與 Graphcore 已將視覺轉換器 (ViT) 模型整合至 Hugging Face Optimum Graphcore 函式庫，讓使用者能在 Graphcore 智慧處理單元（IPU）上微調這些模型。此整合透過結合預訓練的 ViT 檢查點與 IPU 專屬硬體加速，簡化了高效能電腦視覺模型的部署。

視覺轉換器 (ViT) 架構

視覺轉換器 (ViT) 將自注意力機制——最初為自然語言處理（NLP）模型如 BERT 與 GPT 所開發——應用於影像辨識。與使用像素陣列的卷積神經網路（CNN）不同，ViT 模型將輸入影像切割成小的圖塊（視覺標記）。每個圖塊會被線性編碼成向量表示，供轉換器逐一處理。

Key characteristics of the ViT approach include:

  • 特徵提取：預訓練使 ViT 能學習影像的內部表示。可在預訓練的視覺編碼器之上（通常是 [CLS] 標記）加入線性層，以執行下游分類任務。
  • 效能：相較於 CNN，ViT 模型能以較低的計算成本達到更高的辨識準確度。
  • 應用：ViT 被廣泛應用於各種領域，包括目標偵測、分割以及醫療診斷（例如偵測 COVID-19、乳癌與阿茲海默症）。

為 Graphcore IPU 優化 ViT

Graphcore IPU 為 ViT 模型的大規模平行特性而設計。該硬體採用 MIMD（多指令多資料）架構與 IPU‑Fabric 可擴展解決方案，以平行化訓練。

Technical optimizations provided by the Optimum Graphcore library include:

  • 平行化：使用管線平行化可提升每個資料平行實例的批次大小，改善記憶體存取效率，並減少參數聚合的通訊時間。
  • 即用型檢查點：Graphcore 提供已在 IPU 上訓練的模型檢查點與設定檔（例如 Graphcore/vit-base-ipu），免除使用者從頭訓練 ViT 模型的需求，後者通常需要龐大的資料集。
  • 無縫整合：使用者可利用 Hugging Face Model Hub 上的現有檢查點（如 google/vit-base-patch16-224-in21k），並透過 IPUConfig 與 IPUTrainer 類別套用 IPU 專屬設定。

案例研究：多標籤胸部 X 光分類

為展示 Optimum Graphcore 工作流程的效能，我們在 ChestX-ray14 資料集上微調了 ViT 模型，該資料集包含 30,805 位患者的 112,120 張正面 X 光影像。

資料集準備

由於 X 光影像可能同時顯示多種疾病，此任務被視為多標籤分類問題。工作流程包括：

  1. 標籤編碼：將文字標籤轉換為 N-hot 編碼的陣列（布林值），以表示多種同時存在的疾病。
  2. 前處理：使用 AutoImageProcessor 將影像調整為 224x224，將灰階影像轉為 RGB，並以 0.5 的均值與標準差對通道進行正規化。
  3. 資料載入：透過 Hugging Face Datasets 使用 Arrow 檔案格式，以加速訓練期間的載入。

訓練與評估

  • 模型：使用 google/vit-base-patch16-224-in21k 檢查點（在 1400 萬張 ImageNet-21k 圖片上預訓練）。
  • 訓練器：採用 IPUTrainer 類別，負責模型在 IPU 上的編譯以及訓練與評估的管理。
  • 指標：以 ROC 曲線下面積（AUC_ROC）作為主要效能指標，因其對類別不平衡不敏感，且能有效衡量模型區分不同疾病的能力。
  • 結果：訓練日誌顯示損失快速下降，最終在約 0.1 附近穩定，學習率在 25% 的預熱階段後採用餘弦衰減。

透過 Paperspace Gradient 的可存取性

透過與 Paperspace 的合作，開發者可透過 Gradient 的網頁式 Jupyter Notebook 存取以 Graphcore IPU 為基礎的 Hugging Face Optimum 模型。這讓開發者能在 IPU 硬體上試驗 ViT、BERT 與 RoBERTa，而無需本地基礎設施。

Sources