Hugging Face Optimum Graphcore 上の Vision Transformers

Hugging Face と Graphcore は、Vision Transformer（ViT）モデルを Hugging Face Optimum Graphcore ライブラリに統合し、ユーザーが Graphcore Intelligence Processing Units（IPU）上でこれらのモデルをファインチューニングできるようにしました。この統合により、事前学習済み ViT チェックポイントと IPU 固有のハードウェアアクセラレーションを組み合わせることで、高性能なコンピュータビジョンモデルの展開が簡素化されます。

Vision Transformer（ViT）アーキテクチャ

Vision Transformers（ViT）は、自己注意機構を画像認識に適用します。この機構はもともと BERT や GPT などの自然言語処理（NLP）モデル向けに開発されました。ピクセル配列を使用する畳み込みニューラルネットワーク（CNN）とは異なり、ViT モデルは入力画像を小さなパッチ（ビジュアルトークン）に分割します。各パッチは線形にエンコードされ、ベクトル表現となり、トランスフォーマーが個別に処理します。

  • Feature Extraction: 事前学習により ViT は画像の内部表現を学習できます。事前学習済みのビジュアルエンコーダ（通常は [CLS] トークン）の上に線形層を追加することで、下流の分類タスクを実行できます。
  • Performance: ViT モデルは、CNN と比較して計算コストを抑えつつ、より高い認識精度を達成できます。
  • Applications: ViT は、物体検出、セグメンテーション、医療診断（例：COVID-19、乳がん、アルツハイマー病の検出）など、さまざまな分野で利用されています。

Graphcore IPU 向け ViT の最適化

Graphcore の IPU は、ViT モデルの大規模並列性に対応するよう設計されています。ハードウェアは MIMD（Multiple Instruction, Multiple Data）アーキテクチャと IPU‑Fabric スケールアウトソリューションを活用し、トレーニングを並列化します。

Optimum Graphcore ライブラリが提供する技術的最適化には、以下が含まれます。

  • Parallelism: パイプライン並列を使用することで、データ並列インスタンスあたりのバッチサイズが増加し、メモリアクセス効率が向上し、パラメータ集約の通信時間が削減されます。
  • Ready-to-use Checkpoints: Graphcore は、IPU で学習されたモデルチェックポイントと設定ファイル（例：Graphcore/vit-base-ipu）を提供し、通常は大規模データセットが必要な ViT のスクラッチ学習をユーザーが行う必要をなくします。
  • Seamless Integration: ユーザーは google/vit-base-patch16-224-in21k など、Hugging Face Model Hub にある既存のチェックポイントを活用し、IPUConfig と IPUTrainer クラスを通じて IPU 固有の設定を適用できます。

ケーススタディ：マルチラベル胸部X線分類

Optimum Graphcore ワークフローの効率性を示すために、ViT モデルを ChestX-ray14 データセットでファインチューニングしました。このデータセットは、30,805 人の患者からの 112,120 枚の正面ビュー X 線画像を含みます。

データセットの準備

X 線画像は同時に複数の疾患を示すことがあるため、タスクはマルチラベル分類問題として扱われます。ワークフローは以下を含みます。

  1. Label Encoding: テキストラベルを N‑hot エンコードされた配列（ブール値）に変換し、同時に存在する複数の疾患を表現します。
  2. Preprocessing: AutoImageProcessor を使用して画像を 224x224 にリサイズし、グレースケール画像を RGB に変換し、各チャネルを平均 0.5、標準偏差 0.5 で正規化します。
  3. Data Loading: Hugging Face Datasets を介して Arrow ファイル形式を利用し、トレーニング中の高速ロードを実現します。

トレーニングと評価

  • Model: google/vit-base-patch16-224-in21k チェックポイント（14 百万枚の ImageNet‑21k 画像で事前学習）を使用しました。
  • Trainer: IPUTrainer クラスを使用し、IPU 用のモデルコンパイルとトレーニング・評価の管理を行いました。
  • Metric: ROC 曲線下面積（AUC_ROC）を主要な性能指標として実装しました。これはクラス不均衡に影響されず、異なる疾患を区別するモデルの能力を効果的に測定します。
  • Results: トレーニングログは損失が急速に減少し、約 0.1 で安定したことを示しました。学習率は 25% のウォームアップ期間の後、コサイン減衰で変化しました。

Paperspace Gradient を通じたアクセス

Paperspace との提携により、Graphcore IPU を搭載した Hugging Face Optimum モデルへ Gradient のウェブベース Jupyter ノートブックからアクセスできるようになりました。これにより、開発者はローカルインフラを必要とせずに、IPU ハードウェア上で ViT、BERT、RoBERTa を試すことができます。

Sources