利用 NVIDIA TAO 和視覺 AI 模型實現(xiàn)工業(yè)缺陷檢測的變革

效率對于工業(yè)制造來說至關(guān)重要，即便是微小的提升也會帶來顯著的經(jīng)濟效益。美國質(zhì)量協(xié)會的數(shù)據(jù)顯示，“許多企業(yè)機構(gòu)實際的質(zhì)量相關(guān)成本高達銷售收入的 15 - 20%，有些甚至高達總運營成本的 40%?！边@些驚人的統(tǒng)計數(shù)據(jù)揭示了一個嚴峻的現(xiàn)實：工業(yè)應用中的缺陷不僅會危及產(chǎn)品質(zhì)量，還會消耗很大一部分的企業(yè)收入。

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

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

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

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

借助 TAO 套件，我們使用遷移學習訓練了一個模型，該模型在MVTecAnomaly異常數(shù)據(jù)集的“瓶子”類別上達到了 99.67% 的總體準確率、92.3% 的 mIoU、95.8% 的 mF1、97.5 的 mPrecision 和 94.3% 的 mRecall。圖 1 顯示了使用這個訓練有素的模型進行的缺陷掩碼預測。

圖 1. 通過將缺陷圖像與黃金圖像進行比較

來分段預測缺陷物體的缺陷掩碼

第 1 步：設置先決條件

為了按照本文重新創(chuàng)建這些步驟，請執(zhí)行以下操作。

1. 按照《NGC 用戶指南》中提供的步驟在 NGC 目錄上注冊賬戶并生成 API 密鑰：https://docs.nvidia.com/ngc/gpu-cloud/ngc-user-guide/index.html#generating-api-key

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

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

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

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

#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數(shù)據(jù)集中電纜、瓶子

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

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

在 Jupyter Notebook 下載數(shù)據(jù)集后，運行該筆記的第 2.3 節(jié)，將數(shù)據(jù)集處理成 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)

該原始數(shù)據(jù)集專門用于異常檢測。我們將兩者合并成一個包含 283 幅圖像的綜合數(shù)據(jù)集，然后將其分成 253 幅訓練集圖像和 30 幅測試集圖像。兩個數(shù)據(jù)集都包含有缺陷的樣本。

我們確保該測試集包含每個缺陷類中 30% 的缺陷樣本，“瓶子”類別主要包含“無缺陷”圖像，三個缺陷類別中，每個類別各有約 20 幅圖像。

圖 3. 該數(shù)據(jù)集中的輸入樣本，包含測試圖像、黃金圖像和顯示缺陷的分割掩碼。圖中采用從瓶子頂部俯視的視角，攝像頭安裝在瓶口的正上方。

第 2 步：下載 VisualChangeNet 模型

VisualChangeNet 模型是最先進的基于轉(zhuǎn)換器的變化檢測模型。其設計核心是孿生網(wǎng)絡（Siamese Network）。孿生網(wǎng)絡是一種獨特的神經(jīng)網(wǎng)絡架構(gòu)，由兩個或多個相同的子網(wǎng)絡組成。這些“孿生”子網(wǎng)絡接受不同的輸入，但共享相同的參數(shù)和權(quán)重。這種架構(gòu)使 VisualChangeNet 模型能夠比較當前圖像和作為參考的“黃金”圖像之間的特征，從而精準定位變化和變更。該功能使孿生網(wǎng)絡尤其擅長圖像對比和異常檢測等任務。

模型文檔提供更多細節(jié)，如架構(gòu)和訓練數(shù)據(jù)等。我們沒有從頭開始訓練模型，而是從在 NV-ImageNet 數(shù)據(jù)集上訓練而成的預訓練 FAN 骨干開始。我們使用 TAO 套件在 MVTec-AD 數(shù)據(jù)集上針對“瓶子”類別對其進行微調(diào)。

運行筆記第 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 套件訓練模型

在本節(jié)中，我們將詳細介紹如何使用 TAO 套件訓練 VisualChangeNet 模型。您可在模型卡中找到 VisualChangeNet 模型的詳細信息以及支持的預訓練權(quán)重，也可以使用該預訓練 FAN 主干權(quán)重作為微調(diào) VisualChangeNet 的起點，這也是我們在 MVTec-AD 數(shù)據(jù)集上進行微調(diào)時使用的權(quán)重。

