nanoVLM: 使用純 PyTorch 訓練 VLM 的最簡單倉庫

nanoVLM 是使用純 PyTorch 訓練您自己的視覺語言模型 (VLM) 的最簡單方式。它是一個輕量級工具包，讓您可以在免費 Colab 筆記本上啟動 VLM 訓練。

我們受到了 Andrej Karpathy 的 nanoGPT 的啟發，併為視覺領域提供了類似的專案。

其核心在於，nanoVLM 是一個工具包，可幫助您構建和訓練一個能夠理解影像和文字，然後在此基礎上生成文字的模型。nanoVLM 的美妙之處在於它的簡單性。整個程式碼庫有意保持精簡和可讀，使其非常適合初學者或任何想要深入瞭解 VLM 內部機制而不會感到不知所措的人。

在這篇部落格文章中，我們涵蓋了專案背後的核心思想，並提供了與倉庫互動的簡單方法。我們不僅深入探討了專案細節，還將其全部封裝起來，以便您快速上手。

目錄:

• 什麼是視覺語言模型？
• 使用倉庫
• 架構
• 訓練您自己的 VLM
• 在預訓練模型上執行推理
• 結論
• 參考文獻

TL;DR

您可以透過以下步驟使用我們的 nanoVLM 工具包開始訓練視覺語言模型：

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

# Execute the training script
python train.py

這是一個Colab 筆記本，可以幫助您啟動訓練執行，無需本地設定！

什麼是視覺語言模型？

顧名思義，視覺語言模型 (VLM) 是一種多模態模型，它處理兩種模態：視覺和文字。這些模型通常將影像和/或文字作為輸入，並生成文字作為輸出。

根據對影像和文字（輸入）的理解生成文字（輸出）是一種強大的正規化。它支援廣泛的應用，從影像字幕和目標檢測到回答有關視覺內容的問題（如下表所示）。需要注意的是，nanoVLM 僅專注於視覺問答作為訓練目標。

	為影像新增標題	兩隻貓躺在床上，旁邊有遙控器	字幕
	檢測影像中的物件	<locxx><locxx><locxx><locxx>	物體檢測
	分割影像中的物件	<segxx><segxx><segxx>	語義分割
	圖片中有多少隻貓？	2	視覺問答

如果您有興趣瞭解更多關於 VLM 的資訊，我們強烈建議您閱讀我們最新的相關部落格文章：視覺語言模型（更好、更快、更強）

使用倉庫

“空談無用，給我看程式碼”——Linus Torvalds

在本節中，我們將引導您瞭解程式碼庫。在您閱讀時，最好開啟一個標籤頁作為參考。

以下是我們倉庫的資料夾結構。為了簡潔起見，我們刪除了輔助檔案。

.
├── data
│   ├── collators.py
│   ├── datasets.py
│   └── processors.py
├── generate.py
├── models
│   ├── config.py
│   ├── language_model.py
│   ├── modality_projector.py
│   ├── utils.py
│   ├── vision_language_model.py
│   └── vision_transformer.py
└── train.py

架構

.
├── data
│   └── ...
├── models      # 👈 You are here
│   └── ...
└── train.py     

我們根據兩種眾所周知且廣泛使用的架構來建模 nanoVLM。我們的視覺主幹 (models/vision_transformer.py) 是標準的視覺 Transformer，更具體地說是 Google 的 SigLIP 視覺編碼器。我們的語言主幹遵循 Llama 3 架構。

視覺和文字模態透過模態投影模組進行對齊。該模組將視覺主幹生成的影像嵌入作為輸入，並將其轉換為與語言模型的嵌入層中的文字嵌入相容的嵌入。然後將這些嵌入連線起來並饋送到語言解碼器。模態投影模組由畫素混洗操作和線性層組成。

畫素混洗減少了影像標記的數量，這有助於降低計算成本並加快訓練速度，特別是對於對輸入長度敏感的基於 Transformer 的語言解碼器。下圖演示了這一概念。

所有檔案都非常輕量且文件齊全。我們強烈建議您單獨檢視它們，以更好地瞭解實現細節 (models/xxx.py)。

訓練時，我們使用以下預訓練主幹權重：

1. 視覺主幹：google/siglip-base-patch16-224
2. 語言主幹：HuggingFaceTB/SmolLM2-135M

還可以將主幹替換為 SigLIP/SigLIP 2（用於視覺主幹）和 SmolLM2（用於語言主幹）的其他變體。

訓練您自己的 VLM

現在我們熟悉了架構，接下來讓我們談談如何使用 train.py 訓練您自己的視覺語言模型。

.
├── data
│   └── ...
├── models
│   └── ...
└── train.py     # 👈 You are here

您可以開始訓練：

python train.py

此指令碼是您整個訓練管道的一站式商店，包括：

• 資料集載入和預處理
• 模型初始化
• 最佳化和日誌記錄

配置

在此之前，指令碼會從 models/config.py 載入兩個配置類：

• TrainConfig：用於訓練的配置引數，例如學習率、檢查點路徑等。
• VLMConfig：用於初始化 VLM 的配置引數，例如隱藏維度、注意力頭數量等。

資料載入

資料管道的核心是 get_dataloaders 函式。它：

