nanoVLM: 最简洁、最轻量的纯 PyTorch 视觉-语言模型训练代码库
<p><strong>nanoVLM</strong> 是使用纯 PyTorch <strong>训练</strong> 你自己的视觉语言模型 (VLM) 的 <em>最简单</em> 方式。它是一个轻量级 <em>工具包</em> ，让你可以在 免费的 Colab Notebook 上启动 VLM 训练。</p><blockquote>
<p>我们受到了 Andrej Karpathy 的 nanoGPT 的启发，为视觉领域提供了一个类似的项目。</p>
</blockquote>
<p>从本质上讲，nanoVLM 是一个 <strong>工具包</strong>，可以帮助你构建和训练一个能够理解图像和文本，并基于此生成文本的模型。nanoVLM 的魅力在于它的 <em>简洁性</em> 。整个代码库被有意保持 <em>最小化</em> 和 <em>可读性</em> ，使其非常适合初学者或任何想要深入了解 VLM 内部机制而不被复杂性淹没的人。</p>
<p>在这篇博客中，我们将介绍该项目背后的核心思想，并提供与代码库交互的简单方法。我们不仅会深入项目细节，还会将所有内容封装起来，让你能够快速上手。</p>
<h2 id="简要">简要</h2>
<p>你可以按照以下步骤使用我们的 nanoVLM 工具包开始训练视觉语言模型:</p>
<pre><code class="language-bash"># 克隆仓库
git clone https://github.com/huggingface/nanoVLM.git
# 执行训练脚本
python train.py
</code></pre>
<p>这里有一个 Colab Notebook，可以帮助你在无需本地设置的情况下启动训练运行！</p>
<h2 id="什么是视觉语言模型">什么是视觉语言模型？</h2>
<p>顾名思义，视觉语言模型 (VLM) 是一种处理两种模态的多模态模型: 视觉和文本。这些模型通常以图像和/或文本作为输入，生成文本作为输出。</p>
<p>基于对图像和文本 (输入) 的理解来生成文本 (输出) 是一个强大的范式。它支持广泛的应用，从图像字幕生成和目标检测到回答关于视觉内容的问题 (如下表所示)。需要注意的是，nanoVLM 仅专注于视觉问答作为训练目标。</p>
<table>
<tbody><tr>
<td rowspan="4"><img alt=" 一张猫的图片 " width="200" data-src="https://img-s2.andfun.cn/devrel/posts/2025/06/7fc33aafb9355.jpg" class="lazyload"></td>
<td>为图像生成标题</td>
<td>两只猫躺在床上，旁边有遥控器</td>
<td>图像描述</td>
</tr>
<tr>
<td>检测图像中的物体</td>
<td><code><locxx><locxx><locxx><locxx></code></td>
<td>目标检测</td>
</tr>
<tr>
<td>分割图像中的物体</td>
<td><code><segxx><segxx><segxx></code></td>
<td>语义分割</td>
</tr>
<tr>
<td>图像中有多少只猫？</td>
<td>2</td>
<td>视觉问答</td>
</tr>
</tbody></table>
<blockquote>
<p>如果你有兴趣了解更多关于 VLM 的信息，我们强烈建议阅读我们关于该主题的最新博客: 视觉语言模型 (更好、更快、更强)</p>
</blockquote>
<h2 id="使用代码库">使用代码库</h2>
<p>“废话少说，直接看代码” - 林纳斯·托瓦兹</p>
<p>在本节中，我们将引导你了解代码库。在跟随学习时，保持一个 标签页 开启以供参考会很有帮助。</p>
<p>以下是我们仓库的文件夹结构。为简洁起见，我们删除了一些辅助文件。</p>
<pre><code class="language-bash">.
├── 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
</code></pre>
<h2 id="架构">架构</h2>
<pre><code class="language-bash">.
├── data
│ └── ...
├── models # 👈 你在这里
│ └── ...
└── train.py
</code></pre>
<p>我们按照两个知名且广泛使用的架构来建模 nanoVLM。我们的视觉主干网络 ( <code>models/vision_transformer.py</code> ) 是标准的视觉 transformer，更具体地说是谷歌的 SigLIP 视觉编码器。我们的语言主干网络遵循 Llama 3 架构。</p>
<p>视觉和文本模态通过模态投影模块进行 <em>对齐</em> 。该模块将视觉主干网络产生的图像嵌入作为输入，并将它们转换为与语言模型嵌入层的文本嵌入兼容的嵌入。然后将这些嵌入连接起来并输入到语言解码器中。模态投影模块由像素洗牌操作和线性层组成。</p>
<table>
<thead>
<tr>
<th style="text-align: center"><img alt="模型架构图" loading="lazy" src="https://img-s2.andfun.cn/devrel/posts/2025/06/2b2b5c8478a3b.png" class="lazyload"></th>
</tr>
</thead>
<tbody>
<tr>
<td style="text-align: center">模型架构 (来源: 作者)</td>
</tr>
</tbody>
</table>
<p>像素洗牌 减少了图像标记的数量，这有助于降低计算成本并加快训练速度，特别是对于对输入长度敏感的基于 transformer 的语言解码器。下图演示了这个概念。</p>
<table>
<thead>
<tr>
<th style="text-align: center"><img alt="像素洗牌图" loading="lazy" src="https://img-s2.andfun.cn/devrel/posts/2025/06/a1ad48fcd236c.png" class="lazyload"></th>
</tr>
</thead>
<tbody>
<tr>
<td style="text-align: center">像素洗牌可视化 (来源: 作者)</td>
</tr>
</tbody>
</table>
<p>所有文件都非常轻量且有良好的文档说明。我们强烈建议你逐个查看它们，以更好地理解实现细节 ( <code>models/xxx.py</code> )</p>
<p>在训练时，我们使用以下预训练的主干权重:</p>
<ol>
<li>视觉主干: <code>google/siglip-base-patch16-224</code></li>
<li>语言主干: <code>HuggingFaceTB/SmolLM2-135M</code></li>
</ol>
<blockquote>
<p>也可以将主干网络替换为 SigLIP/SigLIP 2 (用于视觉主干) 和 SmolLM2 (用于语言主干) 的其他变体。</p>
</blockquote>
<h2 id="训练你自己的-vlm">训练你自己的 VLM</h2>
<p>现在我们已经熟悉了架构，让我们换个话题，讨论如何使用 <code>train.py</code> 训练你自己的视觉语言模型。</p>
<pre><code class="language-bash">.
├── data
│ └── ...
├── models
│ └── ...
└── train.py # 👈 你在这里
</code></pre>
<p>你可以通过以下命令启动训练:</p>
<pre><code class="language-bash">python train.py
</code></pre>
<p>这个脚本是整个训练流程的一站式解决方案，包括:</p>
<ul>
<li>数据集加载和预处理</li>
<li>模型初始化</li>
<li>优化和日志记录</li>
</ul>
<p><strong>配置</strong></p>
<p>在任何其他操作之前，脚本从 <code>models/config.py</code> 加载两个配置类:</p>
<ul>
<li><code>TrainConfig</code> : 对训练有用的配置参数，如学习率、检查点路径等。</li>
<li><code>VLMConfig</code> : 用于初始化 VLM 的配置参数，如隐藏维度、注意力头数等。</li>
</ul>
<p><strong>数据加载</strong></p>
<p>数据流水线的核心是 <code>get_dataloaders</code> 函数。它:</p>
<ul>
<li>通过 Hugging Face 的 <code>load_dataset</code> API 加载数据集。</li>
<li>组合和洗牌多个数据集 (如果提供)。</li>
<li>通过索引应用训练/验证分割。</li>
<li>将它们包装在自定义数据集 (<code>VQADataset</code> 、<code>MMStarDataset</code> ) 和整理器 (<code>VQACollator</code> 、<code>MMStarCollator</code> ) 中。</li>
</ul>
<blockquote>
<p>这里一个有用的标志是 <code>data_cutoff_idx</code> ，对于在小子集上调试很有用。</p>
</blockquote>
<p><strong>模型初始化</strong></p>
<p>模型通过 <code>VisionLanguageModel</code> 类构建。如果你从检查点恢复，操作非常简单:</p>
<pre><code class="language-python">from models.vision_language_model import VisionLanguageModel
model = VisionLanguageModel.from_pretrained(model_path)
</code></pre>
<p>否则，你将获得一个全新初始化的模型，可选择为视觉和语言预加载主干网络。</p>
<p><strong>优化器设置: 两个学习率</strong></p>
<p>由于模态投影器 ( <code>MP</code> ) 是新初始化的，而主干网络是预训练的，优化器被分成两个参数组，每个都有自己的学习率:</p>
<ul>
<li>MP 使用较高的学习率</li>
<li>编码器/解码器堆栈使用较小的学习率</li>
</ul>
<p>这种平衡确保 MP 快速学习，同时保留视觉和语言主干网络中的知识。</p>
<p><strong>训练循环</strong></p>
<p>这部分相当标准但结构合理:</p>
<ul>
<li>使用 <code>torch.autocast</code> 进行混合精度以提高性能。</li>
<li>通过 <code>get_lr</code> 实现带线性预热的余弦学习率调度。</li>
<li>每批记录令牌吞吐量 (令牌/秒) 以进行性能监控。</li>
</ul>
<p>每 250 步 (可配置)，模型在验证集和 <code>MMStar</code> 测试数据集上进行评估。如果准确率提高，模型将被保存为检查点。</p>
<p><strong>日志记录和监控</strong></p>
<p>如果启用了 <code>log_wandb</code> ，训练统计信息如 <code>batch_loss</code> 、<code>val_loss</code> 、<code>accuracy</code> 和 <code>tokens_per_second</code> 将记录到 Weights & Biases 以进行实时跟踪。</p>
<p>运行使用元数据自动命名，如样本大小、批次大小、epoch 数、学习率和日期，全部由辅助函数 <code>get_run_name</code> 处理。</p>
<p><strong>推送到 Hub</strong></p>
<p>使用以下方法将训练好的模型推送到 Hub，供其他人查找和测试:</p>
<pre><code class="language-python">model.save_pretrained(save_path)
</code></pre>
<p>你可以轻松地使用以下方式推送它们:</p>
<pre><code class="language-python">model.push_to_hub("hub/id")
</code></pre>
<h2 id="在预训练模型上运行推理">在预训练模型上运行推理</h2>
<p>使用 nanoVLM 作为工具包，我们训练了一个 模型并将其发布到 Hub。我们使用了 <code>google/siglip-base-patch16-224</code> 和 <code>HuggingFaceTB/SmolLM2-135M</code> 作为主干网络。该模型在单个 H100 GPU 上对 cauldron 的约 170 万个样本训练了约 6 小时。</p>
<p>这个模型并不旨在与最先进的模型竞争，而是为了揭示 VLM 的组件和训练过程。</p>
<pre><code class="language-bash">.
├── data
│ └── ...
├── generate.py # 👈 你在这里
├── models
│ └── ...
└── ...
</code></pre>
<p>让我们使用 <code>generate.py</code> 脚本在训练好的模型上运行推理。你可以使用以下命令运行生成脚本:</p>
<pre><code class="language-bash">python generate.py
</code></pre>
<p>这将使用默认参数并在图像 <code>assets/image.png</code> 上运行查询 "What is this?"。</p>
<p>你可以在自己的图像和提示上使用此脚本，如下所示:</p>
<pre><code class="language-bash">python generate.py --image path/to/image.png --prompt "你的提示在这里"
</code></pre>
<p>如果你想可视化脚本的核心，就是这些行:</p>
<pre><code class="language-python">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(, 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)
print(f">> Generation {i+1}: {out}")
</code></pre>
<p>我们创建模型并将其设置为 <code>eval</code> 。初始化分词器 (用于对文本提示进行分词) 和图像处理器 (用于处理图像)。下一步是处理输入并运行 <code>model.generate</code> 以生成输出文本。最后，使用 <code>batch_decode</code> 解码输出。</p>
<table>
<thead>
<tr>
<th style="text-align: center">图像</th>
<th style="text-align: center">提示</th>
<th style="text-align: center">生成结果</th>
</tr>
</thead>
<tbody>
<tr>
<td style="text-align: center"><img alt="一张猫的图片" loading="lazy" src="https://img-s2.andfun.cn/devrel/posts/2025/06/7fc33aafb9355.jpg" class="lazyload"></td>
<td style="text-align: center">What is this?</td>
<td style="text-align: center">In the picture I can see the pink color bed sheet. I can see two cats lying on the bed sheet.</td>
</tr>
<tr>
<td style="text-align: center"><img alt="瑜伽" loading="lazy" src="https://img-s2.andfun.cn/devrel/posts/2025/06/66e65074428bb.png" class="lazyload"></td>
<td style="text-align: center">What is the woman doing?</td>
<td style="text-align: center">Here in the middle she is performing yoga</td>
</tr>
</tbody>
</table>
<blockquote>
<p>如果你想在 UI 界面中对训练好的模型运行推理，这里 有一个 Hugging Face Space 供你与模型交互。</p>
</blockquote>
<h2 id="结论">结论</h2>
<p>在这篇博客中，我们介绍了什么是 VLM，探讨了支撑 nanoVLM 的架构选择，并详细解释了训练和推理工作流程。</p>
<p>通过保持代码库轻量级和可读性，nanoVLM 旨在既作为学习工具，又作为你可以在此基础上构建的基础。无论你是想了解多模态输入如何对齐，还是想在自己的数据集上训练 VLM，这个仓库都能让你快速入门。</p>
<p>如果你尝试了它，并在它的基础上尝试构建，或者你只是有问题，我们都很乐意听到你的反馈。祝你探索愉快！</p>
<h2 id="参考文献">参考文献</h2>
<ol>
<li>GitHub - huggingface/nanoVLM: 用于训练/微调小型 VLM 的最简单、最快速的代码库。</li>
<li>视觉语言模型 (更好、更快、更强)</li>
<li>视觉语言模型详解</li>
<li>深入视觉语言预训练</li>
<li>SmolVLM: 重新定义小型高效多模态模型</li>
</ol>
<hr>
<blockquote>
<p>英文原文: https://hf.co/blog/nanovlm</p>
<p>原文作者: Aritra Roy Gosthipaty, Luis Wiedmann, Andres Marafioti, Sergio Paniego, Merve Noyan, Pedro Cuenca, Vaibhav Srivastav</p>
<p>翻译: innovation64</p>
</blockquote><br><br>
来源：https://www.cnblogs.com/huggingface/p/18935098
頁:
[1]