如圖 4 所示，該訓練算法會同步更新所有子網(wǎng)絡的參數(shù)。在 TAO 中，VisualChangeNet 支持將兩幅圖像作為黃金樣本和測試樣本輸入，其目標是檢測“黃金或參考”圖像與“測試”圖像之間的變化。TAO 支持 VisualChangeNet 架構(gòu)的 FAN 骨干網(wǎng)絡。

TAO 支持兩類變化檢測網(wǎng)絡，分別是 VisualChangeNet-Segmentation 和 VisualChangeNet-Classification。在本文中，我們將使用 VisualChangeNet-Segmentation 模型，通過分割 MVTec-AD 數(shù)據(jù)集中兩張輸入圖像之間的變化像素，來演示變化檢測。

圖 4. 檢測瓶子類黃金圖像和測試圖像之間變化的

VisualChangeNet-Segmentation 算法架構(gòu)圖

使用工具套件微調(diào) VisualChangeNet 模型非常簡單，無需編碼經(jīng)驗。只需在 TAO 套件中加載數(shù)據(jù)、設置實驗配置并運行訓練命令即可。

該實驗配置文件定義了 VisualChangeNet 模型架構(gòu)、訓練和評估的超參數(shù)。可在訓練模型前在 Jupyter Notebook 中查看并編輯該配置文件。

我們使用該配置對 VisualChangeNet 模型進行微調(diào)，在該配置中定義了一個帶有預訓練 FAN-Hybrid-Base 骨干（即基準模型）的 VisualChangeNet 模型，然后對該模型進行 30 個歷元訓練，批量大小為 8。下面展示了部分實驗配置以及一些關(guān)鍵參數(shù)。完整的實驗配置可在 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

為了調(diào)整模型的性能，可以修改的一些常用值包括訓練歷元數(shù)、學習率（lr）、優(yōu)化器和預訓練骨干。如要從頭開始訓練，可以將 pretrained_backbone_path 設置為空，但這可能會增加為達到高準確率所需的歷元數(shù)和數(shù)據(jù)量。如需進一步了解該實驗配置文件中的參數(shù)，請參閱《VisualChangeNet用戶指南》：https://docs.nvidia.com/tao/tao-toolkit/text/visual_changenet/index.html

現(xiàn)在數(shù)據(jù)集和實驗配置已經(jīng)準備就緒，我們要開始在 TAO 套件中進行訓練了。運行第 5.1 節(jié)中的代碼塊，使用單顆 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 數(shù)據(jù)集上訓練 VisualChangeNet-Segmentation 模型。在訓練過程中，模型將學習如何識別缺陷對象并輸出顯示缺陷區(qū)域的分割掩碼。訓練日志包含驗證數(shù)據(jù)集上的準確率、訓練損失、學習率和訓練后的模型等，都將保存在實驗配置中設置的結(jié)果目錄中。

第 4 步：評估模型

訓練完成后，我們可以使用 TAO 在驗證數(shù)據(jù)集上對模型進行評估。VisualChangeNet Segmentation 的輸出是 2 張給定輸入圖像的分割變化圖，表示的是像素級缺陷。筆記的第 6 節(jié)將運行評估模型性能的命令。

!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 分數(shù)和 IoU 等。

OA = 有變化/無變化像素的總體準確率（輸入維度 - 256×256）

表 1. MVTec-AD 二進制變化檢測（瓶子類）

中的 VisualChangeNet 模型評估指標

第 5 步：部署模型

可以使用NVIDIA DeepStream或NVIDIA Triton部署這一經(jīng)過微調(diào)的模型。我們將其導出為 .onnx 格式。筆記第 8 節(jié)將運行 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 推理服務器的參考實現(xiàn)。

實時推理性能

推理以 FP16 精度在所提供的未修剪模型上運行。該推理性能在嵌入式Jetson OrinGPU 和數(shù)據(jù)中心 GPU 上使用 trtexec 運行。Jetson 設備以 GPU 頻率最高的 Max-N 配置運行。

請執(zhí)行以下命令來運行 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>

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

總結(jié)

在本文中，我們了解了如何使用 TAO 套件微調(diào) VisualChangeNet 模型，并將其用于分割 MVTec 數(shù)據(jù)集中的缺陷，從而實現(xiàn) 99.67% 的總體準確率。

現(xiàn)在，您還可以使用 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 大會。