nanoVLM 文档解读

一、引言

nanoVLM 是一个用于训练和微调小型视觉语言模型（VLM）的极简代码库。它采用纯 PyTorch 实现，代码简洁易读，旨在为社区提供一个简单易用的 VLM 实现和训练脚本。该框架类似于 Andrej Karpathy 的 nanoGPT，强调教育价值而非追求最新的 SOTA（State of The Art）性能。

二、模型组成

nanoVLM 的模型结构包含以下组件：

1. 视觉骨干网络 ：位于 models/vision_transformer.py，代码行数约 150 行。

2. 语言解码器 ：位于 models/language_model.py，代码行数约 250 行。

3. 模态投影 ：位于 models/modality_projection.py，代码行数约 50 行。

4. VLM 本身 ：位于 models/vision_language_model.py，代码行数约 100 行。

5. 训练循环 ：位于 train.py，代码行数约 200 行。

整个模型定义和训练逻辑代码约 750 行，还包括一些日志记录和参数加载的样板代码。

三、性能表现

使用 SigLIP-B/16-224（85M 参数）和 HuggingFaceTB/SmolLM2-135M 作为骨干网络，构建出一个 222M 参数的 nanoVLM。在单个 H100 GPU 上对约 1.7M 样本的 cauldron 数据集训练约 6 小时后，在 MMStar 数据集上达到 35.3% 的准确率。

四、快速入门

1. 环境搭建

   • 可以使用 uv 作为包管理器，也可自行安装依赖包，包括 torch、numpy、torchvision、pillow、datasets、huggingface-hub、transformers 和 wandb。

   • 通过 git 克隆仓库，使用 uv 或 pip 安装依赖。

2. 训练

   • 使用提供的训练脚本 train.py，默认使用 models/config.py 中的配置。

   • 需先登录 wandb，然后运行 python train.py 开始训练。

3. 推理

   • 使用 generate.py 脚本进行推理。例如，将 assets/image.png 中的图片与问题输入模型，可得到关于图片中猫的描述等输出结果。

五、Hugging Face Hub 集成

1. 加载预训练权重

   • 推荐从 Hugging Face Hub 加载预训练权重，通过 VisionLanguageModel.from_pretrained(“lusxvr/nanoVLM-222M”) 实现。

2. 推送模型至 Hub

   • 训练完成后，可使用 model.push_to_hub(“my-awesome-nanovlm-model”) 将模型推送至 Hub。模型会保存为 config.json（配置文件）、model.safetensors（权重文件），并生成 README.md（模型说明文件）。

3. 本地保存与加载

   • 可通过 model.save_pretrained(“path/to/local/model”) 保存模型到本地文件夹，再使用 VisionLanguageModel.from_pretrained(“path/to/local/model”) 从本地路径加载模型。

六、硬件要求（VRAM 使用情况）

对默认的 nanoVLM 模型（222M 参数）在单个 NVIDIA H100 GPU 上进行测试，不同批次大小的峰值 VRAM 使用情况如下表：

关键结论：

• 即使批次大小为 1，训练默认模型至少需要约 4.5GB VRAM。

• 拥有约 8GB VRAM 时，批次大小可达到 16。

七、贡献指南

1. 纯 PyTorch 实现

   • nanoVLM 致力于保持纯 PyTorch 的轻量级实现。不接受引入 transformers.Trainer、accelerate 或 deepspeed 等依赖的贡献。

2. 新功能开发

   • 若有新功能想法，需先提交问题进行讨论，以确保贡献与项目目标一致。

3. 漏洞修复

   • 鼓励提交修复漏洞的 pull 请求。

八、未来发展计划

项目团队希望在以下方面进行改进和扩展，欢迎相关领域的贡献：

1. 数据打包 ：实现从输入数据创建指定大小的数据包，以优化训练过程。

2. 多 GPU 训练 ：支持在多个 GPU 上进行训练。

3. 多图像支持 ：实现使用多张图像进行训练。

4. 图像分割 ：借鉴 SmolVLM 的方法，通过图像分割支持更高分辨率。

5. VLMEvalKit 集成 ：将 nanoVLM 集成到 VLMEvalKit 中，以便进行更多基准测试。