Hugging Face 发布 nanoVLM：纯 PyTorch 训练视觉语言模型的最简仓库

nanoVLM 是使用纯 PyTorch 训练你自己的视觉语言模型（VLM）的最简单方式。它是一个轻量级工具包，让你可以在一个免费层级的 Colab 笔记本上启动 VLM 训练。

我们受到 Andrej Karpathy 的 nanoGPT 启发，为视觉领域提供了一个类似的项目。

nanoVLM 的核心是一个工具包，帮助你构建和训练一个既能理解图像又能理解文本，并在此基础上生成文本的模型。nanoVLM 的美妙之处在于它的简洁性。整个代码库被有意保持精简和可读，非常适合初学者或任何想要一窥 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 架构。

视觉和文本模态通过一个模态投影模块进行对齐。该模块以视觉骨干网络生成的图像嵌入作为输入，并将其转换为与语言模型嵌入层输出的文本嵌入兼容的嵌入。这些嵌入随后被拼接并送入语言解码器。模态投影模块由一个像素重排操作和一个线性层组成。


模型架构（来源：作者）

像素重排减少了图像 token 的数量，这有助于降低计算成本并加速训练，尤其是对于对输入长度敏感的基于 Transformer 的语言解码器。下图展示了这一概念。


像素重排可视化（来源：作者）

所有文件都非常轻量且文档齐全。我们强烈建议你逐一查看它们，以更好地了解实现细节（models/xxx.py）

在训练过程中，我们使用以下预训练骨干网络权重：

视觉骨干网络：google/siglip-base-patch16-224
语言骨干网络：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 函数。它会：

通过 HF Mirror 的 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 吞吐量（tokens/秒）以监控性能。

每 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 上基于cauldron 的约 170 万条样本训练了约 6 小时。

该模型并非旨在与 SoTA 模型竞争，而是为了揭开 VLM 的组件和训练过程的神秘面纱。

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

让我们使用 generate.py 脚本在训练好的模型上运行推理。你可以使用以下命令运行生成脚本：

python generate.py

这将使用默认参数，并对图像 assets/image.png 运行查询“What is this?”。
你可以像这样在你自己的图像和提示上使用此脚本：

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 解码输出。

图像	提示	生成
	这是什么？	在图片中我能看到粉色的床单。我能看到两只猫躺在床单上。
	那个女人在做什么？	在这里，她正在中间做瑜伽

如果你想在 UI 界面中对训练好的模型运行推理，这里是供你与模型交互的 HF Mirror Space。

结论

在这篇博客文章中，我们介绍了什么是 VLM，探讨了驱动 nanoVLM 的架构选择，并详细解析了训练和推理工作流程。

通过保持代码库轻量且易读，nanoVLM 旨在既作为学习工具，也作为你可以在此基础上构建的基础。无论你是想了解多模态输入是如何对齐的，还是想在自己的数据集上训练 VLM，这个仓库都能为你提供一个良好的起点。

如果你尝试了它、在它之上进行构建，或者只是有问题，我们很乐意听到你的声音。祝你玩得开心！