nanoVLM：用于训练视觉语言模型的极简 PyTorch 工具包

Hugging Face 推出了 nanoVLM，这是一款用纯 PyTorch 编写的极简工具包，旨在揭开训练视觉语言模型（VLM）过程的神秘面纱。受 Andrej Karpathy 的 nanoGPT 启发，nanoVLM 提供了可读、轻量的代码库，使用户能够在免费版 Google Colab 笔记本上训练 VLM。

视觉语言模型（VLM）基础

视觉语言模型是一种多模态架构，能够处理图像和文本输入并生成文本输出。虽然 VLM 可用于图像描述、目标检测、语义分割等多种任务，但 nanoVLM 专注于 视觉问答（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）策略：

  • 更高的 LR： 应用于模态投影器（MP），以促进快速学习。
  • 更低的 LR： 应用于编码器/解码器堆栈，以保留主干中已有的知识。

训练循环与监控

训练使用 torch.autocast 实现混合精度，并采用余弦学习率调度结合线性预热。通过 token 吞吐量（tokens/sec）进行性能监控，并在启用时使用 Weights & Biases（wandb）跟踪批次损失、验证损失和准确率。

推理与预训练模型

Hugging Face 提供了已发布到 Hub 的预训练 nanoVLM 模型（nanoVLM-222M）。该模型使用 cauldron 数据集的 170 万样本，在单个 H100 GPU 上训练约 6 小时。

用户可以使用 generate.py 脚本进行推理，其逻辑流程如下：

  1. 初始化： 加载模型、分词器和图像处理器。
  2. 处理： 对文本提示进行分词，并将图像处理为张量。
  3. 生成： 执行 model.generate 生成文本输出。
  4. 解码： 使用 batch_decode 将生成的 token 转换回可读的文本。

入门指南

要开始训练，用户可以克隆仓库并运行训练脚本：

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

# Execute the training script
python train.py

Sources