NVIDIA TAO 5.5 帶來新基礎模型和增強訓練功能

TAO 與 NVIDIA 硬件和軟件生態系統無縫集成，提供用于高效 AI 模型訓練、部署和推理的工具，并加速 AI 驅動應用程序的上市時間。

圖 1 顯示 TAO 支持 PyTorch、TensorFlow 和 ONNX 等框架。訓練可以在多個平臺上完成，而生成的模型可以部署在 GPU、CPU、MCU 和 DLA 等各種推理平臺上。

NVIDIA 剛剛發布了 TAO 5.5，引入了先進的基礎模型和突破性功能，可增強任何 AI 模型開發。新的功能包括以下內容：

• 多模態傳感器融合模型：將來自多個傳感器的數據集成到統一的鳥瞰圖（BEV）表示中，同時保留幾何和語義信息。
• 帶有文本提示的自動標記：使用文本提示自動創建用于對象檢測和分割的標簽數據集。
• 開放詞匯檢測：使用自然語言描述（而非預定義標簽）識別任何類別中的對象。
• 知識提煉：根據大型網絡的知識創建更小、更高效和更準確的網絡。

在本文中，我們將更詳細地討論 TAO 5.5 的新功能。

NVIDIA TAO 集成了開源、基礎和專有模型，所有模型都在大量專有和商用數據集上進行訓練，使其能夠通用于物體檢測、姿態檢測、圖像分類、分割等任務。TAO 可針對特定用例簡化這些模型的微調，從而輕松進行自定義和商用部署。

大多數模型由 NVIDIA TensorRT 加速，并針對 NVIDIA 硬件的性能進行了優化，從而確保提供強大、高效的 AI 解決方案。

在 TAO 中交換模型主干非常簡單，無需編碼，只需簡單的配置更改即可。這使您能夠嘗試不同的架構，例如 ResNet、Shifted Window (Swin) Transformer 和 Fully Attentional Network (FAN)，從而根據特定需求定制模型。

這種易于定制的方式支持多種應用，例如零售業中的產品識別、醫療保健行業中的醫學影像分類、制造業中的機器人裝配監控以及智能城市中的交通管理。

基礎 DINO 模型

傳統的物體檢測模型僅限于識別預定義類別（閉集檢測）中的物體。這項約束使得這些模型無法用于需要根據人類輸入識別任意物體的應用程序，例如特定類別名稱或詳細的引用表達式。

GroundingDINO 通過將文本編碼器集成到 DINO 模型中，將其轉換為開放式物體檢測器，以解決這一限制。這使模型能夠檢測人類輸入所描述的任何物體。

GroundingDINO 通過使用特征增強器、語言引導的查詢選擇和跨模式解碼器，有效地融合了語言和視覺模式。這使模型能夠泛化預定義類別以外的概念，實現卓越的性能。

將 DINO 與 TAO 相結合

此模型使用 Swin-Tiny 主干進行訓練，在 NVIDIA 專有數據集上使用監督式方式進行商業用途。此外，BERT-Base 用作文本塔的起始權重。最后，GroundingDINO 使用從公開數據集收集的約 180 萬張圖像進行端到端訓練。

在訓練期間使用的所有原始圖像均具有商業許可證，以確保安全的商業用途。

Mask-GroundingDINO 模型

Mask-GroundingDINO 是一個單階段開放詞匯實例分割模型，可圍繞對象的特定實例生成分割掩膜。該模型基于 GroundingDINO 架構構建，并且 Conditional Convolutions for Instance Segmentation (CondInst) 激發了分割圖的靈感。

使用 TAO 對 DINO 進行遮罩訓練

CondInst 提出了用于實例分割的條件卷積頭，該卷積頭可根據輸入或特征圖更新其卷積核權重。我們將最初僅為 CondInst 中的 CNN 設計的掩膜分支和實例感知型掩膜頭擴展到 Transformer 或基于查詢的模型。

在 Mask-GroundingDINO 中，我們選擇使用與 DINO 中相同的編碼器和解碼器設計。然而，也可以與其他 Transformer 編碼器和解碼器建立類似的連接。

• Input: 文本或圖片。
  • 文本可以是類別列表（box、forklift），也可以是自然語言中的場景/物體描述。
• Backbone: 文本或圖片。
  • 文本主干可以是 BERT 或其他 NLP 模型。
  • 圖像主干可以是Swin Transformer模型，也可以是其他基于Transformer或CNN的主干。
