利用 NVIDIA TAO 和视觉 AI 模型实现工业缺陷检测的变革

效率对于工业制造来说至关重要，即便是微小的提升也会带来显著的经济效益。美国质量协会的数据显示，“许多企业机构实际的质量相关成本高达销售收入的 15 - 20%，有些甚至高达总运营成本的 40%。”这些惊人的统计数据揭示了一个严峻的现实：工业应用中的缺陷不仅会危及产品质量，还会消耗很大一部分的企业收入。

但如果企业能够收回这些损失的利润，并将其重新投入到创新和扩张中，那又会怎样呢？这正是 AI 的“用武之地”。

本文将探讨如何使用 NVIDIA TAO 设计出能够精准定位工业应用缺陷，进而提高整体质量的自定义 AI 模型。

NVIDIA TAO 套件是基于 TensorFlow 和 PyTorch 构建的低代码 AI 工具套件。它通过抽象化 AI 模型和深度学习框架的复杂性，来简化并加速模型训练过程。借助 TAO 套件，开发者不仅可以使用预训练模型，还可以针对特定用例对其进行微调。

在本文中，我们使用一个名为 VisualChangeNet 的高级预训练模型进行变化检测，并使用 TAO 工具套件对其进行微调，使其能够检测 MVTec Anomaly 异常检测数据集中的缺陷。这个综合基准数据集由各种工业产品的正常和缺陷样本组成，专门用于机器视觉中的异常检测。

借助 TAO 套件，我们使用迁移学习训练了一个模型，该模型在 MVTec Anomaly 异常数据集的“瓶子”类别上达到了 99.67% 的总体准确率、92.3% 的 mIoU、95.8% 的 mF1、97.5 的 mPrecision 和 94.3% 的 mRecall。图 1 显示了使用这个训练有素的模型进行的缺陷掩码预测。

图 1. 通过将缺陷图像与黄金图像进行比较 来分段预测缺陷物体的缺陷掩码

第 1 步：设置先决条件

为了按照本文重新创建这些步骤，请执行以下操作。

1. 按照《NGC 用户指南》中提供的步骤在 NGC 目录上注册账户并生成 API 密钥：https://docs.nvidia.com/ngc/gpu-cloud/ngc-user-guide/index.html#generating-api-key

1. 按照《TAO 快速入门指南》设置 TAO 启动器。为 MVTec 数据集下载 VisualChangeNet Segmentation Jupyter Notebook。启动 Jupyter Notebook 并运行单元格，以便跟着本文进行操作。

*请注意：VisualChangeNet 模型只能在 5.1 及以上的版本中运行。

3.按照提示进入下载页面，下载并准备 MVTec 异常检测数据集，然后复制 15 个对象类中任何一个的下载链接：https://www.mvtec.com/company/research/datasets/mvtec-ad

1. 将下载链接粘贴到 Jupyter Notebook 第 2.1 部分中的“FIXME”位置，然后运行该笔记本单元。本文主要介绍瓶子对象，但所有 15 个对象都能在该笔记本中使用。图 2 显示了数据集中的缺陷图像样本。

图 2. MVTec 数据集中电缆、瓶子 和晶体管的缺陷图像样本（从左到右）

我们用 MVTec-AD 中的“瓶子”类别展示了如何在工业检测用例中使用 TAO 套件与 VisualChangeNet 来实现自动光学检测。

在 Jupyter Notebook 下载数据集后，运行该笔记的第 2.3 节，将数据集处理成 VisualChangeNet 分割所需的正确格式。

import random import shutil from PIL import Imageos.environ["HOST_DATA_DIR"] = os.path.join(os.environ["LOCAL_PROJECT_DIR"], "data", "changenet")formatted_dir = f"formatted_{mvtec_object}_dataset" DATA_DIR = os.environ["HOST_DATA_DIR"]os.environ["FORMATTED_DATA_DIR"] = formatted_dir #setup dataset folders in expected format formatted_path = os.path.join(DATA_DIR, formatted_dir)a_dir = os.path.join(formatted_path, "A")b_dir = os.path.join(formatted_path, "B")label_dir = os.path.join(formatted_path, "label")list_dir = os.path.join(formatted_path, "list") #Create the expected foldersos.makedirs(formatted_path, exist_ok=True)os.makedirs(a_dir, exist_ok=True)os.makedirs(b_dir, exist_ok=True)os.makedirs(label_dir, exist_ok=True)os.makedirs(list_dir, exist_ok=True)

该原始数据集专门用于异常检测。我们将两者合并成一个包含 283 幅图像的综合数据集，然后将其分成 253 幅训练集图像和 30 幅测试集图像。两个数据集都包含有缺陷的样本。

我们确保该测试集包含每个缺陷类中 30% 的缺陷样本，“瓶子”类别主要包含“无缺陷”图像，三个缺陷类别中，每个类别各有约 20 幅图像。

