技术博客 | NVIDIA TAO 5.5 带来新基础模型和增强训练功能

NVIDIA TAO 是一个旨在简化和加速 AI 模型开发和部署的框架。它使您能够使用预训练模型，使用自己的数据微调模型，并针对特定用例优化模型，而无需深入的 AI 专业知识。

TAO 与 NVIDIA 硬件和软件生态系统无缝集成，提供用于高效 AI 模型训练、部署和推理的工具，并加速 AI 驱动应用程序的上市时间。

图 1 显示 TAO 支持 PyTorch、TensorFlow 和 ONNX 等框架。训练可以在多个平台上完成，而生成的模型可以部署在 GPU、CPU、MCU 和 DLA 等各种推理平台上。

NVIDIA 刚刚发布了 TAO 5.5，引入了先进的基础模型和突破性功能，可增强任何 AI 模型开发。新的功能包括以下内容：

• 多模态传感器融合模型：将来自多个传感器的数据集成到统一的鸟瞰图（BEV）表示中，同时保留几何和语义信息。

• 带有文本提示的自动标记：使用文本提示自动创建用于对象检测和分割的标签数据集。

• 开放词汇检测：使用自然语言描述（而非预定义标签）识别任何类别中的对象。

• 知识提炼：根据大型网络的知识创建更小、更高效和更准确的网络。

在本文中，我们将更详细地讨论 TAO 5.5 的新功能。

表 1：TAO 5.5 中的新基础和多模态模型

NVIDIA TAO 集成了开源、基础和专有模型，所有模型都在大量专有和商用数据集上进行训练，使其能够通用于物体检测、姿态检测、图像分类、分割等任务。TAO 可针对特定用例简化这些模型的微调，从而轻松进行自定义和商用部署。

大多数模型由 NVIDIA TensorRT 加速，并针对 NVIDIA 硬件的性能进行了优化，从而确保提供强大、高效的 AI 解决方案。

在 TAO 中交换模型主干非常简单，无需编码，只需简单的配置更改即可。这使您能够尝试不同的架构，例如 ResNet、Shifted Window (Swin) Transformer 和 Fully Attentional Network (FAN)，从而根据特定需求定制模型。

这种易于定制的方式支持多种应用，例如零售业中的产品识别、医疗保健行业中的医学影像分类、制造业中的机器人装配监控以及智能城市中的交通管理。

基础 DINO 模型

视频 1：使用 NVIDIA GroundingDINO 进行开放式词汇目标检测

传统的物体检测模型仅限于识别预定义类别（闭集检测）中的物体。这项约束使得这些模型无法用于需要根据人类输入识别任意物体的应用程序，例如特定类别名称或详细的引用表达式。

GroundingDINO 通过将文本编码器集成到 DINO 模型中，将其转换为开放式物体检测器，以解决这一限制。这使模型能够检测人类输入所描述的任何物体。

GroundingDINO 通过使用特征增强器、语言引导的查询选择和跨模式解码器，有效地融合了语言和视觉模式。这使模型能够泛化预定义类别以外的概念，实现卓越的性能。

将 DINO 与 TAO 相结合

此模型使用 Swin-Tiny 主干进行训练，在 NVIDIA 专有数据集上使用监督式方式进行商业用途。此外，BERT-Base 用作文本塔的起始权重。最后，GroundingDINO 使用从公开数据集收集的约 180 万张图像进行端到端训练。

在训练期间使用的所有原始图像均具有商业许可证，以确保安全的商业用途。

表 2： GroundingDINO KPIs on the COCO validation dataset

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 笔记，以提供简单的交互式环境，并尽可能减少编码需求。

表 3：DINO KPIs on the COCO 验证数据集

BEVFusion 模型

从自动驾驶到机器人开发和智慧城市，在许多领域中，系统都依赖于各种传感器来感知环境并与其交互。每种传感器类型（例如摄像头、激光雷达或雷达）都提供独特的信息，但也存在内在的局限性。例如，摄像头捕获丰富的视觉细节，但难以处理深度感知，而激光雷达在测量距离方面表现优异，但缺乏语义上下文。

挑战在于如何有效地将这些不同的数据源整合起来，以提高系统的性能和可靠性。

通过将来自多个传感器的数据集成到统一的鸟瞰图（BEV）表示中， BEVFusion 为应对这一挑战提供了解决方案。它通过将多模态特征统一为共享的鸟瞰图表示，从而彻底改变这一过程，同时保留（来自激光雷达的）几何信息和（来自摄像头的）语义信息。

与 TAO 的 BEVFusion

在 TAO 中，我们从 mmdet3d 的 BEVFusion 代码库开始训练模型。原始代码经过增强，可在 3D 空间中处理三个角度（roll、pitch、yaw），而最初仅支持 yaw。经过 TAO 训练的 BEVFusion 模型可在单个摄像头视图和激光雷达中检测人员，并在人员周围创建 3D 边界框。

内容	AP11	AP40
评估集	77.71%	79.35%

表 4： 单摄像头、单激光雷达人员检测 COCO 验证数据集上的 TAO BEVFusion KPIs

CLIP 模型

