如何在NVIDIA Jetson AGX Orin开发工具包上训练nanoVLM

在这篇博文中，我将向您展示如何在NVIDIA Jetson AGX Orin开发工具包上训练hug Face的纳米ovlm。

视觉语言模型(vlm)是功能强大的人工智能模型，可以从图像和文本中学习，使它们能够执行视觉问题回答、图像字幕和多模态推理等任务。

大多数vlm的总体结构如下：

目前，llm和vlm正在Nvidia/AMD gpu的巨大集群上进行训练。这就提出了一个问题：个人如何使用消费级硬件训练模型?

所以，我们很幸运，hugs Face最近开源了nanoVLM项目，专门为设备上的训练和推理设计了一个轻量级的纯PyTorch实现。我们可以训练一个小规模的VLM (222M参数)，它结合了一个ViT编码器(谷歌/siglip-base-patch16-224)和一个LLM (HuggingFaceTB/SmolLM2-135M)来理解文本和视觉输入。它使用HuggingFaceM4/the_cauldron数据集的数据。支持检查点，能够从最近的检查点恢复。nanoVLM：用纯PyTorch训练VLM的最简单的存储库。

本教程将指导您完成使用Kubernetes直接在NVIDIA Jetson AGX Orin Developer Kit上运行纳米ovlm模型训练的过程。

在Nvidia Jetson AGX Orin上设置k3

使用NVIDIA容器运行时支持安装K3s：

这个命令安装K3s，配置它使用Docker，设置kubecconfig文件权限，并禁用Traefik。——docker标志很重要，因为它确保K3s使用docker，它与NVIDIA容器运行时兼容。

验证K3s集群是否正常运行：

您应该看到Jetson AGX Orin被列为集群中的一个节点。

这确认您的单节点K3s集群是可操作的。

让我们安装k9，类似于Lens，它是一个通过CLI用户界面的Kubernetes集群管理工具。我们可以使用下面的命令安装k9s：

开展培训工作

我们现在准备开始训练一个模型!

创建一个名为train的清单文件。并使用以下代码填充它：

由于backoffLimit: 4， Kubernetes将重试pod最多4次。这对于控制批处理作业(如训练循环)中的容错性非常有用，因为这些批处理作业可能由于暂时问题而失败。

要启动作业，运行命令：

有时豆荚不会立即启动。因此，您将在“READY”列中看到0/1，在“STATUS”列中看到Pending。如果你想查看你的pod的详细状态，运行这个命令：

我们还可以用k9检查NanoVLM吊舱。

您应该看到类似于下面的输出：

检查任何启动问题的日志。

记住将POD_NAME替换为训练pod的实际名称。

一旦培训结束，您的工作状态将显示“完成”。下面是输出的一个例子：

您还将看到类似于下面的输出，显示训练进度：

日志还将提供对训练进度的洞察，包括模型加载细节、参数计数和性能度量。

来自wandb的趋势将提供这些指标随时间变化的可视化表示。如果精度提高了，模型将被检查点并保存在本地。

这种特殊的训练持续了大约1天20小时，完成了5个阶段。MMStar的精度为0.1013。这不是为了达到最好的模型精度，而是为了验证Nvidia Jetson Orin Dev Kit可以使用Kubernetes设置处理训练过程的计算需求。

训练成功后，我们可以将您新训练的模型上传到“拥抱脸”，以便分享和将来使用。使用以下Python代码片段，确保你的检查点路径是正确的：

接下来，使用上传的模型进行推理，执行以下命令，如果不相同，将shakhizat/nanoVLM-222M替换为实际的模型ID：

这是我们的演示图像：

下面是使用输入图像运行推理时预期输出的示例：

对于生产就绪的模型来说，这不是一个好结果。然而，这个特殊练习的主要目标并不是达到最先进的精度。相反，这项培训工作是一个至关重要的可行性检查，看看我们是否可以成功地在NVIDIA Jetson AGX Orin开发工具包上运行NanoVLM模型培训过程。这一成功的概念验证为进一步优化和开发设备上VLM培训打开了大门。

本文编译自hackster.io