图 3. 该数据集中的输入样本，包含测试图像、黄金图像和显示缺陷的分割掩码。图中采用从瓶子顶部俯视的视角，摄像头安装在瓶口的正上方。

第 2 步：下载 VisualChangeNet 模型

VisualChangeNet 模型是最先进的基于转换器的变化检测模型。其设计核心是孪生网络（Siamese Network）。孪生网络是一种独特的神经网络架构，由两个或多个相同的子网络组成。这些“孪生”子网络接受不同的输入，但共享相同的参数和权重。这种架构使 VisualChangeNet 模型能够比较当前图像和作为参考的“黄金”图像之间的特征，从而精准定位变化和变更。该功能使孪生网络尤其擅长图像对比和异常检测等任务。

模型文档提供更多细节，如架构和训练数据等。我们没有从头开始训练模型，而是从在 NV-ImageNet 数据集上训练而成的预训练 FAN 骨干开始。我们使用 TAO 套件在 MVTec-AD 数据集上针对“瓶子”类别对其进行微调。

运行笔记第 3 部分以安装 NGC 命令行工具，并从 NGC 下载该预训练骨干。

# Installing NGC CLI on the local machine.## Download and installimport os%env CLI=ngccli_cat_linux.zip!mkdir -p $HOST_RESULTS_DIR/ngccli # # Remove any previously existing CLI installations!rm -rf $HOST_RESULTS_DIR/ngccli/*!wget "https://ngc.nvidia.com/downloads/$CLI" -P $HOST_RESULTS_DIR/ngccli!unzip -u "$HOST_RESULTS_DIR/ngccli/$CLI" -d $HOST_RESULTS_DIR/ngccli/!rm $HOST_RESULTS_DIR/ngccli/*.zipos.environ["PATH"]="{}/ngccli/ngc-cli:{}".format(os.getenv("HOST_RESULTS_DIR", ""), os.getenv("PATH", ""))!mkdir -p $HOST_RESULTS_DIR/pretrained!ngc registry model list nvidia/tao/pretrained_fan_classification_nvimagenet*!ngc registry model download-version "nvidia/tao/pretrained_fan_classification_nvimagenet:fan_base_hybrid_nvimagenet" --dest $HOST_RESULTS_DIR/pretrained

第 3 步：使用 TAO 套件训练模型

在本节中，我们将详细介绍如何使用 TAO 套件训练 VisualChangeNet 模型。您可在模型卡中找到 VisualChangeNet 模型的详细信息以及支持的预训练权重，也可以使用该预训练 FAN 主干权重作为微调 VisualChangeNet 的起点，这也是我们在 MVTec-AD 数据集上进行微调时使用的权重。

如图 4 所示，该训练算法会同步更新所有子网络的参数。在 TAO 中，VisualChangeNet 支持将两幅图像作为黄金样本和测试样本输入，其目标是检测“黄金或参考”图像与“测试”图像之间的变化。TAO 支持 VisualChangeNet 架构的 FAN 骨干网络。

TAO 支持两类变化检测网络，分别是 VisualChangeNet-Segmentation 和 VisualChangeNet-Classification。在本文中，我们将使用 VisualChangeNet-Segmentation 模型，通过分割 MVTec-AD 数据集中两张输入图像之间的变化像素，来演示变化检测。

图 4. 检测瓶子类黄金图像和测试图像之间变化的 VisualChangeNet-Segmentation 算法架构图

使用工具套件微调 VisualChangeNet 模型非常简单，无需编码经验。只需在 TAO 套件中加载数据、设置实验配置并运行训练命令即可。

该实验配置文件定义了 VisualChangeNet 模型架构、训练和评估的超参数。可在训练模型前在 Jupyter Notebook 中查看并编辑该配置文件。

我们使用该配置对 VisualChangeNet 模型进行微调，在该配置中定义了一个带有预训练 FAN-Hybrid-Base 骨干（即基准模型）的 VisualChangeNet 模型，然后对该模型进行 30 个历元训练，批量大小为 8。下面展示了部分实验配置以及一些关键参数。完整的实验配置可在 Jupyter Notebook 中进行查看。

encryption_key: tlt_encodetask: segmenttrain: resume_training_checkpoint_path: null pretrained_model_path: null segment: loss: "ce" weights: [0.5, 0.5, 0.5, 0.8, 1.0] num_epochs: 30 num_nodes: 1 val_interval: 1 checkpoint_interval: 1 optim: lr: 0.0002 optim: "adamw" policy: "linear" momentum: 0.9 weight_decay: 0.01results_dir: "/results"model: backbone: type: "fan_base_16_p4_hybrid" pretrained_backbone_path: /results/pretrained/pretrained_fan_classification_nvimagenet_vfan_base_hybrid_nvimagenet/fan_base_hybrid_nvimagenet.pth

