计算机视觉/视频分析

使用 NVIDIA TAO 6 和 NVIDIA DeepStream 8 构建实时视觉检测工作流

构建可靠的视觉检测流程以实现缺陷检测和质量控制并非易事。制造商和开发者常常面临诸多挑战,例如针对特定领域定制通用视觉 AI 模型、在计算资源受限的边缘设备上优化模型规模,以及高效部署以最大化推理吞吐量。

NVIDIA Metropolis 是一个面向视觉 AI 智能体和应用的开发平台,旨在帮助应对这些挑战。Metropolis 提供了构建跨多个阶段视觉检测工作流的模型和工具,包括:

  • 通过微调自定义视觉基础模型
  • 优化模型以实现实时推理
  • 将模型部署到生产流程中

NVIDIA Metropolis 提供了一个统一的框架,包含用于训练和优化视觉 AI 基础模型NVIDIA TAO 6,以及端到端流分析工具包 NVIDIA DeepStream 8NVIDIA TAO 6NVIDIA DeepStream 8 现已开放下载。您可通过 NVIDIA TAO 文档NVIDIA DeepStream 文档 了解最新功能更新。

本文将为您介绍如何利用 NVIDIA TAO 和 NVIDIA DeepStream 构建端到端的实时视觉检测工作流,具体步骤包括:

  • 使用 TAO 执行自监督式微调,以利用特定领域的未标记数据。
  • 使用 TAO 知识蒸馏优化基础模型,以提高吞吐量和效率。
  • 使用 DeepStream Inference Builder 进行部署,DeepStream Inference Builder 是一款低代码工具,可将模型创意转化为生产就绪型独立应用或可部署的微服务。

如何使用 NVIDIA TAO 通过视觉基础模型扩展自定义模型开发

NVIDIA TAO 支持针对特定领域用例训练、调整和优化大型视觉基础模型的端到端工作流程。该框架专为自定义视觉基础模型而设计,通过微调微服务实现高准确率和卓越性能。

Flow diagram showing an overview of the end-to-end scope of NVIDIA TAO.
图 1。使用 NVIDIA TAO 创建高度准确、自定义的企业就绪型 AI 模型,为您的视觉 AI 应用提供支持

视觉基础模型 (VFM) 是在大量多样化数据集上训练的大规模神经网络,能够捕捉通用且强大的视觉特征表示。得益于这种泛化能力,VFM 成为适用于多种 AI 感知任务(例如图像分类、物体检测和语义分割)的灵活模型骨干。

TAO 提供了一系列强大的基础主干和任务头,可针对工业视觉检测等关键工作负载对模型进行微调。在 TAO 6 中,两个关键的基础主干是 C-RADIOv2(开箱即用精度领先)和 NV-DINOv2。TAO 还支持第三方模型,前提是其视觉主干和任务头架构与 TAO 兼容。

The diagram shows the TAO fine-tuning workflow. It starts with a foundation backbone that learns image features from your dataset, followed by task head layers (classification, detection, segmentation) that use these feature maps to generate final predictions.
图 2。使用 NVIDIA TAO 微调框架、基础模型主干和任务头扩展自定义视觉模型开发

为提高模型准确性,TAO 支持多种模型自定义技术,例如监督式微调 (SFT) 和自监督式学习 (SSL)。其中,SFT 需要针对特定计算机视觉下游任务收集经过精心策划的带标注数据集。然而,获取高质量的标注数据是一个复杂且耗时的过程,通常需要大量人工投入,成本较高。

其次,借助 NVIDIA TAO 6,您可以利用自监督学习来挖掘未标记图像的巨大潜力,从而在标记数据稀缺或成本高昂的情况下加速模型定制流程。

这种方法也称为域自适应使您能够利用无标记数据构建强大的基础模型主干,例如 NV-DINOv2。随后,可将该主干与任务头结合,并使用较小的标注数据集对多种下游检测任务进行微调。

在实际场景中,此工作流意味着模型能够从大量未标记图像中学习缺陷的细微特征,再通过有针对性的监督式微调来优化决策能力,即使在自定义的真实数据集上也能实现卓越的性能。

