nanoVLM：用於訓練視覺語言模型的極簡 PyTorch 工具包

Hugging Face 推出了 nanoVLM，這是一個以純 PyTorch 編寫的極簡工具包，旨在闡明訓練視覺語言模型（VLM） 的過程。靈感來自 Andrej Karpathy 的 nanoGPT，nanoVLM 提供了可讀且輕量的程式碼庫，讓使用者能在免費等級的 Google Colab 筆記本上訓練 VLM。

視覺語言模型（VLM）基礎

視覺語言模型是一種多模態架構，能同時處理圖像與文字輸入並產生文字輸出。雖然 VLM 可用於圖像說明、目標檢測、語意分割等多種任務，nanoVLM 主要聚焦於 Visual Question Answering (VQA) 作為其主要訓練目標。

技術架構

nanoVLM 採用模組化架構，透過投影層對齊兩個預訓練的骨幹模型：

  • Vision Backbone（視覺骨幹）: 使用 Google 的 SigLIP (google/siglip-base-patch16-224) 視覺編碼器。
  • Language Backbone（語言骨幹）: 採用 Llama 3 架構，具體使用 HuggingFaceTB/SmolLM2-135M。
  • Modality Projection Module（模態投影模組）: 此模組對齊視覺與文字模態。它將來自視覺骨幹的圖像嵌入轉換為與語言模型嵌入層相容的嵌入。此過程包括 pixel shuffle 操作——可減少圖像 token 數量以降低計算成本並提升訓練速度——接著是一個線性層。

訓練工作流程與實作

訓練流程由 train.py 管理，負責資料集載入、模型初始化與最佳化。

資料管線

get_dataloaders 函式利用 Hugging Face 的 load_dataset API 來載入、洗牌與切分資料集。它使用自訂資料集（VQADataset、MMStarDataset）與整理器（VQACollator、MMStarCollator）來準備資料。

最佳化策略

為了在訓練預訓練骨幹與新初始化的投影器之間取得平衡，nanoVLM 採用雙學習率（LR）策略：

  • Higher LR（較高學習率）: 套用於 Modality Projector（MP），以促進快速學習。
  • Lower LR（較低學習率）: 套用於編碼器/解碼器堆疊，以保留骨幹中已有的知識。

訓練迴圈與監控

訓練使用 torch.autocast 進行混合精度，並採用餘弦學習率排程搭配線性 warmup。效能透過 token 吞吐量（tokens/sec）監控，若啟用，亦會使用 Weights & Biases（wandb）追蹤批次損失、驗證損失與準確率。

推論與預訓練模型

Hugging Face 提供了已發布至 Hub 的預訓練 nanoVLM 模型（nanoVLM-222M）。此模型使用 cauldron 資料集的 170 萬樣本，在單一 H100 GPU 上訓練約 6 小時。

使用者可以透過 generate.py 腳本執行推論，其邏輯流程如下：

  1. Initialization（初始化）: 載入模型、分詞器與影像處理器。
  2. Processing（處理）: 對文字提示進行分詞，並將影像處理成張量。
  3. Generation（生成）: 執行 model.generate 產生文字輸出。
  4. Decoding（解碼）: 使用 batch_decode 將生成的 token 轉回可讀的文字。

快速開始

要開始訓練，使用者可以克隆此倉庫並執行訓練腳本：

# Clone the repo
git clone https://github.com/huggingface/nanoVLM.git

# Execute the training script
python train.py

Sources