高效地学习有意义且上下文丰富的图像和文本表示是一项重大挑战，对于需要多模态理解的应用尤其如此。零样本学习、图像字幕、视觉搜索、内容调节和多模态交互等用例需要能够准确解释和集成视觉和文本数据的强大模型。

Contrastive Language-Image Pre-training (CLIP) 模型采用双编码器架构来同时处理图像和文本，从而解决了这一挑战。CLIP 利用对比学习最大限度地提高了正确的图像 – 文本对之间的相似性，同时最大限度地减少了错误对之间的相似性，使模型能够学习捕获各种概念的一般表示。

这种方法使 CLIP 能够在各种应用程序中表现出色，例如生成描述性图像描述、根据文本查询执行视觉搜索，以及实现零样本学习，模型可以利用其学习表示将图像分类为未见类别。

使用 TAO 的 NVCLIP

在 TAO 中，我们在 NVIDIA 数据集上提供预训练的 TensorRT 加速 CLIP 模型。我们使用了大约 700M 张图像来训练此模型。为进行评估，我们使用了 ImageNet 数据集的 50,000 张验证图像。

您可以将模型从 NVIDIA Jetson 部署到 NVIDIA Ampere 架构的 GPU 上。

模型	Top-1 精度
ViT-H -336	0.7786
升 336	0.7629

表 5： 使用 ViT-H-336 和 ViT-L-336 主干的 NVCLIP 零样本 ImageNet 验证准确率

SEGIC 模型

视频 3： 使用视觉提示进行上下文分割借助 NVIDIA TAO

SEGIC 是一款创新性的端到端上下文中分割框架，可以借助单视觉基础模型（VFM）彻底改变上下文中分割。

与传统方法不同，SEGIC 捕获目标图像与上下文示例之间的密集关系，提取几何、视觉和元指令以指导最终蒙版预测。这种方法显著降低了标记和训练成本，同时在一次性分割基准测试中实现了先进的性能。

SEGIC 的通用性扩展到各种任务，包括视频对象分割和开放词汇分割，使其成为图像分割领域的强大工具。

采用 TAO 的 SEGIC

对于在 TAO 中训练 SEGIC 模型，我们提供了一个预训练的 ONNX 模型和一个使用 NVIDIA Triton 推理服务器的部署方法。对于训练，我们使用了以下资源：

• 图像编码器： 使用 NVIDIA 专有图像训练的 DINOv2-L。

• 元描述编码器： CLIP-B，使用 NVIDIA 的专有图像进行训练。

模型	mIOU
SEGIC：DINOv2L，CLIP-B，mask_decoder	69.57

表 6：根据 COCO20i 数据集评估的 SEGIC 模型 KPIs

基础姿态模型

准确的 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 的框架，使用边界框生成高质量的实例分割蒙版。

这种组合可显著减少创建详细标签所需的工作量，从而更轻松、更快速地构建用于训练高级模型的可靠数据集。

视频 5： 使用文本提示通过 NVIDIA TAO 进行自动标注

使用 generate 命令生成自动标记：

tao dataset auto_label generate [-h] -e [results_dir= ] [num_gpus= ]

借助知识提炼实现高效 AI

知识提炼是机器学习中的一种技术，在该技术中，更小、更高效的模型（学生）学习模仿更大、更复杂的模型（教师）的行为。这一过程有助于减少训练和微调时间，因为学生模型使用教师已经捕获的知识（包括软标签等更细致入微的信息），而无需为任何新任务从头开始训练。

通过提炼这些知识，学生模型可以在显著减少计算资源的情况下实现类似的性能，使其成为在资源受限的环境中部署的理想选择，并加速训练过程。

TAO 支持知识提炼。指定不同层之间的绑定，以计算提炼损失。有关更多信息，请参阅 Optimizing the Training Pipeline 。

我们提供了 dino_distillation.ipynb 参考 notebook，以展示知识提炼，以及如何使用 L2 损失从 DINO 和 FAN 小至 DINO 和 ResNet-50 提取中间特征图。

distill.yaml: model: backbone: resnet_50 # Student model train_backbone: True distill: teacher: backbone: fan_small train_backbone: False num_feature_levels: 4 # Number of feature-maps to map from teacher pretrained_teacher_model_path: /workspace/tao-experiments/dino/pretrained_dino_coco_vdino_fan_small_trainable_v1.0/dino_fan_small_ep12.pth bindings: - teacher_module_name: 'model.backbone.0.body' student_module_name: 'model.backbone.0.body' criterion: L2

在 distill.yaml 文件中，您可以定义教师模型、学生模型主干和损失函数。训练模型：

!tao model dino distill \ -e $SPECS_DIR/distill.yaml \ results_dir=$RESULTS_DIR/

在训练学生模型后，您可以使用 TAO evaluate 工具评估提炼模型，使用 inference 工具可视化推理，使用 export 工具将训练好的模型导出到 ONNX，并最终生成 gen_trt_engine 文件以将 DINO 和 ResNet-50 模型导出到 TensorRT。

TAO 5.5 入门

全球各地的开发者正在使用 NVIDIA TAO 加速其视觉 AI 应用的 AI 训练。使用 TAO 5.5 的新功能来增强您的 AI 应用。