效率對于工業制造來說至關重要，即便是微小的提升也會帶來顯著的經濟效益。美國質量協會的數據顯示，“許多企業機構實際的質量相關成本高達銷售收入的 15 - 20%，有些甚至高達總運營成本的 40%。”這些驚人的統計數據揭示了一個嚴峻的現實：工業應用中的缺陷不僅會危及產品質量，還會消耗很大一部分的企業收入。

但如果企業能夠收回這些損失的利潤，并將其重新投入到創新和擴張中，那又會怎樣呢？這正是 AI 的“用武之地”。

本文將探討如何使用NVIDIA TAO設計出能夠精準定位工業應用缺陷，進而提高整體質量的自定義 AI 模型。

NVIDIA TAO 套件是基于TensorFlow和 PyTorch構建的低代碼 AI 工具套件。它通過抽象化 AI 模型和深度學習框架的復雜性，來簡化并加速模型訓練過程。借助 TAO 套件，開發者不僅可以使用預訓練模型，還可以針對特定用例對其進行微調。

在本文中，我們使用一個名為 VisualChangeNet 的高級預訓練模型進行變化檢測，并使用 TAO 工具套件對其進行微調，使其能夠檢測MVTecAnomaly異常檢測數據集中的缺陷。這個綜合基準數據集由各種工業產品的正常和缺陷樣本組成，專門用于機器視覺中的異常檢測。

借助 TAO 套件，我們使用遷移學習訓練了一個模型，該模型在MVTecAnomaly異常數據集的“瓶子”類別上達到了 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

2. 按照《TAO 快速入門指南》設置 TAO 啟動器。為 MVTec 數據集下載 VisualChangeNet Segmentation Jupyter Notebook。啟動 Jupyter Notebook 并運行單元格，以便跟著本文進行操作。

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

3. 按照提示進入下載頁面，下載并準備MVTec 異常檢測數據集，然后復制 15 個對象類中任何一個的下載鏈接：https://www.mvtec.com/company/research/datasets/mvtec-ad

4. 將下載鏈接粘貼到 Jupyter Notebook 第 2.1 部分中的“FIXME”位置，然后運行該筆記本單元。本文主要介紹瓶子對象，但所有 15 個對象都能在該筆記本中使用。圖 2 顯示了數據集中的缺陷圖像樣本。

#Download the data
import os
MVTEC_AD_OBJECT_DOWNLOAD_URL = "FIXME"
mvtec_object = MVTEC_AD_OBJECT_DOWNLOAD_URL.split("/")[-1].split(".")[0]
os.environ["URL_DATASET"]=MVTEC_AD_OBJECT_DOWNLOAD_URL
os.environ["MVTEC_OBJECT"]=mvtec_object
!if[!-f$HOST_DATA_DIR/$MVTEC_OBJECT.tar.xz];thenwget$URL_DATASET-O$HOST_DATA_DIR/$MVTEC_OBJECT.tar.xz;elseecho"imagearchivealreadydownloaded";fi

圖 2. MVTec數據集中電纜、瓶子

和晶體管的缺陷圖像樣本（從左到右）

我們用 MVTec-AD 中的“瓶子”類別展示了如何在工業檢測用例中使用 TAO 套件與 VisualChangeNet 來實現自動光學檢測。

在 Jupyter Notebook 下載數據集后，運行該筆記的第 2.3 節，將數據集處理成 VisualChangeNet 分割所需的正確格式。

import random 
import shutil 
from PIL import Image
os.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 folders
os.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 install
import 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/*.zip
os.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*
!ngcregistrymodeldownload-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_encode
task: segment
train:
  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.01
results_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 OrinGPU 和數據中心 GPU 上使用 trtexec 運行。Jetson 設備以 GPU 頻率最高的 Max-N 配置運行。

請執行以下命令來運行 trtexec：

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

此處顯示的性能僅為推理性能。流式傳輸視頻數據的端到端性能可能會因硬件和軟件的其他瓶頸而有所變化。

總結

在本文中，我們了解了如何使用 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

GTC 2024 將于 2024 年 3 月 18 至 21 日在美國加州圣何塞會議中心舉行，線上大會也將同期開放。點擊“閱讀原文”或掃描下方海報二維碼，立即注冊 GTC 大會。