• 特征增強器：自注意力層、文本到圖像交叉注意力層和圖像到文本交叉注意力層的堆棧，可以促進多模態對齊。
• 查詢選擇模塊：根據文本特征和圖像特征之間的點積選擇熱門的 K 個查詢。
• Mask FCN 頭：一個具有動態內核的幾個卷積層。
• 控制器：用于預測Mask FCN頭中動態內核權重的幾個卷積層。

該模型在偽標記的 Open Images 數據集上使用商用 GroundingDINO 預訓練模型進行了微調，可用于商業用途。該模型使用 Swin-Tiny 主干，其分割頭在大約 83 萬張帶有偽標記真值遮罩的圖像上進行了端到端微調。

您可以原樣使用此模型，也可以使用 TAO 對其進行微調。如果需要進行微調，我們提供了 Mask-GroundingDINO Jupyter 筆記，以提供簡單的交互式環境，并盡可能減少編碼需求。

BEVFusion 模型

從自動駕駛到機器人開發和智慧城市，在許多領域中，系統都依賴于各種傳感器來感知環境并與其交互。每種傳感器類型（例如攝像頭、激光雷達或雷達）都提供獨特的信息，但也存在內在的局限性。例如，攝像頭捕獲豐富的視覺細節，但難以處理深度感知，而激光雷達在測量距離方面表現優異，但缺乏語義上下文。

挑戰在于如何有效地將這些不同的數據源整合起來，以提高系統的性能和可靠性。

通過將來自多個傳感器的數據集成到統一的鳥瞰圖（BEV）表示中，BEVFusion 為應對這一挑戰提供了解決方案。它通過將多模態特征統一為共享的鳥瞰圖表示，從而徹底改變這一過程，同時保留（來自激光雷達的）幾何信息和（來自攝像頭的）語義信息。

與 TAO 的 BEVFusion

在 TAO 中，我們從 mmdet3d 的?BEVFusion?代碼庫開始訓練模型。原始代碼經過增強，可在 3D 空間中處理三個角度（roll、pitch、yaw），而最初僅支持 yaw。經過 TAO 訓練的 BEVFusion 模型可在單個攝像頭視圖和激光雷達中檢測人員，并在人員周圍創建 3D 邊界框。

CLIP 模型

高效地學習有意義且上下文豐富的圖像和文本表示是一項重大挑戰，對于需要多模態理解的應用尤其如此。零樣本學習、圖像字幕、視覺搜索、內容調節和多模態交互等用例需要能夠準確解釋和集成視覺和文本數據的強大模型。

Contrastive Language-Image Pre-training (CLIP) 模型采用雙編碼器架構來同時處理圖像和文本，從而解決了這一挑戰。CLIP 利用對比學習最大限度地提高了正確的圖像 – 文本對之間的相似性，同時最大限度地減少了錯誤對之間的相似性，使模型能夠學習捕獲各種概念的一般表示。

這種方法使 CLIP 能夠在各種應用程序中表現出色，例如生成描述性圖像描述、根據文本查詢執行視覺搜索，以及實現零樣本學習，模型可以利用其學習表示將圖像分類為未見類別。

使用 TAO 的 NVCLIP

在 TAO 中，我們在 NVIDIA 數據集上提供預訓練的 TensorRT 加速 CLIP 模型。我們使用了大約 700M 張圖像來訓練此模型。為進行評估，我們使用了 ImageNet 數據集的 50,000 張驗證圖像。

您可以將模型從 NVIDIA Jetson 部署到 NVIDIA Ampere 架構的 GPU 上。

SEGIC 模型

SEGIC 是一款創新性的端到端上下文中分割框架，可以借助單視覺基礎模型（VFM）徹底改變上下文中分割。

與傳統方法不同，SEGIC 捕獲目標圖像與上下文示例之間的密集關系，提取幾何、視覺和元指令以指導最終蒙版預測。這種方法顯著降低了標記和訓練成本，同時在一次性分割基準測試中實現了先進的性能。

SEGIC 的通用性擴展到各種任務，包括視頻對象分割和開放詞匯分割，使其成為圖像分割領域的強大工具。

采用 TAO 的 SEGIC

對于在 TAO 中訓練 SEGIC 模型，我們提供了一個預訓練的 ONNX 模型和一個使用 NVIDIA Triton 推理服務器的部署方法。對于訓練，我們使用了以下資源：

• 圖像編碼器：使用 NVIDIA 專有圖像訓練的 DINOv2-L。
• 元描述編碼器：CLIP-B，使用 NVIDIA 的專有圖像進行訓練。

基礎姿態模型