为了调整模型的性能，可以修改的一些常用值包括训练历元数、学习率（lr）、优化器和预训练骨干。如要从头开始训练，可以将 pretrained_backbone_path 设置为空，但这可能会增加为达到高准确率所需的历元数和数据量。如需进一步了解该实验配置文件中的参数，请参阅《VisualChangeNet用户指南》：https://docs.nvidia.com/tao/tao-toolkit/text/visual_changenet/index.html

现在数据集和实验配置已经准备就绪，我们要开始在 TAO 套件中进行训练了。运行第 5.1 节中的代码块，使用单颗 GPU 启动 VisualChangeNet 训练。

print("Train model")!tao model visual_changenet train \\ -e $SPECS_DIR/experiment.yaml \\ train.num_epochs=$NUM_EPOCHS \\ dataset.segment.root_dir=$DATA_DIR \\ model.backbone.pretrained_backbone_path=$BACKBONE_PATH

该单元将开始在 MVTec 数据集上训练 VisualChangeNet-Segmentation 模型。在训练过程中，模型将学习如何识别缺陷对象并输出显示缺陷区域的分割掩码。训练日志包含验证数据集上的准确率、训练损失、学习率和训练后的模型等，都将保存在实验配置中设置的结果目录中。

第 4 步：评估模型

训练完成后，我们可以使用 TAO 在验证数据集上对模型进行评估。VisualChangeNet Segmentation 的输出是 2 张给定输入图像的分割变化图，表示的是像素级缺陷。笔记的第 6 节将运行评估模型性能的命令。

!tao model visual_changenet evaluate \\ -e $SPECS_DIR/experiment.yaml \\ evaluate.checkpoint=$RESULTS_DIR/train/changenet.pth \\ dataset.segment.root_dir=$DATA_DIR

TAO 中的评估命令将返回验证集的多个 KPI，比如在缺陷类（缺陷像素）上的准确度、精确度、召回率、F1 分数和 IoU 等。

OA = 有变化/无变化像素的总体准确率（输入维度 - 256×256）

表 1. MVTec-AD 二进制变化检测（瓶子类） 中的 VisualChangeNet 模型评估指标

第 5 步：部署模型

可以使用 NVIDIA DeepStream 或 NVIDIA Triton 部署这一经过微调的模型。我们将其导出为 .onnx 格式。笔记第 8 节将运行 TAO 导出命令。

!tao model visual_changenet export \\ -e $SPECS_DIR/experiment.yaml \\ export.checkpoint=$RESULTS_DIR/train/changenet.pth \\ export.onnx_file=$RESULTS_DIR/export/changenet.onnx

输出的 .onnx 模型与训练好的 .pth 模型保存在同一目录下。如要部署到 Triton，请查看 GitHub 上的 tao-toolkit-triton 资源库。该项目提供了将许多 TAO 模型（包括 VisualChangeNet Segmentation）部署到 Triton 推理服务器的参考实现。

实时推理性能

推理以 FP16 精度在所提供的未修剪模型上运行。该推理性能在嵌入式 Jetson Orin GPU 和数据中心 GPU 上使用 trtexec 运行。Jetson 设备以 GPU 频率最高的 Max-N 配置运行。

请执行以下命令来运行 trtexec：

/usr/src/tensorrt/bin/trtexec --onnx= --minShapes=input0:1x3x512x512,input1:1x3x512x512 --maxShapes=input0:8x3x512x512,input1:8x3x512x512 --optShapes=input0:4x3x512x512,input1:4x3x512x512 --saveEngine=

此处显示的性能仅为推理性能。流式传输视频数据的端到端性能可能会因硬件和软件的其他瓶颈而有所变化。

总结

在本文中，我们了解了如何使用 TAO 套件微调 VisualChangeNet 模型，并将其用于分割 MVTec 数据集中的缺陷，从而实现 99.67% 的总体准确率。

现在，您还可以使用 NVIDIA TAO 检测制造工作流程中的缺陷。

如要开始使用，请先：

• 从 NVIDIA NGC 目录下载 VisualChangeNet 模型。
• 
  
• 按照《TAO 快速入门指南》设置 TAO 启动器：https://docs.nvidia.com/tao/tao-toolkit/text/tao_toolkit_quick_start_guide.html
• 
  
• 从 GitHub 下载 VisualChangeNet Segmentation Notebook：https://github.com/NVIDIA/tao_tutorials/tree/main/notebooks/tao_launcher_starter_kit/visual_changenet
• 
  
• 请通过 NVIDIA 文档进一步了解 NVIDIA TAO 工具套件：https://docs.nvidia.com/tao/tao-toolkit/text/overview.html#:~:text=NVIDIA%20TAO%20Toolkit%20is%20a,and%20the%20deep%20learning%20framework