A diagram showing the two stages to effectively adapt and finetune a large scale trained foundation model to a specific downstream task.
图 3。用于调整基础模型以适应特定下游用例的端到端工作流程

通过基础模型微调提高 PCB 缺陷检测准确性

例如,我们采用大规模无标注印刷电路板(PCB)图像,应用TAO基础模型自适应工作流,对视觉基础模型进行微调,以实现缺陷检测。Starting with NV-DINOv2,这是一个基于7亿张通用图像训练而成的通用模型,我们利用包含约70万张无标注PCB图像的数据集,通过自监督学习(SSL)对模型进行针对PCB应用的定制化优化。这一过程有效推动模型从广泛泛化能力向特定领域专业性能的转化。

完成领域适应后,我们利用带标注的 PCB 数据集,通过线性探测优化特定任务的分类头部以提升准确性,并进一步进行全模型微调,以同时调整主干网络和分类头部。第一个数据集包含约 600 个训练样本和 400 个测试样本,用于将图像分类为“OK”(正常)或“Defect”(缺陷),其中缺陷类别包括缺失、偏移、倒置、焊接不良以及异物等模式。

特征图显示,经过调整的 NV-DINOv2 甚至在进行下游微调之前,便能清晰地区分组件以及前景与背景(图 4 和图 5)。该模型在将集成电路(IC)引脚等复杂结构从背景中分离方面表现出色,这是通用模型难以实现的任务。

