如何在NVIDIA Jetson AGX Orin開發(fā)工具包上訓練nanoVLM
在這篇博文中,我將向您展示如何在NVIDIA Jetson AGX Orin開發(fā)工具包上訓練hug Face的納米ovlm。
視覺語言模型(vlm)是功能強大的人工智能模型,可以從圖像和文本中學習,使它們能夠執(zhí)行視覺問題回答、圖像字幕和多模態(tài)推理等任務。
大多數(shù)vlm的總體結(jié)構(gòu)如下:
目前,llm和vlm正在Nvidia/AMD gpu的巨大集群上進行訓練。這就提出了一個問題:個人如何使用消費級硬件訓練模型?
所以,我們很幸運,hugs Face最近開源了nanoVLM項目,專門為設(shè)備上的訓練和推理設(shè)計了一個輕量級的純PyTorch實現(xiàn)。我們可以訓練一個小規(guī)模的VLM (222M參數(shù)),它結(jié)合了一個ViT編碼器(谷歌/siglip-base-patch16-224)和一個LLM (HuggingFaceTB/SmolLM2-135M)來理解文本和視覺輸入。它使用HuggingFaceM4/the_cauldron數(shù)據(jù)集的數(shù)據(jù)。支持檢查點,能夠從最近的檢查點恢復。nanoVLM:用純PyTorch訓練VLM的最簡單的存儲庫。
本教程將指導您完成使用Kubernetes直接在NVIDIA Jetson AGX Orin Developer Kit上運行納米ovlm模型訓練的過程。
在Nvidia Jetson AGX Orin上設(shè)置k3
使用NVIDIA容器運行時支持安裝K3s:
這個命令安裝K3s,配置它使用Docker,設(shè)置kubecconfig文件權(quán)限,并禁用Traefik?!猟ocker標志很重要,因為它確保K3s使用docker,它與NVIDIA容器運行時兼容。
驗證K3s集群是否正常運行:
您應該看到Jetson AGX Orin被列為集群中的一個節(jié)點。
這確認您的單節(jié)點K3s集群是可操作的。
讓我們安裝k9,類似于Lens,它是一個通過CLI用戶界面的Kubernetes集群管理工具。我們可以使用下面的命令安裝k9s:
開展培訓工作
我們現(xiàn)在準備開始訓練一個模型!
創(chuàng)建一個名為train的清單文件。并使用以下代碼填充它:
由于backoffLimit: 4, Kubernetes將重試pod最多4次。這對于控制批處理作業(yè)(如訓練循環(huán))中的容錯性非常有用,因為這些批處理作業(yè)可能由于暫時問題而失敗。
要啟動作業(yè),運行命令:
有時豆莢不會立即啟動。因此,您將在“READY”列中看到0/1,在“STATUS”列中看到Pending。如果你想查看你的pod的詳細狀態(tài),運行這個命令:
我們還可以用k9檢查NanoVLM吊艙。
您應該看到類似于下面的輸出:
檢查任何啟動問題的日志。
記住將POD_NAME替換為訓練pod的實際名稱。
一旦培訓結(jié)束,您的工作狀態(tài)將顯示“完成”。下面是輸出的一個例子:
您還將看到類似于下面的輸出,顯示訓練進度:
日志還將提供對訓練進度的洞察,包括模型加載細節(jié)、參數(shù)計數(shù)和性能度量。
來自wandb的趨勢將提供這些指標隨時間變化的可視化表示。如果精度提高了,模型將被檢查點并保存在本地。
這種特殊的訓練持續(xù)了大約1天20小時,完成了5個階段。MMStar的精度為0.1013。這不是為了達到最好的模型精度,而是為了驗證Nvidia Jetson Orin Dev Kit可以使用Kubernetes設(shè)置處理訓練過程的計算需求。
訓練成功后,我們可以將您新訓練的模型上傳到“擁抱臉”,以便分享和將來使用。使用以下Python代碼片段,確保你的檢查點路徑是正確的:
接下來,使用上傳的模型進行推理,執(zhí)行以下命令,如果不相同,將shakhizat/nanoVLM-222M替換為實際的模型ID:
這是我們的演示圖像:
下面是使用輸入圖像運行推理時預期輸出的示例:
對于生產(chǎn)就緒的模型來說,這不是一個好結(jié)果。然而,這個特殊練習的主要目標并不是達到最先進的精度。相反,這項培訓工作是一個至關(guān)重要的可行性檢查,看看我們是否可以成功地在NVIDIA Jetson AGX Orin開發(fā)工具包上運行NanoVLM模型培訓過程。這一成功的概念驗證為進一步優(yōu)化和開發(fā)設(shè)備上VLM培訓打開了大門。
本文編譯自hackster.io