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

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

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

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

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

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

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

圖 1. 通過(guò)將缺陷圖像與黃金圖像進(jìn)行比較

來(lái)分段預(yù)測(cè)缺陷物體的缺陷掩碼

第 1 步：設(shè)置先決條件

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

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

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

   *請(qǐng)注意：VisualChangeNet 模型只能在 5.1 及以上的版本中運(yùn)行。

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

4. 將下載鏈接粘貼到 Jupyter Notebook 第 2.1 部分中的“FIXME”位置，然后運(yùn)行該筆記本單元。本文主要介紹瓶子對(duì)象，但所有 15 個(gè)對(duì)象都能在該筆記本中使用。圖 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 中的“瓶子”類(lèi)別展示了如何在工業(yè)檢測(cè)用例中使用 TAO 套件與 VisualChangeNet 來(lái)實(shí)現(xiàn)自動(dòng)光學(xué)檢測(cè)。

在 Jupyter Notebook 下載數(shù)據(jù)集后，運(yùn)行該筆記的第 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ù)集專(zhuān)門(mén)用于異常檢測(cè)。我們將兩者合并成一個(gè)包含 283 幅圖像的綜合數(shù)據(jù)集，然后將其分成 253 幅訓(xùn)練集圖像和 30 幅測(cè)試集圖像。兩個(gè)數(shù)據(jù)集都包含有缺陷的樣本。

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

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

第 2 步：下載 VisualChangeNet 模型

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

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

運(yùn)行筆記第 3 部分以安裝 NGC 命令行工具，并從 NGC 下載該預(yù)訓(xùn)練骨干。

# 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 套件訓(xùn)練模型

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

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

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

圖 4. 檢測(cè)瓶子類(lèi)黃金圖像和測(cè)試圖像之間變化的

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

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

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

我們使用該配置對(duì) VisualChangeNet 模型進(jìn)行微調(diào)，在該配置中定義了一個(gè)帶有預(yù)訓(xùn)練 FAN-Hybrid-Base 骨干（即基準(zhǔn)模型）的 VisualChangeNet 模型，然后對(duì)該模型進(jìn)行 30 個(gè)歷元訓(xùn)練，批量大小為 8。下面展示了部分實(shí)驗(yàn)配置以及一些關(guān)鍵參數(shù)。完整的實(shí)驗(yàn)配置可在 Jupyter Notebook 中進(jìn)行查看。

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)整模型的性能，可以修改的一些常用值包括訓(xùn)練歷元數(shù)、學(xué)習(xí)率（lr）、優(yōu)化器和預(yù)訓(xùn)練骨干。如要從頭開(kāi)始訓(xùn)練，可以將 pretrained_backbone_path 設(shè)置為空，但這可能會(huì)增加為達(dá)到高準(zhǔn)確率所需的歷元數(shù)和數(shù)據(jù)量。如需進(jìn)一步了解該實(shí)驗(yàn)配置文件中的參數(shù)，請(qǐng)參閱《VisualChangeNet用戶(hù)指南》：https://docs.nvidia.com/tao/tao-toolkit/text/visual_changenet/index.html

現(xiàn)在數(shù)據(jù)集和實(shí)驗(yàn)配置已經(jīng)準(zhǔn)備就緒，我們要開(kāi)始在 TAO 套件中進(jìn)行訓(xùn)練了。運(yùn)行第 5.1 節(jié)中的代碼塊，使用單顆 GPU 啟動(dòng) VisualChangeNet 訓(xùn)練。

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

該單元將開(kāi)始在 MVTec 數(shù)據(jù)集上訓(xùn)練 VisualChangeNet-Segmentation 模型。在訓(xùn)練過(guò)程中，模型將學(xué)習(xí)如何識(shí)別缺陷對(duì)象并輸出顯示缺陷區(qū)域的分割掩碼。訓(xùn)練日志包含驗(yàn)證數(shù)據(jù)集上的準(zhǔn)確率、訓(xùn)練損失、學(xué)習(xí)率和訓(xùn)練后的模型等，都將保存在實(shí)驗(yàn)配置中設(shè)置的結(jié)果目錄中。

第 4 步：評(píng)估模型

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

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

TAO 中的評(píng)估命令將返回驗(yàn)證集的多個(gè) KPI，比如在缺陷類(lèi)（缺陷像素）上的準(zhǔn)確度、精確度、召回率、F1 分?jǐn)?shù)和 IoU 等。

OA = 有變化/無(wú)變化像素的總體準(zhǔn)確率（輸入維度 - 256×256）

表 1. MVTec-AD 二進(jìn)制變化檢測(cè)（瓶子類(lèi)）

中的 VisualChangeNet 模型評(píng)估指標(biāo)

第 5 步：部署模型

可以使用NVIDIA DeepStream或NVIDIA Triton部署這一經(jīng)過(guò)微調(diào)的模型。我們將其導(dǎo)出為 .onnx 格式。筆記第 8 節(jié)將運(yùn)行 TAO 導(dǎo)出命令。

!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 模型與訓(xùn)練好的 .pth 模型保存在同一目錄下。如要部署到 Triton，請(qǐng)查看 GitHub 上的 tao-toolkit-triton 資源庫(kù)。該項(xiàng)目提供了將許多 TAO 模型（包括 VisualChangeNet Segmentation）部署到 Triton 推理服務(wù)器的參考實(shí)現(xiàn)。

實(shí)時(shí)推理性能

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

請(qǐng)執(zhí)行以下命令來(lái)運(yù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>

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

總結(jié)

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

現(xiàn)在，您還可以使用 NVIDIA TAO 檢測(cè)制造工作流程中的缺陷。

如要開(kāi)始使用，請(qǐng)先：

• 從 NVIDIA NGC 目錄下載 VisualChangeNet 模型。

• 按照《TAO 快速入門(mén)指南》設(shè)置 TAO 啟動(dòng)器：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

• 請(qǐng)通過(guò)NVIDIA 文檔進(jìn)一步了解 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 日在美國(guó)加州圣何塞會(huì)議中心舉行，線上大會(huì)也將同期開(kāi)放。點(diǎn)擊“閱讀原文”或掃描下方海報(bào)二維碼，立即注冊(cè) GTC 大會(huì)。