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：重新定义小型高效多模态模型