準確的 6D 物體姿態估計和跟蹤對于機器人、增強現實和自動駕駛等各種應用至關重要。然而，現有方法通常需要大量微調，并且受其對基于模型或無模型設置的依賴程度的限制。這給在不同場景和對象中實現穩健性能帶來了挑戰，尤其是在 CAD 模型不可用或處理新對象時。

FoundationPose 通過為 6D 對象姿態估計和跟蹤提供統一的基礎模型來解決此問題，該模型支持基于模型和無模型的設置。它可以在測試時立即應用于新對象，利用 CAD 模型或一些參考圖像。

通過將神經隱式表示用于新的視圖合成，并利用LLM（一種基于Transformer的新型架構）和對比學習等先進技術進行大規模合成訓練，FoundationPose實現了強大的可泛化性，并在各種公共數據集中優于專業方法。

例如，在增強現實應用中，FoundationPose 無需大量手動調整，即可準確估計和跟蹤物體姿態，從而將虛擬物體無縫集成到真實環境中。

基礎姿勢與 TAO

我們使用 NVIDIA Triton 推理服務器提供了預訓練的 ONNX 模型和部署方案。我們的訓練數據由使用 Google Research 和 Objaverse 掃描對象中的 3D 資產以高質量逼真度渲染的場景組成，每個數據點都包括 RGB 圖像、深度信息、物體姿態、攝像頭姿態、實例分割和 2D 邊界框，并具有廣泛的域隨機化。然后，我們使用 NVIDIA Isaac Sim 創建了約 100 萬張合成圖像。

在全球 BOP 排行榜中，截至 2024 年 3 月，基于模型的新型物體姿態估計排名第一（圖 3）。

基于提示的自動標注，用于對象檢測和分割。

擁有良好標注的數據集對于訓練和微調模型至關重要，尤其是在物體檢測和分割等任務中。然而，為新類別和實例獲取已標注的數據集需要花費大量時間和精力。這對于實例分割而言尤其如此，其中使用精確的蒙版注釋每個對象可能比繪制簡單的邊界框花費長達10倍。

這就是一個簡單易用、基于提示的自動標記方法可以幫助您

• GroundingDINO：開放詞匯的物體檢測模型，使用物體類別等提示生成邊界框。
• Mask Auto-Labeler：基于 Transformer 的框架，使用邊界框生成高質量的實例分割蒙版。

這種組合可顯著減少創建詳細標簽所需的工作量，從而更輕松、更快速地構建用于訓練高級模型的可靠數據集。

我們提供用于自動標記的完整端到端 Jupyter Notebook (text2box.ipynb)，無需編碼。我們提供兩個規格文件，一個用于 bbox 標簽，另一個用于分割，您可以在其中定義要標記的對象，例如 [person, helmet]，未標記數據集的路徑以及用于存儲標簽的結果目錄。

使用 generate 命令生成自動標記：

借助知識提煉實現高效 AI

知識提煉是機器學習中的一種技術，在該技術中，更小、更高效的模型（學生）學習模仿更大、更復雜的模型（教師）的行為。這一過程有助于減少訓練和微調時間，因為學生模型使用教師已經捕獲的知識（包括軟標簽等更細致入微的信息），而無需為任何新任務從頭開始訓練。

通過提煉這些知識，學生模型可以在顯著減少計算資源的情況下實現類似的性能，使其成為在資源受限的環境中部署的理想選擇，并加速訓練過程。

TAO 支持知識提煉。指定不同層之間的綁定，以計算提煉損失。有關更多信息，請參閱 Optimizing the Training Pipeline。

我們提供了 dino_distillation.ipynb 參考 notebook，以展示知識提煉，以及如何使用 L2 損失從 DINO 和 FAN 小至 DINO 和 ResNet-50 提取中間特征圖。

在 distill.yaml 文件中，您可以定義教師模型、學生模型主干和損失函數。訓練模型：

在訓練學生模型后，您可以使用 TAO evaluate 工具評估提煉模型，使用 inference 工具可視化推理，使用 export 工具將訓練好的模型導出到 ONNX，并最終生成 gen_trt_engine 文件以將 DINO 和 ResNet-50 模型導出到 TensorRT。

TAO 5.5 入門

全球各地的開發者正在使用 NVIDIA TAO 加速其視覺 AI 應用的 AI 訓練。使用 TAO 5.5 的新功能來增強您的 AI 應用。

有關更多信息，請參閱以下資源：

• 下載 NVIDIA TAO
• 開始使用 NVIDIA TAO
• 在論壇中提出技術問題