• 透過 Hugging Face 的 load_dataset API 載入資料集。
• 組合並打亂多個數據集（如果提供）。
• 透過索引應用訓練/驗證分割。
• 將它們封裝在自定義資料集（VQADataset、MMStarDataset）和整理器（VQACollator、MMStarCollator）中。

此處有一個有用的標誌是 data_cutoff_idx，它在小資料集上除錯時非常有用。

模型初始化

模型透過 VisionLanguageModel 類構建。如果您要從檢查點恢復，就像這樣簡單：

from models.vision_language_model import VisionLanguageModel

model = VisionLanguageModel.from_pretrained(model_path)

否則，您將獲得一個全新初始化的模型，可以選擇預載入視覺和語言主幹。

最佳化器設定：兩個學習率

由於模態投影器 (MP) 是全新初始化的，而主幹是預訓練的，因此最佳化器被分為兩個引數組，每個組都有自己的學習率：

• MP 的學習率更高
• 編碼器/解碼器堆疊的學習率更小

這種平衡確保 MP 快速學習，同時保留視覺和語言主幹中的知識。

訓練迴圈

這部分是相當標準的，但結構經過精心設計：

• 使用 torch.autocast 進行混合精度以提高效能。
• 透過 get_lr 實現帶有線性預熱的餘弦學習率排程。
• 每批次的 token 吞吐量（token/秒）被記錄下來，用於效能監控。

每隔 250 步（可配置），模型在驗證和 MMStar 測試資料集上進行評估。如果準確率提高，模型將儲存檢查點。

日誌和監控

如果啟用 log_wandb，訓練統計資料（如 batch_loss、val_loss、accuracy 和 tokens_per_second）將記錄到 Weights & Biases，用於即時跟蹤。

執行會自動命名，使用樣本大小、批大小、epoch 計數、學習率和日期等元資料，所有這些都由輔助函式 get_run_name 處理。

推送到 Hub

使用以下命令將訓練好的模型推送到 Hub，供他人查詢和測試：

model.save_pretrained(save_path)

您可以輕鬆地使用以下命令推送它們：

model.push_to_hub("hub/id")

在預訓練模型上執行推理

我們使用 nanoVLM 作為工具包，訓練了一個模型併發布到 Hub。我們使用 google/siglip-base-patch16-224 和 HuggingFaceTB/SmolLM2-135M 作為主幹。該模型在單個 H100 GPU 上訓練了約 6 小時，使用了約 170 萬個 cauldron 樣本。

此模型並非旨在與 SoTA 模型競爭，而是旨在揭開 VLM 元件和訓練過程的神秘面紗。

.
├── data
│   └── ...
├── generate.py     # 👈 You are here
├── models
│   └── ...
└── ...

讓我們使用 generate.py 指令碼在訓練好的模型上執行推理。您可以使用以下命令執行生成指令碼：

python generate.py

這將使用預設引數並在影像 assets/image.png 上執行查詢“這是什麼？”。
您可以在您自己的影像和提示上使用此指令碼，如下所示：

python generate.py --image path/to/image.png --prompt "You prompt here"

如果您想視覺化指令碼的核心，它就是這些行：

model = VisionLanguageModel.from_pretrained(source).to(device)
model.eval()

tokenizer = get_tokenizer(model.cfg.lm_tokenizer)
image_processor = get_image_processor(model.cfg.vit_img_size)

template = f"Question: {args.prompt} Answer:"
encoded = tokenizer.batch_encode_plus([template], return_tensors="pt")
tokens = encoded["input_ids"].to(device)

img = Image.open(args.image).convert("RGB")
img_t = image_processor(img).unsqueeze(0).to(device)

print("\nInput:\n ", args.prompt, "\n\nOutputs:")
for i in range(args.generations):
    gen = model.generate(tokens, img_t, max_new_tokens=args.max_new_tokens)
    out = tokenizer.batch_decode(gen, skip_special_tokens=True)[0]
    print(f"  >> Generation {i+1}: {out}")

我們建立模型並將其設定為 eval。初始化分詞器，用於對文字提示進行分詞，以及影像處理器，用於處理影像。下一步是處理輸入並執行 model.generate 以生成輸出文字。最後，使用 batch_decode 對輸出進行解碼。

如果您想在使用者介面中對訓練好的模型進行推理，這裡是 Hugging Face Space，供您與模型互動。

結論

在這篇部落格文章中，我們詳細介紹了 VLM 是什麼，探討了 nanoVLM 的架構選擇，並詳細闡述了訓練和推理工作流程。

透過保持程式碼庫的輕量化和可讀性，nanoVLM 旨在作為一種學習工具和您可以在此基礎上構建的基礎。無論您是想了解多模態輸入如何對齊，還是想在自己的資料集上訓練 VLM，這個倉庫都能為您提供一個良好的開端。

如果您嘗試使用它，在此基礎上進行構建，或者有任何問題，我們都期待您的反饋。祝您玩得愉快！

參考文獻

1. GitHub - huggingface/nanoVLM: 用於訓練/微調小型 VLM 的最簡單、最快速的倉庫。
2. 視覺語言模型 (更好、更快、更強)
3. 視覺語言模型詳解
4. 深入探究視覺語言模型
5. SmolVLM：重新定義小型高效多模態模型