Two side-by-side images comparing the features from a generic NV-DINOv2 model (versus a domain adapted NV-DINOv2 model when computed for an PCB image for the OK class.
图 4。使用领域适应性 NV-DINOv2 (左) 和通用 NV – DINOv2 (右) 比较 OK 类的特征图
Two side-by-side images comparing the features from a generic NV-DINOv2 model (versus a domain adapted NV-DINOv2 model when computed for an PCB image for the “Defect” class.
图 5。使用领域适应性 NV-DINOv2 (左) 和通用 NV – DINOv2 (右) 比较缺陷类的特征图

分类准确率提升了4.7个百分点,从93.8%提高至98.5%。

Plot showing evolution of accuracy over the number of epochs during training when starting from a generic NV-DINOv2 vs an NV-DINOv2 checkpoint that’s domain adapted on unlabeled images.
图 6。适用于领域的 NV-DINOv2 与通用 NV – DINOv2 的准确度比较

适用于不同领域的 NV-DINOv2 还展现出强大的视觉理解能力,并能在同一领域中提取相关图像特征。这表明,通过下游监督式微调,即使使用较少的标注数据,也能实现相近甚至更高的准确性。

在某些情况下,收集多达 70 万张未标记图像的数据可能仍然具有挑战性。然而,即使使用较小的数据集,您依然能够从 NV-DINOv2 域自适应中受益。

图 7 显示了实验的运行结果,该实验仅使用 10 万张图像对 NV-DINOv2 进行调整,其性能优于标准的 NV-DINOv2 模型。

Plot comparing accuracy convergence over the duration of the training (in epochs) for a when starting from a generic NV-DINOv2 (in green), domain adapted NV-DINOv2 with 100k images (in blue)
and a domain adapted NV-DINOv2 with 700k images (in orange).
图 7。不同 NV-DINOv2 模型分类准确性比较 for 分类

此示例说明了如何使用 NVIDIA TAO 和 NV-DINOv2 对无标签域数据进行自监督学习,以实现可靠且准确的 PCB 缺陷检测,同时降低对大量标注样本的依赖。

如何优化视觉基础模型以提高吞吐量

优化是部署深度学习模型的重要一步。许多生成式 AI 和视觉基础模型可能包含上亿个参数,导致计算量极为庞大,对于用于工业视觉检测或实时交通监控系统等实时应用的大多数边缘设备而言,往往难以承受。

NVIDIA TAO 利用从大型基础模型中获得的知识,采用名为“知识蒸馏”(Knowledge Distillation)的技术,将其优化为更小的模型规模。知识蒸馏通常能够在不牺牲准确性的前提下,将大型、高精度的教师模型压缩为更小、更高效的學生模型。该过程通过让学生模型不仅模仿教师模型的最终预测,还模仿其内部特征表示和决策边界,从而在资源受限的硬件上实现实际部署,并支持可扩展的模型优化。

NVIDIA TAO 通过强大支持多种蒸馏形式(包括主干、logit 以及空间/特征蒸馏),进一步拓展了知识蒸馏的能力。TAO 的一项显著特性是其专为物体检测设计的单级蒸馏方法。借助这一简化流程,学生模型(通常更小且更高效)能够在单一统一的训练阶段,直接从教师模型学习主干特征表示以及任务特定的预测结果。这种方法在不牺牲精度的前提下,显著降低了推理延迟和模型规模。

将单阶段蒸馏应用于实时 PCB 缺陷检测模型

蒸馏使用 TAO 的有效性在 PCB 缺陷检测数据集 上进行了评估,该数据集包含 9,602 张训练图像和 1,066 张测试图像,涵盖六个具有挑战性的缺陷类别:缺失孔、鼠咬、开路、短路、短刺和假铜。他们采用两个不同的教师模型进行蒸馏评估。实验基于从 ImageNet-1K 预训练权重初始化的主干网络,并以目标检测中常用的 COCO 平均精度(mAP)作为评估指标。

Flow diagram with icons labeled (clockwise from bottom center) Data, Teacher Model, Knowledge, and Student Model.
图 8。使用 NVIDIA TAO 将蒸馏知识从更大的教师模型转换为更小的学生模型

在第一组实验中,我们采用 ResNet 系列骨干网络,在师生架构下进行了相同的蒸馏实验,结果学生模型的准确率不仅达到了教师模型的水平,甚至有所超越。

基准实验作为与 TAO 中的 RT-DETR 模型相关的 train 动作运行。以下代码段展示了可用于执行训练实验的可行实验规格文件。

model:
  backbone: resnet_50
  train_backbone: true
  num_queries: 300
  num_classes: 7

train:
  num_gpus: 1
  epochs: 72
  batch_size: 4
  optim:
    lr: 1e-4
    lr_backbone: 1.0e-05

dataset:
  train_data_sources:
    - image_dir: /path/to/dataset/images/train
      json_file: /path/to/dataset/annotations/train.json
  val_data_sources:
    image_dir: /path/to/dataset/images/val
    json_file: /path/to/dataset/annotations/val.json
  test_data_sources:
    image_dir: /path/to/dataset/images/test
    json_file: /path/to/dataset/annotations/test.json
  batch_size: 4
  remap_coco_categories: false
  augmentation:
    multiscales: [640]
    train_spatial_size: [640, 640]
    eval_spatial_size: [640, 640]

To run train , use the following command:

To run train, use the following command:

tao model rtdetr train -e /path/to/experiment/spec.yaml results_dir=/path/to/results/dir model.backbone=backbone_name model.pretrained_backbone_path=/path/to/the/pretrained/model.pth

您可以更改主干,方法是将 model.backbone 参数重写为主干名称,并将 model.pretrained_backbone_path 重写为主干的预训练检查点文件的路径。

蒸馏实验在 TAO 中作为与 RT-DETR 模型相关的 IG14Q 动作运行。要配置蒸馏实验,可将以下 IG15Q 元素添加到原始训练实验 IG16Q 文件中。

distill:
  teacher:
    backbone: resnet_50
  pretrained_teacher_model_path: /path/to/the/teacher/checkpoint.pth

使用以下示例命令运行蒸馏:

tao model rtdetr distill -e /path/to/experiment/spec/yaml results_dir=/path/to/results/dir model.backbone=backbone_namemodel.pretrained_backbone_path=/path/to/pretrained/backbone/checkpoint.pth distill.teacher.backbone=teacher_backbone_name distill.pretrained_teacher_model_path=/path/to/the/teacher/model.pth
Graph showing a ResNet50 teacher model distilled into a lighter ResNet18 student model, achieving a 5% accuracy gain.
图 9。将 ResNet50 模型蒸为更轻的 ResNet18 模型,可将准确率提高 5%

在边缘部署模型时,推理加速和内存限制可能是关键考虑因素。TAO 不仅能够在同一主干系列中提取检测功能,还能跨不同主干系列提取检测功能。

Graph showing a ConvNeXt teacher model distilled into a lighter ResNet34-based student model, achieving a 3% accuracy gain.
图 10。将 ConvNeXt 模型蒸为基于 ResNet34 的较轻模型,可将准确率提高 3%

在本示例中,我们采用基于 ConvNeXt 的 RT-DETR 模型作为教师模型,并将其知识蒸馏至基于 ResNet34 的轻量级模型。通过单阶段蒸馏,TAO 将准确率提升了 3%,模型大小缩减了 81%,从而实现了更高的吞吐量和低延迟推理。

如何使用 DeepStream 8 推理构建器打包和部署模型

现在,借助 TAO 中经过训练和蒸馏的 RT-DETR 模型,下一步是将其部署为推理微服务。全新的 NVIDIA DeepStream 8 推理构建器 是一款低代码工具,可将模型构想快速转化为独立应用或可部署的微服务。

要使用推理构建器,请提供 YAML 配置、Dockerfile 以及可选的 OpenAPI 定义。推理构建器将基于这些输入生成 Python 代码,自动串联数据加载、GPU+ 加速预处理、推理及后处理等阶段,并可生成用于微服务部署的 REST 端点。

其目标是根据用户提供的模型和配置文件,自动生成推理服务代码、API 层以及部署所需的构件。通过简单的配置,推理构建器即可处理与服务器、请求处理和数据流相关的复杂任务,从而避免了手动编写大量重复性代码的需求。

视频 1。了解如何使用 NVIDIA DeepStream 推理构建器部署 AI 模型

第 1 步:定义配置

  • 创建 config.yaml 文件以绘制模型和推理工作流
  • (可选) 如果需要明确的 API 模式定义,则加入 openapi.yaml 文件

第 2 步:执行 DeepStream 推理构建器

  • 将配置提交至推理构建器
  • 此实用程序利用推理模板、服务器模板和实用程序 (例如编解码器) 自主生成项目代码
  • 输出构成了一个全面的软件包,包括推理逻辑、服务器代码和辅助实用程序
  • 输出 infer.tgz,一种打包的推理服务

第 3 步:检查生成的代码

该软件包扩展为一个精心策划的项目,包含:

  • 配置:config/
  • 服务器逻辑:server/
  • 推理库:lib/
  • 实用程序:资产管理器、编解码器、响应人员等

第 4 步:构建 Docker 镜像

  • 使用参考 Dockerfile 容器化服务
  • 执行 docker build -t my-infer-service

第 5 步:使用 Docker Compose 进行部署

  • 使用 Docker Compose:docker-compose up 启动服务
  • 随后,该服务将在容器内加载您的模型

第六步:服务用户

  • 推理微服务现已投入运行
  • 最终用户或应用可以向公开的 API 端点发送请求,并直接从您的模型接收预测

如需详细了解 NVIDIA DeepStream 推理构建器,请访问 NVIDIA-AI-IOT/deepstream_tools on GitHub。

用于实时视觉检查的其他应用程序

除了识别 PCB 缺陷之外,您还可以将 TAO 和 DeepStream 应用于汽车和物流等行业的异常检测。有关具体用例,请参阅使用合成数据大幅缩短制造 AI 部署时间以及 NVIDIA TAO

开始构建实时视觉检测工作流

借助 NVIDIA DeepStream 和 NVIDIA TAO,开发者正在推动视觉 AI 从快速原型设计到大规模部署的演进。

DeepStream 8.0 为开发者提供推理构建器等强大工具,简化工作流的创建,并提升复杂环境中的跟踪准确性。TAO 6 通过领域适应性、自监督微调和知识蒸馏,充分释放基础模型的潜力。

这转化为更快的迭代周期、更好地利用无标记数据,以及生产就绪的推理服务。

准备好了吗?

下载 NVIDIA TAO 6,体验全新功能。欢迎前往 NVIDIA TAO 开发者论坛提问并参与讨论。

下载 NVIDIA DeepStream 8探索新功能。在 NVIDIA DeepStream 开发者论坛 中提问并参与讨论。

 

标签